Efficient Search-Space Pruning for
Integrated Fusion
and Tiling Transformations

Xiaoyang Gao, Sriram Krishnamoorthy,
Swarup Kumar Sahoo, Chi-Chung Lam,
P. Sadayappan
Ohio State University
Gerald Baumgartner, J. Ramanujam,
Louisiana State University
1

Introduction
 Integrated framework to determine a variety ofloop transformations:
Loop fusion
 Loop tiling Loop permutation
 Concrete performance models

 Reduction in the space of possible solutions

2

Context
 Tensor Contraction Engine (TCE):TCE A domain-specific compiler used in Quantum Chemistry.
Transform high-level math. specification to efficient parallel
programs optimized for target machines.

Input:- Sequence of tensor contraction expressions
Output:- Parallel Fortran code3

Four-index Transform
B(a, b, c, d ) = ∑C1(d , s) *C2(c, r) *C3(b, q) *C4(a, p) * A( p, q, r, s)
p,q,r , sOperation-minimal form

Producer-consumer relationship4Observations
Sequence of fully permutable loop nests
 Often, arrays are too large to fit into physical memory
 Array access expressions are loop indices In each contraction, indices form three disjoint groups,
each group appearing in exactly two array references
 C[i,j] += A[i,k] * B[k,j] T[i,j] += A[k,l] * B[i,j,k,l] A producer loop nest cannot be fused with consumer ifsummation index is the outermost loop in the producer.
5

Problem Statement
Objective: Given a tensor expression and machine parameters,
determine the appropriate loop transformations, and the position andordering of I/O placements to minimize disk I/O cost.

Problem Addressed:
 Several loop transformations are applied.
 Their effects on I/O cost are interrelated.
 Space of possible solutions too large to exhaustively searchApproach: Pruning of the search space to achieve better solution pereffort expended.
In this paper, we focus on the integration of loop fusion and tiling.

6

Operation Tree

Operation Tree: A binary
tree represents a sequenceof tensor contractions.
Leaf: Input arrays
 Root: Output array Interior node: Intermediateor output arrays, producedby the tensor contraction of
their immediate children Edge: Producer-consumer
relationship between tensorcontractions

B = SUM(T3*C1)

T3
T3 = SUM(T2*C2)

T2 = SUM(T1*C3)
T2

T1 = SUM(A*C4)AC1C2C3C47

Problem Statement
 Input : Operation Tree

 Output: Candidate loop structures

 Objective: Minimize number of loop structures to beconsidered while maximizing search space explored.

8

Fusion Enumeration Space A natural approach

All combinations of common loops in related loop
nests (producers and consumers in a contraction)
 Very large solution space.
 Key observation
Given any fused structure

A canonical fusion structure can be generated
All common loops in the loop nests are fused
All loops are tiled and tile sizes set appropriately

9

Two-index Transform
for ifor j,nT[n] += A[i,j]*C2[n,j]for m,nB[m,n] += T[n]*C1[m,i]for nfor j,iT[i] += A[i,j]*C2[n,j]for m,iB[m,n] += T[i]*C1[m,i]for i,nfor j
T += A[i,j]*C2[n,j]for mB[m,n] += T*C1[m,i]T[i,n] = A[i,j] * C2[n,j]B[m,n] = T[i,n] * C1[m,i]for it1, nt1
for j, it2, nt2
T[it2, nt2] += A[it1+it2, j] * C2[nt1+nt2, j]
for m, it2, nt2
B[m, nt1+nt2] += T[it2,nt2] * C1[m, it1+it2]

Fuse all common loops

10Two-index Transform (Contd.)for i
for j,nT[n] += A[i,j]*C2[n,j]for m,nB[m,n] += T[n]*C1[m,i]for it1, nt1=1
for j, it2=1, nt2
T[it2, nt2] += A[it1+it2, j] * C2[nt1+nt2, j]
for m, it2=1, nt2
B[m, nt1+nt2] += T[it2,nt2] * C1[m, it1+it2]

Fusion + tiling to reduce number of candidate loop structures11

Cut-point and Fused Sub-tree
To fuse or not-to-fuseCut-point: For a fusion structure, an intermediate node
not fused with its consumer, is a cut-point in theoperation tree.

