Networks (GNN), Physical Simulation, and Quantum Chemistry [Hamilton et al. 2017;Hu et al. 2020;Rahman et al. 2021] make use of sparse tensor contractions.

Due to the compressed data formats that store the sparse tensors, their contractions are realized as non-affine loop nests, where bounds depend on the input, and accesses are indirect. The non-affine loop nests of sparse tensor contractions prevent us from directly applying classical affine loop transformation frameworks to reduce load imbalances and bad locality for performance enhancement. This challenge has given rise to specialized compilers for sparse tensor computations [Bik et al. 2022;Bik and Wijshoff 1993;Kjolstad et al. 2019Kjolstad et al. , 2017;;Kotlyar et al. 1997;Senanayake et al. 2020;Tian et al. 2021;Venkat et al. 2015] and various abstractions for the schedule -realization of computation (e.g., loop structure, parallelization etc.) -to separate it from the computation. Schedule abstractions ( § 2.4) make it convenient to realize a plethora of ways to materialize a computation using transformations such as loop reordering, loop fusion/fission, loop tiling, loop parallelization, etc.

Choosing a better-performing schedule for a sparse tensor contraction is not straightforward. Therefore, it is more challenging than finding a schedule for its dense counterpart, which is realized as affine loop nests (i.e., There exist well-studied analytical cost models of schedules and machines for dense tensor computations). The schedule selection heavily depends on sparse tensor inputs (number of non-zero values and sparsity structure), making it difficult to pick a performant one for sparse tensor computations. Hence, the simplest method to evaluate the cost of a schedule is to execute it on a given machine using the provided sparse tensor inputs to measure the time it takes to finish the execution. The sheer number of schedules makes it an arduous time-consuming process, rendering it impractical to execute all schedules in search of the best one. Also, it is important to note that there may not be a single best schedule for all sparse tensor inputs and machines, and some schedules may be asymptotically better than others.

The challenge in finding a performant schedule for sparse tensor contractions arises from two main factors: vast space of schedules and heavy dependency on sparse tensor inputs. We provide a systematic way to completely explore the vast space of schedule at compile-time rooted in transformations (i.e., loop reordering and loop/kernel fusion/fission), which makes it convenient to realize the schedule. The exploration of the schedule space is augmented with machine-independent pruning strategies and symbolic sparse tensor input attributes at compile-time to filter most of the schedules and keep a handful of schedules to be evaluated at run-time with machine-dependent parameters and concrete sparse tensor input attributes to select a performant schedule. As there are asymptotically superior schedules, the pruning strategy encompasses comparing schedules for both time and auxiliary memory complexity, which depends on the attributes of sparse tensor inputs. To the best of our knowledge, prior work does not optimize for both time and auxiliary memory complexity [Ahrens et al. 2022;Kanakagari and Solomonik 2023].

Consider this example of sparse tensor times matrix contraction: = B 1 . This computation can be expressed using a simple linear loop nest with a time complexity of (nnz(B ) ). Alternatively, the contraction can be expressed as = and = -two separate computations with a total time complexity of (nnz(B ) + ) and a dense temporary . Another schedule can be obtained from the observation that the outer loops of the first computation (producer) can be fused with the second computation (consumer) (i.e., loop fusion). This schedule reduces the overall time complexity to ( (nnz(B ) + nnz(B ) )) 2 and a scalar temporary, which is asymptotically superior to both of the previous schedules in time and memory complexity. The last schedule has a branching loop structure (i.e., imperfectly nested loop nest) that is different from the other two schedules, which have simple loop structures (i.e., perfectly nested loop nests). However, the last schedule dominates the other two schedules in terms of both time and memory complexity ( § 3). Therefore, to explore schedules with multi-level branching loop structures, which are of asymptotically superior time and auxiliary memory complexity, we introduce the extended representation of branched iteration graphs [Dias et al. 2022] and a new scheduling directive to realize such schedules ( § 4). Furthermore, we explore the schedule space of a given sparse tensor computation and present strategies based on partially ordered sets (posets) that can be combined with user-defined constraints at compile-time to prune the schedule space ( § 5). Contributions of this paper are as follows:

Recursive extension of branched iteration graph We generalize the branched iteration graph (BIR) representation of SparseLNR [Dias et al. 2022] to support schedules with multiple levels of imperfectly nested loops and new scheduling primitives to realize the schedules by recursively applying loop/kernel fusion/fission with loop reorder. Complete schedule space exploration We provide a strategy to explore the schedules of a given sparse tensor contraction guaranteed to cover the complete space of schedules with loop structures, including multi-level branching (i.e., multiple levels of imperfectly nested loops), attainable using loop/kernel fusion/fission. Novel auto-scheduler We introduce a novel poset-based auto-scheduler to prune the space of schedules to create a Pareto frontier wrt. both time and auxiliary memory complexity. We use a Satisfiability Modulo Theory (SMT) solver to compare the symbolic time and memory complexity with user-defined constraints.

The rest of the paper is organized as follows. We provide the necessary background in Section 2 and in Section 3, we motivate the problem. The multi-level branched iteration graph and the scheduling primitives are introduced in Section 4. We discuss schedule exploration and selection in Section 5. Evaluation of our auto-scheduler is presented in Section 6. We conclude the paper in Section 8 with a discussion.

This section discusses the necessary background on sparse tensor access constraints, tensor index notation, iteration graph representation, and scheduling primitives to understand the challenge in auto-scheduling for sparse tensor contraction.

There are several compressed data formats used to store sparse tensors: Compressed Sparse Row (CSR), Sorted Coordinate (Sorted COO), Compressed Sparse Fiber (CSF), etc., to name a few. These formats are abstracted by level format [Chou et al. 2018], a tree structure that shows the order in which index arrays must be traversed to retrieve an element. The sparse tensor access constraints are imposed by the order of access of the index arrays in compressed data formats. For example, if A is in CSR format, the row index should be traversed to get to the column index, which results in a dependency between and , the indices traverse rows and columns of A, respectively. Therefore, the loops and belong to cannot be freely reordered. TACO Format Abstraction [Chou et al. 2018] describes the level formats in detail.

