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