xref: /llvm-project/mlir/test/Transforms/parametric-tiling.mlir (revision d4fd20258f63d30be638b04f10eaa469707759f0)
1// RUN: mlir-opt -test-extract-fixed-outer-loops='test-outer-loop-sizes=7' %s | FileCheck %s --check-prefixes=COMMON,TILE_7
2// RUN: mlir-opt -test-extract-fixed-outer-loops='test-outer-loop-sizes=7,4' %s | FileCheck %s --check-prefixes=COMMON,TILE_74
3
4// COMMON-LABEL: @rectangular
5func.func @rectangular(%arg0: memref<?x?xf32>) {
6  %c2 = arith.constant 2 : index
7  %c44 = arith.constant 44 : index
8  %c1 = arith.constant 1 : index
9  // Range of the original loop:
10  //   (upper - lower + step - 1) / step
11  // where step is known to be %c1.
12  // COMMON:      %[[diff:.*]] = arith.subi %c44, %c2
13  // COMMON:      %[[adjustment:.*]] = arith.subi %c1, %c1_{{.*}}
14  // COMMON-NEXT: %[[diff_adj:.*]] = arith.addi %[[diff]], %[[adjustment]]
15  // COMMON-NEXT: %[[range:.*]] = arith.divui %[[diff_adj]], %c1
16
17  // Ceildiv to get the parametric tile size.
18  // COMMON:       %[[sum:.*]] = arith.addi %[[range]], %c6
19  // COMMON-NEXT:  %[[size:.*]] = arith.divui %[[sum]], %c7
20  // New outer step (original is %c1).
21  // COMMON-NEXT:      %[[step:.*]] = arith.muli %c1, %[[size]]
22
23  // Range of the second original loop
24  //   (upper - lower + step - 1) / step
25  // where step is known to be %c2.
26  // TILE_74:      %[[diff2:.*]] = arith.subi %c44, %c1
27  // TILE_74:      %[[adjustment2:.*]] = arith.subi %c2, %c1_{{.*}}
28  // TILE_74-NEXT: %[[diff2_adj:.*]] = arith.addi %[[diff2]], %[[adjustment2]]
29  // TILE_74-NEXT: %[[range2:.*]] = arith.divui %[[diff2_adj]], %c2
30
31  // Ceildiv to get the parametric tile size for the second original scf.
32  // TILE_74:      %[[sum2:.*]] = arith.addi %[[range2]], %c3
33  // TILE_74-NEXT: %[[size2:.*]] = arith.divui %[[sum2]], %c4
34  // New inner step (original is %c2).
35  // TILE_74-NEXT:     %[[step2:.*]] = arith.muli %c2, %[[size2]]
36
37  // Updated outer loop(s) use new steps.
38  // COMMON: scf.for %[[i:.*]] = %c2 to %c44 step %[[step]]
39  // TILE_74:scf.for %[[j:.*]] = %c1 to %c44 step %[[step2]]
40 scf.for %i = %c2 to %c44 step %c1 {
41    // Upper bound for the inner loop min(%i + %step, %c44).
42    // COMMON:      %[[stepped:.*]] = arith.addi %[[i]], %[[step]]
43    // COMMON-NEXT: %[[ub:.*]] = arith.minsi %c44, %[[stepped]]
44    //
45    // TILE_74:      %[[stepped2:.*]] = arith.addi %[[j]], %[[step2]]
46    // TILE_74-NEXT: %[[ub2:.*]] = arith.minsi %c44, %[[stepped2]]
47
48    // Created inner scf.
49    // COMMON:scf.for %[[ii:.*]] = %[[i]] to %[[ub:.*]] step %c1
50
51    // This loop is not modified in TILE_7 case.
52    // TILE_7: scf.for %[[j:.*]] = %c1 to %c44 step %c2
53    //
54    // But is modified in TILE_74 case.
55    // TILE_74:scf.for %[[jj:.*]] = %[[j]] to %[[ub2]] step %c2
56   scf.for %j = %c1 to %c44 step %c2 {
57      // The right iterator are used.
58      // TILE_7:  memref.load %arg0[%[[ii]], %[[j]]]
59      // TILE_74: memref.load %arg0[%[[ii]], %[[jj]]]
60      memref.load %arg0[%i, %j]: memref<?x?xf32>
61    }
62  }
63  return
64}
65
66// COMMON-LABEL: @triangular
67func.func @triangular(%arg0: memref<?x?xf32>) {
68  %c2 = arith.constant 2 : index
69  %c44 = arith.constant 44 : index
70  %c1 = arith.constant 1 : index
71  // Range of the original outer loop:
72  //   (upper - lower + step - 1) / step
73  // where step is known to be %c1.
74  // COMMON:      %[[diff:.*]] = arith.subi %c44, %c2
75  // COMMON:      %[[adjustment:.*]] = arith.subi %c1, %c1_{{.*}}
76  // COMMON-NEXT: %[[diff_adj:.*]] = arith.addi %[[diff]], %[[adjustment]]
77  // COMMON-NEXT: %[[range:.*]] = arith.divui %[[diff_adj]], %c1
78
79  // Ceildiv to get the parametric tile size.
80  // COMMON:       %[[sum:.*]] = arith.addi %[[range]], %c6
81  // COMMON-NEXT:  %[[size:.*]] = arith.divui %[[sum]], %c7
82  // New outer step (original is %c1).
83  // COMMON-NEXT:  %[[step:.*]] = arith.muli %c1, %[[size]]
84
85  // Constant adjustment for inner loop has been hoisted out.
86  // TILE_74:      %[[adjustment2:.*]] = arith.subi %c2, %c1_{{.*}}
87
88  // New outer scf.
89  // COMMON: scf.for %[[i:.*]] = %c2 to %c44 step %[[step]]
90
91  // Range of the original inner loop
92  //   (upper - lower + step - 1) / step
93  // where step is known to be %c2.
94  // TILE_74:      %[[diff2:.*]] = arith.subi %[[i]], %c1
95  // TILE_74-NEXT: %[[diff2_adj:.*]] = arith.addi %[[diff2]], %[[adjustment2]]
96  // TILE_74-NEXT: %[[range2:.*]] = arith.divui %[[diff2_adj]], %c2
97
98  // Ceildiv to get the parametric tile size for the second original scf.
99  // TILE_74:      %[[sum2:.*]] = arith.addi %[[range2]], %c3
100  // TILE_74-NEXT: %[[size2:.*]] = arith.divui %[[sum2]], %c4
101  // New inner step (original is %c2).
102  // TILE_74-NEXT:     %[[step2:.*]] = arith.muli %c2, %[[size2]]
103
104  // New inner scf.
105  // TILE_74:scf.for %[[j:.*]] = %c1 to %[[i]] step %[[step2]]
106 scf.for %i = %c2 to %c44 step %c1 {
107    // Upper bound for the inner loop min(%i + %step, %c44).
108    // COMMON:      %[[stepped:.*]] = arith.addi %[[i]], %[[step]]
109    // COMMON-NEXT: %[[ub:.*]] = arith.minsi %c44, %[[stepped]]
110    // TILE_74:      %[[stepped2:.*]] = arith.addi %[[j]], %[[step2]]
111    // TILE_74-NEXT: %[[ub2:.*]] = arith.minsi %[[i]], %[[stepped2]]
112    //
113    // Created inner scf.
114    // COMMON:scf.for %[[ii:.*]] = %[[i]] to %[[ub:.*]] step %c1
115
116    // This loop is not modified in TILE_7 case.
117    // TILE_7: scf.for %[[j:.*]] = %c1 to %[[ii]] step %c2
118    //
119    // But is modified in TILE_74 case.
120    // TILE_74:scf.for %[[jj:.*]] = %[[j]] to %[[ub2]] step %c2
121   scf.for %j = %c1 to %i step %c2 {
122      // The right iterator are used.
123      // TILE_7:  memref.load %arg0[%[[ii]], %[[j]]]
124      // TILE_74: memref.load %arg0[%[[ii]], %[[jj]]]
125      memref.load %arg0[%i, %j]: memref<?x?xf32>
126    }
127  }
128  return
129}
130