The notation that describes tensor contraction operations is based on the Einstein Summation (Einsum) convention. This notational convention implies summation over a set of repeated indices. For example, the expression ( , ) = ( , ) • ( , ) implies summation over the repeated index

Fig. 1. An example of an iteration graph for sparse matrix-matrix multiplication and corresponding code.

and equivalent to the standard mathematical notation =foot_3 . We use both these notations interchangeably in the text. Since this computation can be performed using a simple linear triply nested loop, its iteration time complexity is ( ), where , , and are the loop bounds. If B is sparse, then the iteration time complexity is (nnz(B ) ), where nnz is the number of non-zero elements.

Consider the example sparse matrix-matrix multiplication (SpMM), = B . An iterator that iterates through all of , , and can read each value of B , , multiply each value sharing the same , and store the result in . An example iteration graph is shown in Figure 1a, and this internal representation (IR) is used to generate code in Figure 1b. The nodes in the iteration graph represent indices in the Einsum notation. This is an acyclic graph where the edges represent the dimensions of tensors and how they map to indices. Since B is sparse, 1 and 2 , incidents on indices and must not change the order, and other indices can appear in any order as they traverse dense tensors (e.g., 1 and 2 ). No other sparse tensor access constraints ( § 2.1) are imposed. TACO [Kjolstad et al. 2017] describes the concept of iteration graphs in detail.

A schedule describes one way of realizing a computation, and multiple schedules can realize the same computation. For example, we can change the loop order in Figure 1 to get the order , , instead of , , . TACO and Sparse Iteration Space Framework [Senanayake et al. 2020] describe the importance of abstractions to separate the computation from the schedule. The algorithmic and scheduling languages describe the computation and schedule, respectively, and scheduling primitives form the scheduling language. Some of the scheduling primitives are as follows: reorder to reorder the loops; split to split a loop for tiling; collapse to collapse one loop onto another; parallelize and vectorize for parallel execution. TACO-Workspaces [Kjolstad et al. 2019] introduces precompute to add dense intermediaries to schedules. In Section 4, we introduce a new scheduling primitive called loopfuse, which can produce loop nests with branched loops (i.e., imperfectly nested loops) combined with the reorder directive.

It may not be straightforward to decide whether to apply transformations across multiple kernels. The decision depends both on the iteration complexity of the final loop nests and the working set SparseAuto: An Auto-scheduler for Sparse Tensor Computations using Recursive Loop Nest Restructuring 290:5 1 for perm(l,m,n,i,j_pos,k_pos):

• nnz( ), Memory: 0 Loop Depth: 6, Memory Depth: 0 1 for perm(m, l): 2 T<k> = 0 3 for i, j_pos, k_pos:

• (nnz( ) + ), Memory: Loop Depth: 5, Memory Depth: 1 1 for l: 2

T<k,j> = 0 3 for i, j_pos, k_pos: 4 T(k,j) += B(i,j,k) * C(i,l) 5

for perm(j, m, k, n):

), Memory: Loop Depth: 5, Memory Depth: 2

for l: 2 T<k,j> = 0 3 for i, j_pos, k_pos: 4 T(k,j) += B(i,j,k) * C(i,l) 5 for perm(m, k): 6 t = 0 7 for j: 8 t += T(k,j) * D(j,m) 9 for n: 10 A(l,m,n) += t * E(k,n) (d) = ( ( B • ) • ) • Time: • (nnz( ) + ( + )), Memory: Loop Depth: 4, Memory Depth: 2 1 for l: 2 T<j,k> = 0 3 for i, j_pos, k_pos: 4 T(k,j) += B(i,j,k) * C(i,l) 5 T<m,k> = 0 6 for perm(j, m, k): 7 T(m,k) += T(j,k) * D(j,m) 8 for perm(m, k, n): 9 A(l,m,n) += T(m,k) * E(k,n) (e) = ( ( B • ) • ) • Time: • (nnz( ) + ( + )), Memory: + Loop Depth: 4, Memory Depth: 2 Fig. 2. Different schedules of executing = B • • • .

Here, the code snippet 2a has a perfectly nested loop structure while all the other code snippets has a nested loop structure. Here, _ refers to the non-affine loop associated with the index . The loop _ is non-affine because B is sparse. The code snippets 2b and 2c has one level of branching whereas the code snippets 2d and 2e has a branch nesting depth of two.

sizes. 4 If the working set sizes are small and fit into the cache then, it is better to use the version with lower iteration complexity. Otherwise, it is better to use the schedule with lower auxiliary memory. Hence, an auto-scheduler that only looks at the iteration complexity or only the auxiliary memory complexity may choose the wrong schedule as the final output or prune a good schedule from the search space in the process.

There may be many schedules to perform a tensor contraction, and which one to choose depends on your viewpoint. Consider the following example involving a sparse tensor B:

Figure 2a refers to performing the computation using a simple loop nest of depth 6. The same computation can be written as in figures 2b, 2c, 2d, and 2e with branching loop nests of depth 4 or 5. In this section, we will discuss the performance of these different schedules. We will evaluate all the schedules with the same loop structure, but with different index ordering and report the best one. For the loop structure in Figure 2d, we will evaluate both the inner loop order of , and , . Similarly, for Figure 2b, we will evaluate four different loop orders, two of them by interchanging the inner loops , and two of them by interchanging the outer loops , .

From the asymptotic time complexity viewpoint, an auto-scheduler might lean towards pruning Schedule 2a. This is due to its loop nesting depth of 6 and time complexity of (nnz( ) ), in contrast to the schedule in Figure 2b with a loop nesting depth of 5 and asymptotic time complexity of (nnz( ) + ) or the schedule in Figure 2d with a loop nesting depth of 4 and asymptotic time complexity of (nnz( ) + ( + )). Notably, the schedule in Figure 2e has the same asymptotic time complexity as the schedule in Figure 2d, while the asymptotic time complexity of the schedule in Figure 2c is (nnz( ) + ). These schedules can be placed on a asymptotic time complexity vs. auxiliary memory complexity space plot as shown in Figure 3, relative to each other. From the perspective of asymptotic time complexity, an autoscheduler might favor either the schedule in Figure 2d or Figure 2e, both having a loop depth of 4, the lowest among the five schedules in Figure 2. Comparing these two schedules, Figure 2d uses one 2D auxiliary memory for storing intermediate results between branched loop nests, while Figure 2e uses a 2D and a 1D auxiliary memory. Consequently, the former has lower memory complexity than the latter. In summary, the schedule in Figure 2d dominates Figure 2e, as both schedules share the same asymptotic time complexity, but the former is better in terms of auxiliary memory complexity.

