  Tensor algebra is widely used in many applications, such as scientific
computing, machine learning, and data analytics. The tensors represented
real-world data are usually large and sparse. There are tens of storage formats
designed for sparse matrices and/or tensors and the performance of sparse
tensor operations depends on a particular architecture and/or selected sparse
format, which makes it challenging to implement and optimize every tensor
operation of interest and transfer the code from one architecture to another.
We propose a tensor algebra domain-specific language (DSL) and compiler
infrastructure to automatically generate kernels for mixed sparse-dense tensor
algebra operations, named COMET. The proposed DSL provides high-level
programming abstractions that resemble the familiar Einstein notation to
represent tensor algebra operations. The compiler performs code optimizations
and transformations for efficient code generation while covering a wide range
of tensor storage formats. COMET compiler also leverages data reordering to
improve spatial or temporal locality for better performance. Our results show
that the performance of automatically generated kernels outperforms the
state-of-the-art sparse tensor algebra compiler, with up to 20.92x, 6.39x, and
13.9x performance improvement, for parallel SpMV, SpMM, and TTM over TACO,
respectively.
