• Corpus ID: 238253309

An Attempt to Generate Code for Symmetric Tensor Computations

@article{Shi2021AnAT,
  title={An Attempt to Generate Code for Symmetric Tensor Computations},
  author={Jessica Shi and Stephen Chou and Fredrik Kjolstad and Saman P. Amarasinghe},
  journal={ArXiv},
  year={2021},
  volume={abs/2110.00186}
}

• Jessica Shi, Stephen Chou, +1 author Saman P. Amarasinghe
• Published 1 October 2021
• Computer Science
• ArXiv

Symmetric matrices, a frequently studied topic in linear algebra, can be extended to higher dimensions through symmetric tensors, which arise in domains such as computational physics and chemistry [5, 11]. The fundamental mathematical appeal of the study of symmetry also renders these tensors useful in contexts ranging from a rather beautiful equivalence with homogenous polynomials [3] to more concrete applications, including decompositions [9] and finding eigenvalues [8]. Knowing about a tensor’s symmetry tells us which of its values are guaranteed to be equal. If leveraged effectively, this can be utilized to improve both storage usage and computation time: values that are the same do not need to be redundantly stored or computed. Many existing works focus on optimizing one specific type of tensor computation involving symmetry. On the other hand, compilers such as taco [6] can produce code for a broad class of tensor algebra expressions, with support for per-dimension level formats [2]. However, taco does not presently provide a way to indicate or apply information about which and how tensors are symmetric. Our goal is thus to work towards a compiler-based approach for symmetric tensor computations, relying on taco for aspects of code generation unrelated to symmetry. The remainder of this document is outlined as follows. We start in Section 2 by reviewing standard definitions, which characterize a tensor’s symmetry by a partition on its dimensions, such that coordinates along dimensions that belong to the same part can be permuted without changing the value. Then, we consider in Section 3 the problem of representing symmetric tensors in storage. Some values can be omitted, but when the remaining values are then placed in a dense array, we need new formulas to index into them. While non-symmetric tensors can be geometrically viewed as rectangular prisms, the lower half of, say, a symmetric matrix looks like a triangle. As we move into higher dimensions, the triangles become simplexes, and as a result, our formulas make heavy use of simplicial numbers. At this point, we will be able to reason about the symmetries of individual tensors, so we proceed a step further in Section 4 to provide abstractions for thinking about the symmetry of a computation, which will ultimately inform the structure of the code that we generate. Interesting cases arise when the various input tensors in the computation are symmetric along dimensions with different indices. Section 5 details how to generate a loop structure to handle these cases. Our method is to separate the for loops that iterate over a particular dimension based on whether the values in the corresponding region of the input are stored or omitted, in order to determine how to access the tensors correctly. We end by observing that the storage savings exploited in the resulting code come at the cost of poor computational performance.