Comparing the schedules in Figures 2d and 2c from the asymptotic memory complexity perspective, both exhibit an auxiliary memory complexity of ( ). The time complexity of the former, Figure 2d, is superior with ( + ) being better than ( ) for larger values of and . Consequently, Figure 2d dominates Figure 2c. For the sake of brevity, comparisons involving Figures 2e and 2c with other schedules are omitted in the following paragraphs.

Consider the comparison of the schedules in Figures 2a, 2b, and 2d when the bounds change in the range as follows; 1 ≤ ≤ 1800, 1 ≤ ≤ 1600, 400 ≤ ≤ 4000, 8 ≤ ≤ 256, 8 ≤ ≤ 256, 8 ≤ ≤ 256, and 0.001 ≤ sparsity( ) ≤ 0.01. Note that the schedule in Figure 2a dominates both the schedules in Figures 2b and 2d in terms of the auxiliary memory usage because no auxiliary memory is used in the Schedule 2a. Although the loop depth is four for the schedule in Figure 2d, within the given ranges of bounds and the sparsity of , we cannot claim that it is the best in all cases. Let us look at some cases by changing the loop bounds and sparsity for tensor . The evaluation configuration is explained in Section 6.

Case 1 : Lowest loop depth schedule (Figure 2d) is the best In this case, we set the loop bounds for the schedules in Figure 2 to specific values: = 1800, = 800, = 1000, = 64, = 16, = 32 5 , and sparsity( ) = 0.08. Under these conditions, the iteration time complexities follow the 5 , , , , , and are the loop bounds of loops with indices , , , , , and , respectively.

Proc. ACM Program. Lang., Vol. 8, No. OOPSLA2, Article 290. Publication date: October 2024.

SparseAuto: An Auto-scheduler for Sparse Tensor Computations using Recursive Loop Nest Restructuring 290:7 inequality Φ( )foot_5 < 7.5 * Φ( ) ≈ Φ( ) < 237.5 * Φ( ) ≈ Φ( ). The corresponding execution times follow the inequality Ψ( )foot_6 = 2.48 < Ψ( ) = 6.26 < Ψ( ) = 32.40 . The schedule in Figure 2d exhibits the lowest loop depth and iteration time complexity. An auto-scheduler that factors in loop depth could choose the best schedule in this case.

Case 2 : Effect of the size of auxiliary memory Adjusting the loop bounds to = 1600, = 2000 and sparsity( ) = 0.02 while maintaining other loop bounds as in the previous example, the schedule in Figure 2d now incurs an auxiliary temporary memory requirement of 12.21 (compared to 3.05 in Case 1 ). This consumes more than 50% of the last-level cache (LLC). The iteration time complexities follow the inequality:

Execution times for the schedules follow the inequality, Ψ( ) = 8.60 < Ψ( ) = 10.20 < Ψ( ) = 33.51 . The schedule in Figure 2d, with the minimum loop depth, exhibits the lowest iteration time complexity as in the previous example, but the schedule in Figure 2b, with a loop depth of 5, performs better. It is evident from this case that a good auto-scheduler must consider the sizes of the auxiliary memory arrays used in the computation. Consequently, an auto-scheduler solely reliant on loop depth would fail in this scenario.

. This scenario is an example where the schedule with the highest loop depth (Schedule ) executes the fastest. An auto-scheduler that factors in loop depth would discard this schedule in favor of the schedules with lower loop depths. This case highlights the need for a robust auto-scheduler to consider factors beyond loop depth.

Case 4 : Neither the lowest loop depth, nor the highest loop depth schedule (Figure 2b) is the best Setting loop bounds and sparsity as = 265, = 1207, = 479, = 251, = 234, = 42, and sparsity( ) ≈ 0.0033, Figure 2b performs the fastest at Ψ( ) = 513 , followed by Ψ( ) = 1.14 for Figure 2d, and Ψ( ) = 1.66 for Figure 2a. In this scenario, auxiliary memories account for less than 12% of the LLC. For these values, Φ( ) = 1.13 × Φ( ) and Φ( ) = 40.4 × Φ( ). Execution times align with iteration complexities, and auxiliary memory usage is reasonably modest. Unlike previous cases, where the best loop or memory depth proved to be the most efficient, this instance underscores the need for schedulers to consider multiple factors beyond loop and auxiliary memory depth when pruning the search space.

The insights drawn from the motivating example and our approach to schedule selection can be summarized as follows.

Multi-Level Branched Loop Nests Nested loop computations with reduced loop depth (as in Case 1 ) are crucial. However, existing scheduling languages lack support for multi-level branched loop nests. To address this, we extend the branched iteration graph (BIG) [Dias et al. 2022;Kjolstad et al. 2017] to accommodate recursive, multi-level branched iteration graphs with multi-dimensional temporary buffers. We also enhance the scheduling language to support recursive fusion by adapting TACO 's code generation strategies.

Both time and auxiliary memory complexities contribute to the schedule's execution time. An effective auto-scheduler needs to consider both aspects when selecting a schedule (as seen in Cases 2 -4 ). If a schedule's auxiliary memory takes up a large portion of the last-level cache, it tends to perform worse than the alternatives (as observed in Case 2 ). To address this, we introduce an auto-scheduler that employs an SMT solver. The solver is guided by the constraints of sparse computations and reasons about the partial orders of time and auxiliary memory complexity. This approach effectively prunes the search space, leading to the selection of schedules that dominate others in both time and memory complexity.

