Sparse tensor networks are commonly used to represent contractions over sparse tensors. Tensor contractions are higher-order analogs of matrix multiplication. Tensor networks arise commonly in many domains of scientific computing and data science. After a transformation into a tree of binary contractions, the network is implemented as a sequence of individual contractions. Several critical aspects must be considered in the generation of efficient code for a contraction tree, including sparse tensor layout mode order, loop fusion to reduce intermediate tensors, and the interdependence of loop order, mode order, and contraction order. We propose CoNST, a novel approach that considers these factors in an integrated manner using a single formulation. Our approach creates a constraint system that encodes these decisions and their interdependence, while aiming to produce reduced-order intermediate tensors via fusion. The constraint system is solved by the Z3 SMT solver and the result is used to create the desired fused loop structure and tensor mode layouts for the entire contraction tree. This structure is lowered to the IR of the TACO compiler, which is then used to generate executable code. Our experimental evaluation demonstrates very significant (sometimes orders of magnitude) performance improvements over current state-of-the-art sparse tensor compiler/library alternatives.

翻译：摘要：稀疏张量网络常被用于表示对稀疏张量的缩并操作。张量缩并是矩阵乘法的高阶泛化。在科学计算和数据科学的许多领域中，张量网络普遍存在。经过转换为二叉树形式的二元缩并操作后，该网络实现为一系列独立的缩并步骤。在生成缩并树的高效代码时，必须考虑若干关键因素，包括稀疏张量布局的模序、通过循环融合减少中间张量、以及循环顺序、模序与缩并顺序的相互依赖关系。我们提出了一种新颖方法CoNST，通过单一公式化框架综合考量这些因素。该方法构建了一个约束系统，对这些决策及其相互依赖关系进行编码，同时旨在通过融合生成降阶的中间张量。该约束系统由Z3 SMT求解器求解，其结果用于为整个缩并树创建所需的融合循环结构与张量模布局。该结构被降级为TACO编译器的中间表示，进而生成可执行代码。实验评估表明，与当前最先进的稀疏张量编译器/库方案相比，我们的方法实现了非常显著的性能提升（有时可达数量级）。