Optimizing Sparse Tensor Contraction with TiledArray

[ChemChuan]

Dear contributors: I am working on tensor contraction problems involving various formulas, such as: $X(a,b)=\sum{c,d}T(c,d)*H(a,b,c,d)$ $X(a,b)=\sum{c,d}T(a,c)T(b,d)H(a,b,c,d)$ $X(a,b)=\sum_{c,d,e,f}T(a,c)T(b,d)T(e,f)*H(c,d,e,f)$

Here are some details:

1. The range of the state indicator a, b, c, d, e, f is 0 to 15.
2. For example, $T(a,c)$ is $16*16=256$ matrix, but only 25 elements are non-zero.
3. And $H(c,d,e,f)$ is a (16,16,16,16) 4D array, with several hundred non-zero elements.

I used to use the einsum function in libtorch in the past, but there is no further optimization for the einusm of this sparse matrix in libtorch. Can tiledarray be applied to these types of situations? If possible, could you provide me with an example or reference.

Thanks for your help!

[evaleev]

@ChemChuan I'd say for the specific problem sizes and tensor expressions TiledArray will not be useful. I'm making several assumptions here:

• your tensors have unstructured sparsity, and
• the problem sizes and tensor expressions are representative of your problem sizes.

Rationale:

• TA is designed primarily for problems where sparsity is structured. Such sparsity is ubiquitous in physical simulation where one can figure out a priori how to express the problem in a way that tensors are block-sparse. In other contexts (e.g. data analytics) figuring out the sparsity structure involves effort. I'm assuming your are talking about unstructured sparsity here, with which TiledArray does not deal with out of the box. Please correct me if my assumption is incorrect.
• In my limited experience unstructured sparsity greatly reduces the FLOP throughput, which means that the degree of sparsity (% of nonzeros in the tensor) must be << 1% to make sparsity exploitable. In TA we focus on the most common type of structured sparsity pattern (block-sparse tensors) that still allows to attain high performance due to tensor blocks being themselves dense, hence resulting in high arithmetic intensity kernels in most important cases (tensor contraction, etc.)

I have followed up with my collaborators who are working on developing tools for evaluation of (structured and unstructured) sparse tensor networks such as described here, they may provide further comments here.

[smr97]

@ChemChuan, it seems you want to solve a sparse tensor network with element-wise sparsity in all inputs. As Prof. Valeyev mentioned, we built CoNST exactly for this purpose. It is an IR generator for TACO (another tensor algebra compiler). Please take a look here: https://github.com/HPCRL/CoNST The expressions you're looking at seem to be somewhat similar to MTTKRP, so maybe this example would be the best to follow along: https://github.com/HPCRL/CoNST/blob/master/test_tree_mttkrp.py

CoNST generates a TACO kernel that will be compiled and run using the TACO compiler (https://github.com/tensor-compiler/taco).

Finally, it would be useful to know more about your application from a research standpoint. I'm wondering if there's some paper about these expressions that you want to run. What domain/problem do they come from and is there an open source repository of the data being used for these contractions?

[ChemChuan]

Thank you very much for your response and your collaborators. For my tensors $T(a,b)$ or $H(a,b,c,d)$, I can determine in advance which elements are nonzero, but their positions may not follow some specific patterns and are scattered throughout the tensor. For example, nonzero elements in $T(a,b)$ are $(a,b) = (1,1),(1,2),(3,3),(4,5),(6,15),(7,13),(8,14),(9,10),(10,11)...$ Would this be classified as structured sparsity (or block-sparse) or unstructured sparsity?

In this case, would it be feasible to use the CoNST tool mentioned above to generate my code and then process it with the TACO compiler?

I am studying the materials mentioned above. Thank you again for your and your collaborators' help!