Tensor contractions can be materialized using a simple linear loop nest where there would be a corresponding loop for each of the indices in the Einsum expression. This loop nest is represented as a linear iteration graph (LIG) as explained in Section 4.1, which is used for sparse code generation. However, this simple loop nest must respect the sparse tensor access constraints. For example, if a sparse tensor is in CSR format, the row should be accessed first. In this section, we describe an algorithm to recursively generate a branched iteration graph (BIG), the transformation required to convert a LIG into a multi-level BIG ( § 4.2), and how this transformation can cover a plethora of possible loop nests for the computation ( § 4.5).

Consider a tensor contraction:

Here, 1 . . . denote the input tensors; denotes the output tensor; idx denotes the indices that need to be contracted from the tensor expression. The example tensor contraction can be materialized in several ways, two of which are as follows (access indices are omitted for brevity):

1 loop 1 . . . loop . . . loop . . . loop : 2 += 1 * ... * * ... * * ... * 1 loop 1 . . . loop . . . loop . . . loop : 2 += 1 * ... * * ... * * ... *

There are two main differences between the two materializations: loops and are swapped, as well as tensors and in the expression. Therefore, the orders of accessing elements of input tensors and storing elements of the output tensor differ. But in general, any permutation of 1 , 2 , . . . , , and any permutation of 1 , . . . , yields the correct output tensor , when we complete all the iterations. This observation also holds when some of the tensors are sparse, although the index order must satisfy the sparse tensor access constraints. Overall, materializations like the ones above belong to an equivalence class because they produce the same output.

We define a linear iteration graph (LIG) as a loop nest with no two loops having the same depth from the root of the nest and an index order that respects all the sparse tensor access constraints. Hence, we consider any permutation of the loops and input tensors that do not violate the sparse access constraints as a representative of an equivalence class since it produces the same result for a given tensor contraction.

In this section, we describe the multi-level BIG transformation. We demonstrate how the algorithm works for the example tensor contraction from Section 3 and by showing how the LIG in Figure 5 transforms into a BIG (5a → 5b and 5a → 5c → 5d → 5e).

The tensor contraction ( , , ) = B( , , ) * ( , ) * ( , ) * ( , ) has a default iteration graph (Figure 5a) which implies generated code in Figure 2a. The IR of this iteration graph is shown in the listing below: Transform Iteration Graph 5a → 5b Let us split the computation into two parts, the producer and the consumer. The first one, (idx temp ) = B( , , ) * ( , ) * ( , ), produces the intermediate temporary tensor with indices idx temp , which is consumed in the second one, ( , , ) = (idx temp ) * ( , ), to generate the output . Here, idx temp = { , , } is obtained by evaluating: Indices(B( , , ) * ( , ) * ( , )) ∩ (Indices( ( , , )) ∪ Indices( ( , ))). Hence, the auxiliary memory required is (

). The split computations from the original one have iteration graphs with the following index orders: l ) m ) i ) j ) k for the producer and l ) m ) n ) k for the consumer.

These ones preserve the index order in the original iteration graph l ) m ) n ) i ) j ) k (Figure 5a). As both iteration graphs share the same indices and at the beginning, they can be fused into the BIG in Figure 5b, facilitating code generation in Figure 2b. Since the unfused sections of the producer and consumer graphs include , , and , respectively, the fused iteration graph would require extra memory of , obtained from { , , } ∩ { , }, to pass the intermediate results between the two computations. Notably, only an auxiliary memory of size is required after fusion, compared to the one with size before fusion.

Transform Iteration Graph 5a → 5c The original tensor contraction can be split into two computations in a different way. For example, (idx temp ) = B( , , ) * ( , ) and ( , , ) = (idx temp ) * ( , ) * ( , ) where idx temp = { , , }. This would result in the producer and consumer iteration graphs l ) i ) j ) k and l ) m ) n ) j ) k , respectively. Since they have a common index at the beginning of the iteration graph, they can be fused to generate the BIG in Figure 5c. Fusing the iteration graph reduces the auxiliary memory requirement to as opposed to the before. After fusion, the IR is shown in the listing below (notice the addition of the temporary 1 and the where clause to combine the producer and consumer computations): Transform Iteration Graph 5c → 5d We notice that, after fusion, each of the unfused sections of the producer and consumer iteration graphs can be treated as separate iteration graphs by keeping all the fused indices fixed in the computation. For example, take the consumer computation, ( , , ) = ( , , ) * ( , ) * ( , ). We can rewrite the computation as (_, , )+ = (_, , ) * ( , ) * ( , ) by fixing , with corresponding iteration graph m ) n ) j ) k . This can be split into two computations: ′ ( , )+ = (_, , ) * ( , ) and (_, , )+ = ′ ( , ) * ( , ), with corresponding iteration graphs m ) j ) k and m ) n ) k , respectively. Both of these iteration graphs have the same first index m that can be fused. The iteration graph in this configuration (not shown in Figure 5) would be able to generate the code in Figure 2e. This configuration would require auxiliary memory of size because the unfused part of each iteration graph shares the common index { } = { , } ∩ { , }. The iteration graph in Figure 5c can be transformed to the iteration graph in Figure 5d with the inner consumer index order m ) k ) n ) j , by reordering the consumer part of Figure 5c.

Transform Iteration Graph 5d → 5e Splitting the consumer computation as described previously yields the producer and consumer iteration graphs m ) k ) j and m ) k ) n , respectively, which can be fused to generate the multi-level BIG in Figure 5e. Since the unfused sections of the producer and consumer graphs do not share any common indices, it only requires a scalar auxiliary memory to pass the intermediate results between the producer and consumer. The final IR is shown in the listing below (notice the use of two temporaries 1 and 2, and the nested combination of where clauses):

1 forall(l, where( 2 forall(m, forall(k, where( 3 forall(n, A(l,m,n) += t2 * E(k,n)), 4 forall(j, t2 += t1(j,k) * D(j,m))))), 5 forall(i, forall(j, forall(k, t1(j,k) += B(i,j,k) * C(i,l))))))

Algorithm 1 shows the pseudo-code for the transformation described previously in Section 4.2. This algorithm takes several inputs: the original iteration graph (LIG or BIG), the ℎ to an inner producer/consumer section, the position to split ( : ), the input tensors in the contraction, and a boolean flag to indicate whether the producer expression is on the left or the right after the split. Input variable ℎ is used in line 1 to access the inner producer/consumer graph sections, which helps to apply the transformation recursively to the inner linear graph sections.