Fused Sub-tree: Cut-points divide an operation tree into
several sub-trees. A sub-tree without any interior cutpoints is a fused sub-tree.

12Fused Sub-tree and Cut-point (4index)
Loop Structure:

T1 = SUM(A*C4)

for a,r,q,s,pT1(a,q,r,s)+=A(p,q,r,s)*C4(a,p)

B
B==SUM(T3*C1)
SUM(T3*C1)

A

for a,b
for r,sfor q
T2(r,s)+=T1(a,q,r,s)*C3(b,q)
for cT3(c,s)+=T2(r,s)*C2(c,r)
for c,d,s
B(a,b,c,d)+=T3(c,s)*C1(d,s)

C4T3 = SUM(T2*C2)T2 = SUM(T1*C3)

T1

C1

C2

C313

Integrated Framework
Input: Operation Tree
Procedure:
Operation Tree Partitioning
Loop Structures Enumeration
Intra-Tile Loop Placements
Disk I/O Placements and Orderings
Tile Size Selection
Code Generation
Output: Fortran Code14

Operation Tree Partitioning Partition the operation tree using cut-points

 Each intermediate tree node is potentially a cut-point Operation tree with M intermediate nodes – 2M

fusion structures

15

Fused Sub-tree Enumeration
Three choices for each contraction
 Fuse all loops common to any two of the threenodes involved in the contraction


The two producer nests and the consumer nest

Fusing the loops of the producer loop-nests places thesummation indices as the outermost Fusion structure cannot be extended – a cut-point

All fusion sub-structures to be enumerated are chains

16Fused Sub-tree Enumeration
 Dynamic programming solution to constructfusion structures hierarchically

At any interior node of operation tree,
Extend fusion structures of the producer nests to theconsumer or Fuse the loops of the producer and terminate the fusionstructure.
17

Loop Structure Enumeration
1.

Fusion sub-trees form a chain of contractions.

2.

All possible enumerations of loop structures parenthesization problem

3.

For each parenthesization, a maximally fused loop
structure is created by a recursive construction
procedure.
Maximally fused loop: Each loop nest in which two subnest
have as many common loops as possible.
18

Maximally fused loop structure1.

4index:B(a, b, c, d ) = ∑C1(d , s) *C2(c, r) *C3(b, q) *C4(a, p) * A( p, q, r, s)p, q, r , s

2.

Contraction sequence:
T 1(a, q, r, s) = ∑ C 4(a, p) * A(p, q, r, s)p

T 2(a, b, r, s) = ∑ C 3(b, q) * T 1(a, q, r, s)q

T 3(a, b, c, s) = ∑ C 2(c, r) * T 2(a, b, r, s)r

B(a, b, c, d) = ∑ C1(d, s) * T 3(a, b, c, s)
s

3.
4.

Contraction chain: T1 T2 T3 B
Parenthesizations: (T1(T2(T3B))), ((T1(T2T3))B), (T1((T2T3)B) ),(((T1T2)T3)B) , ((T1T2)(T3B)), (T1(T2(T3B)))

19

Maximally fused loop structure (Contd.)
Maximally fused loop structure for ((T1(T2T3))B):

5.

(T2T3)(T1(T2T3))

a,b,r,sa,r,sq (T2)c (T3)p,q (T1)q (T2)((T1(T2T3))B)a,sbrc (T3)p,q (T1)

b,c,d (B)bq (T2)c (T3)

20

Experimental Evaluation
 Determined the reduction in the number ofpossible loop structures before and after pruning.
 Evaluated on representative expressions fromthree quantum chemistry codes:
Four-index transform (4index)
 CCSD computation (CCSD)
 CCSDT computation (CCSDT)
21

Experimental Evaluation
Expressions

Total loopstructures

Loopstructures
after pruning

Reduction

4index

241

598%

CCSD

69

297%

CCSDT

182

598%
22Conclusions
 Partitioned an operation tree into fused sub-trees.

 Determined candidate loop structures as

parenthesizations of candidate fusion chains.
 Search space of possible loop structures isdrastically reduced.

23

Thank You!

24