The split operation occurs in lines 2-4. For example, given the expression = * * and = 2, the splits are = * and = * . If = 1, then the splits are = and = * * . The algorithm initially deduces the indices of the temporary resulting from the split (line 5) using the equation

)). This equation calculates the indices in the producer that also appear in either the consumer or the output. The algorithm generates corresponding split expressions in lines 6-7. Subsequently, the producer and consumer graphs are computed in lines 8-9, preserving the index order of the original graph. Then, the algorithm determines the fusible outer loops (lines 10-14) and shared indices (lines 15-16). Finally, it produces the expressions for the producer and consumer in the fused iteration graph in lines 17-18, and the original iteration graph is replaced with the fused iteration graph in lines 19-20. One step of the transformation is linear time with respect to the number of indices in the graph.

Figure 6 shows the implementation of the transformations described in Section 4.2 using the scheduling language. The schedule description in Figure 6a can be used to transform the iteration graph in Figure 5a to the one in Figure 5b. The schedule description in Figure 6b can be used to transform the original iteration graph in Figure 5a to the one in Figure 5e by doing multiple transformations.

• loopfuse scheduling directive splits and fuses LIGs. It takes three parameters:

path identifies linear graph sections, consumer or producer sections, of a BIG. The path parameter must direct to a linear graph section for the transformation to be applied. Here, {} means accessing the root of an iteration graph, {0} means accessing the producer section, and {1} means accessing the consumer section. If the BIG has multiple levels, {0, 1} would access the producer of the first level and then the consumer of the second level. loc specifies the split position in the inner computation. For example, if the inner computation is = * * and loc=2, then the split is = * and = * , and if loc=1, then the split is = and = * * . pol designates the first or second half of the contraction as the producer. If pol=True, then the producer is on the left, and the consumer is on the right, and if pol=False, then vice versa. For example, if the expression is = * * , loc=2 and pol=True, then the split is = * and = * , and if pol=False, then the split is = * and = * . • reorder scheduling directive reorders indices of a linear graph section. It takes two parameters:

1 A(l,m,n) = B(i,j,k) * C(i,l) * D(j,m) * E(k,n); 2 // Index stmt of 5a 3 IndexStmt stmt = A.getAssignment(). concretize(); 4 // Apply transformation 5 stmt = stmt // 5a -> 5b 6 .loopfuse(loc = 2, pol = True, path = {});

(a) 5a → 5b 1 A(l,m,n) = B(i,j,k) * C(i,l) * D(j,m) * E(k,n); 2 // Index stmt of 5a 3 IndexStmt stmt = A.getAssignment().concretize(); 4 // Apply transformation 5 stmt = stmt 6 .loopfuse(loc = 3, pol = True, path = {}) // 5a -> 5c 7 .reorder(order = {m, k, n, j}, path = {1}) // 5c -> 5d 8 .loopfuse(loc = 2, pol = True, path = {1}); // 5d -> 5e (b) 5a → 5c → 5d → 5e Fig. 6. Transformation on the loop contraction path identifies an inner linear graph section.

order specifies the new order of the indices in the linear graph section.

The reorder and loopfuse directives can be used together to obtain the desired multi-level BIG. These two directives can be used in conjunction to generate all possible loop trees for a given tensor contraction.

This section provides a proof sketch for the completeness of Algorithm 1: we argue that the algorithm can generate all possible loop structures for a given tensor contraction using the loopfuse and reorder scheduling directives.

Constraints Two essential constraints ensure the validity of BIG constructed by Algorithm 1 and the equivalence of the BIG to the initial LIG. First, the BIG must not violate any sparse tensor access constraints present in the initial computation. For instance, in all the iteration graphs in Figure 5, the iteration order 1 ) 2 ) 3 (i.e., ) ) ), is consistently maintained when contracting the sparse tensor B with other tensors. Second, a permutation of indices in producer and consumer loops is required to establish identical orders of shared indices (i.e., indices in the temporary tensor).

The second constraint can be understood as follows. Let and be the indices present in the temporary tensor. Let and be the sets of all valid (in the sense of the first constraint) permutations of indices in the loops of producer and consumer, respectively. Let the focus be on

If either condition1 or condition2 is satisfied, then we say that the BIG is valid. If neither condition1 nor condition2 is satisfied, then we say the second constraint is violated and the BIG is invalid. The temporary tensors introduced by the Algorithm 1 are dense. Hence, they do not impose extra constraints, and establishing identical orders of shared indices is not impeded by the temporaries.

Equivalence of BIG and LIG In the context of reasoning about LIGs and BIGs, having a transformation from BIG to LIG, complementary to the one performed by Algorithm 1, proves beneficial. This transformation involves examining the innermost subtree (with no nested subtrees inside of it), as illustrated in Figure 7. Denoting shared indices between the producer and consumer as , producer loops as , and consumer loops as , the transformation requires an order of complying with sparse tensor access constraints in both and , adhering to the second constraint. By keeping the outer loops constant in inner branch computations (details omitted for simplicity), the BIG illustrated in Figure 7a is realized. The contraction for is depicted in Figure 7b, where \ denotes contracted indices. Substituting ( ) inside the consumer (Figure 7c), the consumer computation (

) is moved inside the summation operation, ensuring none of its indices contain the ones in \ (Figure 7d). Further, contracting indices are reintegrated into consumer graph loops (Figure 7e), satisfying all constraints and resulting in a LIG. This recursive process applies to multi-level BIGs, yielding a LIG equivalent to the initial BIG.

Fineteness of the space A conservative upper bound on the number of BIGs can be established by considering the number of input tensors in the tensor contraction ( ), and the number of indices in the tensor contraction ( ). The input tensors can be permuted in ! ways. A BIG can be built by recursively splitting the computation into producer and consumer sections. Since the number of input tensors are , the number of binary trees that can be built is bounded by 2 . The indices can be permuted in ! ways. Since there are 2 splits, indices can be permuted at each split giving ! 2 permutations. At each split operation, there is a choice to fuse the indices or not. Since the number of indices is , the number of ways to fuse the indices is bounded by + 1 for each of those binary trees. Therefore, the total number of BIGs is bounded by ! × 2 × ! 2 × ( + 1).

Completeness Consider the different schedules for a given tensor contraction as points in a space. If you can reach a schedule from another, then they are connected in this space. We established that a BIG can be linearized. Therefore, every BIG is connected to a LIG. As outlined in Section 4.1, LIG schedules are equivalent, and we end up connecting all the points. Focusing on the linearization procedure for a BIG, the movement of from producer to consumer involves reordering the loops of producer and consumer such that after reordering the loops, the shared indices have the same relative ordering. consists of some input tensors in the original tensor contraction. In other words, input tensors in the producer computation are some combination of the input tensors. This combination can be obtained by permuting the input tensors in the original expression and splitting from a specific position. Since each valid BIG can be transformed to a LIG, it is possible to traverse in the direction of LIG to BIG by using the transformation in Sections 4.2-4.4. Therefore, by ( 1) permuting all schedules in our equivalence class, (2) applying the transformation in Section 4.2 to obtain BIGs, and (3) recursively applying ( 1) and ( 2) on inner producer and consumer sections, we can generate all possible iteration graphs (loop structures of schedules) for that computation.

We build an auto-scheduler that, given a tensor contraction, explores the schedule space and chooses a memory-and time-efficient schedule. The main function of the scheduler is pruning the schedule space. The scheduler decides on what schedules to prune by creating a Pareto frontier of schedules using partially ordered sets of time and memory complexity.

The complete pruning pipeline is shown in Figure 8. The pipeline starts by generating schedules in the search space ( § 5.1). The following stages are divided into two parts. The first three stages are executed during compile-time with symbolic expressions ( § 5.2), and the last two stages are executed with concrete expressions at run-time ( § 5.3). We add the language support as an extension to the TACO/SparseLNR compilers. We implement the auto-scheduler in Python. The abstract schedules selected by the auto-scheduler is then described using the language definition provided in the Section 4.

For a given expression and a LIG, we generate all the loop index orders that conform to the sparse tensor access constraints. Then, we split the tensor contraction at different positions for all those index orders. Once tensor contraction is split, we infer the temporary indices and call the same function recursively for both the consumer and producer sub-computations. After those sub-computations return the consumer schedules sch and producer schedules sch , we combine those two schedule spaces as sch × sch to create super schedules that completely describe the initial computation. If the producer and consumer sections can be fused as explained in Section 4.2, we merge the sub-computations to create fused schedules, which we add to the list of schedules.

Consider the previous example ( § 3-4). First, we create different permutations of input tensors. Since tensor contractions are commutative in Einsum notation, ( , , ) = B( , , ) * ( , ) * ( , ) * ( , ) is equivalent to ( , , ) = B( , , ) * ( , ) * ( , ) * ( , ). For each of these permutations, we create permutations of indices that conform to sparse tensor access constraints. These two steps combined create the complete set of LIGs ( § 4.1). For each LIG, we split the input tensors at different positions to generate producers and consumers.

There are two ways in which we can split the input tensors. = * * * can be split as (a)

= * , and = * * : the result of the producer is directly used in consumer, and (b) 1 = * , 2 = * , and = 1 * 2: the consumer expects results of two producers. This procedure can be repeated (recursive application of the algorithm) on the producer and consumer sub-computations. Out of these two ways, the first one is more interesting because it opens up avenues for loop fusion. If the producer and consumer graphs contain the same indices, then we can fuse them. Consider ( , , ) = B( , , ) * ( , ) * ( , ) * ( , ) with index order ) ) ) ) ) , split between , and . The fusion of them would result in ) ⟨ ( , );

:

2). After the fusion operation, we remove the fused indices in the inner computation, for example, (_, , )+ = ( , ) * ( , _) * ( , ), to recursively call the schedule generation procedure, and combine the results with the outer loops.

Stage 1: Memory-Depth-Based Pruning The first pruning stage analyzes the dimensions of the temporary tensors used in the schedule. For example, Schedule 2a does not use any temporaries, and hence, its memory depth is 0. In Schedule 2b, the temporary ⟨ ⟩ has the memory depth of 1, whereas Schedules 2d and 2e have the memory depth of 2 since they employ 2D temporary tensors. If we split the computation as before to = * and = * * but do not fuse (see Figure 5b for the fused version), the memory depth would be 3. The memory depth 3 is too high for most realistic scenarios, and hence, we use it as a threshold to prune schedules like that. The memory-depth-based heuristic goes first in the pruning pipeline because it usually discards many unrealistic schedules. At the end of this stage, we compute the symbolic iteration time and auxiliary memory complexity for each of the schedules. Then, we allocate the schedules into groups using the symbolic (time, auxiliary memory) complexity tuples.

We note that this stage can be replaced by a more sophisticated memory-volume-based pruning mechanism that uses an SMT solver. With such an approach, the user can provide a heuristic upper bound value on the total auxiliary memory use. This type of pruning would guarantee that the schedules with higher depth but lesser auxiliary memory are not pruned and the schedules with lesser depth but higher auxiliary memory are pruned. However, this would require further exploration on how to select the upper bound value which often depends on the execution environment (i.e., machine parameters, such as cache size, etc.). Moreover, a memory-based pruning stage capable of reasoning about the actual volume of auxiliary memory will require the complete information about the loop bounds. Knowing either piece of information-machine parameters or loop bounds-blurs the boundary between compile-time and run-time stages and goes beyond our approach. We leave this idea for future work.

Stage 2: Poset of Loop-and Memory-Depth-Based Pruning In Section 3, we explain that using only the loop depth could prune potentially useful schedules. Therefore, at this stage, we consider both the loop depth and memory depth for pruning to create a Pareto frontier of schedules for the next stage. We use a poset to remove the schedules that are worse in terms of both loop depth and memory depth. The poset-based pruning mechanism ensures that pruned schedules contain linear loop nest schedules (with no branches), including the default TACO schedule, as long as there are no other schedules with a scalar auxiliary memory and lesser time complexity. This guarantees that we end up with a superior schedule when compared to the default schedule, if such a schedule exists. The memory depth heuristic we use in Stage 1 (Section 5.2) ensures that we do not prune schedules that are likely to have lesser loop depth than the fused simple linear loop nest schedules.

The poset-based pruning mechanism can be formally written as follows. This stage removes a schedule from the set of schedules (received from Stage 1) if there exists ∈ such that

where ( ) and ( ) are the loop and memory depths of the schedule , respectively. This type of pruning ensures that we do not remove schedules that are likely to be better in the Pareto frontier of schedules. We allocate each schedule to a different (time, memory) bucket at the end of this stage. In other words, the schedules in a bucket have the same iteration time and memory complexity but differ in the order of loops.

Stage 3: Poset of Time-and Memory-Complexity-Based Pruning Using an SMT-Solver In some cases, the user (e.g., performance engineer) may know during the compile-time some information about the loop bounds or sparsities of the tensors used in the computation. For example, if it is a graph neural network computation, the user may know that the feature size of the nodes is in the range of [16,256], or the graph size is in the order of 10 and the sparsity of the graph is in the range of [0.001, 0.01]. Then, they can provide those ranges, and the auto-scheduler can use an SMT solver to reason about the time and auxiliary memory complexities of the schedules using symbolic cost expressions that it builds for every schedule ( § 3). Note that we assume that the tensors have uniform sparsities.

There are three types of constraints that we can provide the SMT solver; Listing 1. Formulation of the Pareto frontier using the SMT solver.

(1) Range constraints: the range in which dense loop bounds and sparsities can vary (e.g., line 2 in Listing 1). ( 2) Inferred constraints: non-zero values in a sparse tensor are always less than the number of elements in its dense representation, and the non-affine loop that iterates through the sparse tensor will vary between 0 and the dense loop bound that defines the sparse tensor (e.g., line 3 in Listing 1). (3) User-defined constraints: other special constraints that the user may know about the loop bounds or sparsities. For instance, the user may know that one loop bound is twice of another (e.g., line 4 in Listing 1).

After providing the constraints known at compile-time, we check if one schedule is dominated by at least another schedule in terms of both iteration time and auxiliary memory complexity; if so, we remove that schedule from the Pareto frontier (See line 10 in Listing 1). In other words, the system removes a schedule ( ) if there exists at least one schedule ( ) in the schedule space such that for all possible loop bounds and sparsities, the time and memory complexities of are worse than or equal to and there is no set of loop bounds and sparsities for which the time and memory complexities of are better than . The system does not remove a schedule if there exists at least one set values of loop bounds and sparsities for which the time and memory complexities of are better than . This gurantees that the system does not over-prune the schedules. This procedure can be formally written as follows. Let user-defined constraints of loop bounds and sparsities be Φ, let Z3 be the SMT solver, and the schedules from Stage 2 be . Provide Φ to Z3. Remove from if ∃ ∈ s.t. ∃ loop bounds and sparsities s.t. Here, ( ) and ( ) denote the symbolic iteration time and auxiliary memory complexities of the schedule , respectively. An auto-scheduler could have this stage alone without the previous stage in Section 5.2. But, when any of the above conditions are , the Z3 solver takes a long time to return. The previous stage compares a lot more schedules than this stage. Therefore, using a poset-based pruning strategy with absolute depth values, which is computationally efficient, is beneficial compared to using an SMT solver alone. Nonetheless, this stage is important because we can further reduce the number of schedules evaluated at run-time with the information available at compile-time.

Furthermore, the user could use this stage alone by removing both the memory depth-based and poset-based pruning stages. Although this removes the dependency on using memory depth as a heuristic to prune the schedules, it takes a long time for the solver to prune the schedule and does not work when the number of generated schedules is large.

Stage 4: Evaluated Time-Complexity-Based Pruning At the first stage of filtration at run-time, we evaluate the symbolic cost expressions with real values available at the run-time and select schedules that have the least iteration time complexity such that the auxiliary memory requirement is less than 50% of the last level cache (LLC) from the Pareto frontier. We take 50% as a rough margin for the selection criterion, assuming that 50% of LLC is available for the other input and output tensors in the computation. Multiple schedules with the same iteration time complexity can exist due to the same branched loop nest structure with different loop reorderings. These schedules are then passed to the next stage for further pruning.

Stage 5: Evaluated Cache-Access-Based Pruning At the second stage of run-time filtration of schedules, we prune the schedules based on cache behavior. We have included this stage here for completeness, and it is not our primary focus. Many remaining schedules may share equivalent time and memory complexity due to loop reorderings, such as the loop orders , , and , in the outer loops of Figure 2b. Since some schedules have the same time and memory complexity, if one of those schedules is not filtered away by previous stages, both of them will remain unpruned. Thus, we have a simple model that assigns a cache access cost to each schedule. This cache model takes two criteria into account. One, it looks at the leaves in the BIG and the leaf loop index. If the leaf loop index is and if a tensor in the expression at that leaf branch has as the last index (e.g., ( , )) or index is not present in the tensor (e.g., ( , )), then the cost of access is zero. If is present and not in the last accessed index (e.g., ( , )), then we take the cache access cost as since elements are accessed locations apart. We assign costs to all the leaves in the BIG and sum those to calculate a final cache access cost. We consider the last index of the leaves in the BIG because it has the highest impact on temporal and spatial cache locality. Two, we give precedence to the schedules that have the same index order as the loop order in the iteration graph. For instance, if tensors in the computation have ( , ) and ( , ), it would favor the loop order , , over , , . If both these criteria are the same for two schedules, we randomly pick one of them.

We assess SparseAuto using a collection of sparse tensor kernels in comparison to the schedules from TACO. We compare the results of SparseAuto with Pigeon [Ahrens et al. 2022] and SpTTN-Cyclops [Kanakagari and Solomonik 2023] qualitatively and quantitatively when applicable.

Experimental Setup. We conducted the experiments on a machine with four Non-Uniform Memory Access (NUMA) nodes of Intel(R) Xeon(R) CPU E5-4650 8-core processor (32-cores in total), operating at 2.70 GHz, with 32KB L1 data cache, 256KB L2 cache per core, and 80MB LLC shared between 4 NUMA nodes. Code compilation utilized GCC 11.4.0 with with optimization flags -O3 --ffast-math. The process involved a warm-up run, followed by 31 executions of the kernel computation. The results reported are the median values, accompanied by the corresponding standard deviation across the 31 runs. Parallel executions were performed on 32 threads using OpenMP.

Datasets. In the evaluation, we employ numerous real-world tensors sourced from the SuiteSparse Collection [Davis and Hu 2011], Network Repository [Rossi et al. 2015], Formidable Repository of Open Sparse Tensors and Tools [Smith et al. 2017], and the 1998 DARPA Intrusion Detection Evaluation Dataset [Lippmann et al. 2000]. The tensors and matrices used in the evaluation are shown in Table 1. These tensors span a wide range of sizes and sparsities. Sparse inputs to the kernels used the Compressed Sparse Fiber (CSF) format.

Kernels. We compare the performance of SparseAuto and TACO using kernels in Table 2. The kernel naming conventions are as follows: ⟨ , ⟩ indicates that the kernel is a combination of and , and the kernel can be decomposed into these two sub-kernels, each capable of being executed sequentially. The evaluation incorporates various combinations of the following kernels.

Sampled Dense-Dense Matrix Multiplication and Sparse Matrix-Matrix Multiplication are used in graph neural networks. In this context, the operation is used in computing attention values along the edges of a graph, then is used after the operation to transform the feature vector of each node, and the operation is used for multiplication with a weight matrix [Dias et al. 2022]. ⟨3 ⟩ Tensor-Times Matrix Contractions are used in Tucker Decompositions [Tucker 1966]. Matrizied Tensor Times Khatri-Rao product (

) is used in sparse computations such as signal processing and computer vision [Choi et al. 2018]. Sparse Tensor Times Matrix (

) operation is used in data mining and data analytics applications and is a sub-computation in Tucker Decomposition [Tucker 1966].

Number of schedules and overheads of each stage. Table 3 shows the schedule counts of each stage in the pruning pipeline, and Table 4 shows the corresponding execution times for each of these stages. These tables expose a correlation between execution times and the number of schedules to process. Also,

Table 4 shows that Depth Poset-based pruning (Stage 2) helps to save on expensive SMT work (Stage 3) because the column Stage 3 (Skipping Stage 2) is always longer than Stage 2 and 3 combined.

Table 2 shows the selected kernels after the compile-time pruning stages. Table 3 shows the number of schedules after each stage in the pruning pipeline. Furthermore, we bypass the second stage in the pruning pipeline and directly apply the SMT solver-based pruning in Stage 3 to the output from Stage 1. We observe that for kernel 2 , the number of schedules spared when Stage 2 is bypassed is 32 compared to the 8 schedules spared with Stage 2. For other kernels, the number of final schedules is the same with or without Stage 2. This indicates the effectiveness of the Depth Poset-based pruning in Stage 2. We also see that some of the schedules in Stage 2 are pruned in Stage 3, indicating the effectiveness of the SMT solver-based pruning in Stage 3.

Figure 9 shows the execution times and speedups of the selected schedules against the default TACO schedule. We observe orders of magnitude better performance compared to TACO. Although we do not reason about the effects of parallel execution in our auto-scheduler, we report the parallel performance of schedules by parallelizing the outer loops using OpenMP for completeness. We observe that parallel executions of the schedules have similar gains over TACO. We report only the serial execution times for ⟨ , ⟩ because the output of the kernel is sparse.

Comparison with Pigeon [Ahrens et al. 2022]. Pigeon introduces an auto-scheduler based on the time complexity of the schedule. They explore the search space with data layout transforms and transposes of sparse tensors, targetting an o ine schedule selection. Given A , they would consider schedules having both A , and A , with corresponding index orders of , and , in Table 6. Auxiliary memory requirements for different kernels. The right four columns show the number of elements required to store the intermediate results. The '-' label shows that the given schedule is not available with the given framework. pruning. Their system does not optimize for both time and auxiliary memory complexities, operates completely o ine, executes multiple schedules at the last stage, in other words they design the system for complete o ine schedule selection, their search space exploration algorithm does not explore the schedules with multi-level branch nests, and lacks the use of user-defined constraints at compile time for search space pruning. Although they introduce a good cost model, their system has the disadvantages described in Section 3. Their framework, evaluated on TACO, faces limitations in supporting intermediate temporaries with more than one dimension and includes schedules with data layout transformations in their auto-scheduler.

SpTTN-Cyclops [Kanakagari and Solomonik 2023] presents another auto-scheduler for sparse tensor contractions, offering a fully automated framework without user intervention in schedule selection. Unlike our approach, they do not emphasize poset-based pruning and opt for minimum loop depth schedules and then a maximum number of dense loops, which may not always be optimal, as discussed in Section 3. They lack support for user-defined constraints using an SMT solver to analyze schedule complexities for search space pruning. Their run-time loop generation algorithm affects evaluation time, while our method minimizes the number of schedules evaluated during run time.

Auto-scheduling is a challenging problem due to the vast number of potential schedules -ranging from thousands to hundreds of thousands -for a given computation. Factors such as time complexity, memory usage, cache behavior, and parallelism must all be considered. Most systems rely on heuristic-based approaches or empirical evaluations to identify optimal schedules. We advocate for a systematic approach that dedicates considerable time to o ine schedule generation and analysis; by investing hours in this process, most schedules can be eliminated, leaving only a few for evaluation at run-time. This viewpoint suggests that scheduling as a service (SchaaS) could benefit expensive computations, particularly scientific workloads, and lead to faster compute times and more efficient resource utilization.

We have introduced SparseAuto, a framework for recursive loop nest restructuring, including a scheduling language for sparse tensor contractions. An auto-scheduler for sparse tensor contractions is implemented, leveraging the said scheduling language to generate schedules. Among the many factors influencing schedule performance, we focus on two machine-independent criteria: time complexity and auxiliary memory usage, arising from variations in loop structures within sparse tensor contractions. SparseAuto employs a poset-based approach to prune the search space and utilizes an SMT solver for analyzing the symbolic cost of a schedule. Our findings demonstrate that SparseAuto delivers noteworthy performance enhancements.