squiid/llvm
1
0
Fork 0
Code Issues Pull requests Activity

Extend map canonicalization to propagate constant operands

Browse source 
- extend canonicalizeMapAndOperands to propagate constant operands into
  the map's expressions (and thus drop those operands).
- canonicalizeMapAndOperands previously only dropped duplicate and
  unused operands; however, operands that were constants were
  retained.

This change makes IR maps/expressions generated by various
utilities/passes even simpler; also makes some of the test checks more
accurate and simpler -- for eg., 0' instead of symbol(%{{.*}}).

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes tensorflow/mlir#107

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/107 from bondhugula:canonicalize-maps c889a51486d14fbf7db489f224f881e7e1ff7d72
PiperOrigin-RevId: 266085289
This commit is contained in:
Uday Bondhugula 2019-08-29 01:13:01 -07:00 committed by A. Unique TensorFlower
parent bc2a543225
commit 4bb6f8ecdb
9 changed files with 89 additions and 75 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
mlir/examples/Linalg/Linalg4/Example.cpp
View file

@ -80,8 +80,8 @@ TEST_FUNC(matmul_tiled_loops) {
// CHECK: affine.for %{{.*}} = 0 to (d0) -> (d0)(%[[M]]) step 8 {
// CHECK: affine.for %{{.*}} = 0 to (d0) -> (d0)(%[[N]]) step 9 {
// CHECK: affine.for %{{.*}} = 0 to (d0) -> (d0)(%[[K]]) {
// CHECK: affine.for %{{.*}} = max (d0)[s0] -> (s0, d0)(%{{.*}})[%{{.*}}] to min (d0)[s0] -> (s0, d0 + 8)(%{{.*}})[%[[M]]] {
// CHECK: affine.for %{{.*}} = max (d0)[s0] -> (s0, d0)(%{{.*}})[%{{.*}}] to min (d0)[s0] -> (s0, d0 + 9)(%{{.*}})[%[[N]]] {
// CHECK: affine.for %{{.*}} = max (d0) -> (0, d0)(%{{.*}}) to min (d0)[s0] -> (s0, d0 + 8)(%{{.*}})[%[[M]]] {
// CHECK: affine.for %{{.*}} = max (d0) -> (0, d0)(%{{.*}}) to min (d0)[s0] -> (s0, d0 + 9)(%{{.*}})[%[[N]]] {
// CHECK-NEXT: %{{.*}} = cmpi "eq", %{{.*}}, %{{.*}} : index
// CHECK-NEXT: %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
// CHECK-NEXT: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : f32

22
mlir/lib/Dialect/AffineOps/AffineOps.cpp
View file

@ -602,7 +602,7 @@ canonicalizePromotedSymbols(AffineMap *map,
for (unsigned i = 0, e = map->getNumInputs(); i != e; ++i) {
if (i < map->getNumDims()) {
if (isValidSymbol((*operands)[i])) {
// This is a valid symbols that appears as a dim, canonicalize it.
// This is a valid symbol that appears as a dim, canonicalize it.
dimRemapping[i] = getAffineSymbolExpr(oldNumSyms + nextSym++, context);
remappedSymbols.push_back((*operands)[i]);
} else {
@ -653,6 +653,7 @@ void mlir::canonicalizeMapAndOperands(
unsigned nextDim = 0;
for (unsigned i = 0, e = map->getNumDims(); i != e; ++i) {
if (usedDims[i]) {
// Remap dim positions for duplicate operands.
auto it = seenDims.find((*operands)[i]);
if (it == seenDims.end()) {
dimRemapping[i] = getAffineDimExpr(nextDim++, context);
@ -667,18 +668,29 @@ void mlir::canonicalizeMapAndOperands(
SmallVector<AffineExpr, 8> symRemapping(map->getNumSymbols());
unsigned nextSym = 0;
for (unsigned i = 0, e = map->getNumSymbols(); i != e; ++i) {
if (usedSyms[i]) {
if (!usedSyms[i])
continue;
// Handle constant operands (only needed for symbolic operands since
// constant operands in dimensional positions would have already been
// promoted to symbolic positions above).
IntegerAttr operandCst;
if (matchPattern((*operands)[i + map->getNumDims()],
m_Constant(&operandCst))) {
symRemapping[i] =
getAffineConstantExpr(operandCst.getValue().getSExtValue(), context);
continue;
}
// Remap symbol positions for duplicate operands.
auto it = seenSymbols.find((*operands)[i + map->getNumDims()]);
if (it == seenSymbols.end()) {
symRemapping[i] = getAffineSymbolExpr(nextSym++, context);
resultOperands.push_back((*operands)[i + map->getNumDims()]);
seenSymbols.insert(std::make_pair((*operands)[i + map->getNumDims()],
symRemapping[i]));
seenSymbols.insert(
std::make_pair((*operands)[i + map->getNumDims()], symRemapping[i]));
} else {
symRemapping[i] = it->second;
}
}
}
*map =
map->replaceDimsAndSymbols(dimRemapping, symRemapping, nextDim, nextSym);
*operands = resultOperands;

44
mlir/test/AffineOps/canonicalize.mlir
View file

@ -5,25 +5,27 @@
// CHECK-DAG: [[MAP1:#map[0-9]+]] = (d0) -> (d0 + 1)
// Affine maps for test case: compose_affine_maps_1dto2d_with_symbols
// CHECK-DAG: [[MAP4:#map[0-9]+]] = (d0)[s0] -> (d0 - s0)
// CHECK-DAG: [[MAP6:#map[0-9]+]] = (d0)[s0] -> (d0 * 2 - s0 + 1)
// CHECK-DAG: [[MAP7:#map[0-9]+]] = (d0)[s0, s1] -> (d0 * 2 + s0 - s1)
// CHECK-DAG: [[MAP4:#map[0-9]+]] = (d0) -> (d0 - 4)
// CHECK-DAG: [[MAP4b:#map[0-9]+]] = (d0) -> (d0 - 7)
// CHECK-DAG: [[MAP7:#map[0-9]+]] = (d0) -> (d0 * 2 - 3)
// CHECK-DAG: [[MAP7a:#map[0-9]+]] = (d0) -> (d0 * 2 + 1)
// Affine map for test case: compose_affine_maps_d2_tile
// CHECK-DAG: [[MAP8:#map[0-9]+]] = (d0, d1)[s0] -> ((d0 ceildiv s0) * s0 + d1 mod s0)
// CHECK-DAG: [[MAP8:#map[0-9]+]] = (d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)
// CHECK-DAG: [[MAP8a:#map[0-9]+]] = (d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)
// Affine maps for test case: compose_affine_maps_dependent_loads
// CHECK-DAG: [[MAP9:#map[0-9]+]] = (d0)[s0] -> (d0 + s0)
// CHECK-DAG: [[MAP10:#map[0-9]+]] = (d0)[s0] -> (d0 * s0)
// CHECK-DAG: [[MAP11:#map[0-9]+]] = (d0)[s0, s1] -> ((d0 + s1) ceildiv s0)
// CHECK-DAG: [[MAP12:#map[0-9]+]] = (d0)[s0] -> ((d0 - s0) * s0)
// CHECK-DAG: [[MAP9:#map[0-9]+]] = (d0) -> (d0 + 3)
// CHECK-DAG: [[MAP10:#map[0-9]+]] = (d0) -> (d0 * 3)
// CHECK-DAG: [[MAP11:#map[0-9]+]] = (d0) -> ((d0 + 7) ceildiv 3)
// CHECK-DAG: [[MAP12:#map[0-9]+]] = (d0) -> (d0 * 7 - 49)
// Affine maps for test case: compose_affine_maps_diamond_dependency
// CHECK-DAG: [[MAP13A:#map[0-9]+]] = (d0) -> ((d0 + 6) ceildiv 8)
// CHECK-DAG: [[MAP13B:#map[0-9]+]] = (d0) -> ((d0 * 4 - 4) floordiv 3)
// Affine maps for test case: partial_fold_map
// CHECK-DAG: [[MAP15:#map[0-9]+]] = ()[s0, s1] -> (s0 - s1)
// CHECK-DAG: [[MAP15:#map[0-9]+]] = ()[s0] -> (s0 - 42)
// Affine maps for test cases: symbolic_composition_*
// CHECK-DAG: [[map_symbolic_composition_a:#map[0-9]+]] = ()[s0] -> (s0 * 512)
@ -89,20 +91,20 @@ func @compose_affine_maps_1dto2d_with_symbols() {
%c4 = constant 4 : index
%x0 = affine.apply (d0)[s0] -> (d0 - s0) (%i0)[%c4]
// CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP4]](%{{.*}})[%{{.*}}]
// CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP4]](%{{.*}})
// CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I0]], [[I0]]{{\]}}
%v0 = load %0[%x0, %x0] : memref<4x4xf32>
// Test load[%x0, %x1] with symbol %c4 captured by '%x0' map.
%x1 = affine.apply (d0) -> (d0 + 1) (%i0)
%y1 = affine.apply (d0, d1) -> (d0+d1) (%x0, %x1)
// CHECK-NEXT: [[I1:%[0-9]+]] = affine.apply [[MAP6]](%{{.*}})[%{{.*}}]
// CHECK-NEXT: [[I1:%[0-9]+]] = affine.apply [[MAP7]](%{{.*}})
// CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I1]], [[I1]]{{\]}}
%v1 = load %0[%y1, %y1] : memref<4x4xf32>
// Test load[%x1, %x0] with symbol %c4 captured by '%x0' map.
%y2 = affine.apply (d0, d1) -> (d0 + d1) (%x1, %x0)
// CHECK-NEXT: [[I2:%[0-9]+]] = affine.apply [[MAP6]](%{{.*}})[%{{.*}}]
// CHECK-NEXT: [[I2:%[0-9]+]] = affine.apply [[MAP7]](%{{.*}})
// CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I2]], [[I2]]{{\]}}
%v2 = load %0[%y2, %y2] : memref<4x4xf32>
@ -110,7 +112,7 @@ func @compose_affine_maps_1dto2d_with_symbols() {
%c5 = constant 5 : index
%x2 = affine.apply (d0)[s0] -> (d0 + s0) (%i0)[%c5]
%y3 = affine.apply (d0, d1) -> (d0 + d1) (%x2, %x0)
// CHECK: [[I3:%[0-9]+]] = affine.apply [[MAP7]](%{{.*}})[%{{.*}}, %{{.*}}]
// CHECK: [[I3:%[0-9]+]] = affine.apply [[MAP7a]](%{{.*}})
// CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I3]], [[I3]]{{\]}}
%v3 = load %0[%y3, %y3] : memref<4x4xf32>
}
@ -138,8 +140,8 @@ func @compose_affine_maps_2d_tile() {
((d0 * s0) + d2) (%x0, %x1, %x2, %x3)[%c4, %c8]
%x41 = affine.apply (d0, d1, d2, d3)[s0, s1] ->
((d1 * s1) + d3) (%x0, %x1, %x2, %x3)[%c4, %c8]
// CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP8]](%{{.*}}, %{{.*}})[%{{.*}}]
// CHECK: [[I1:%[0-9]+]] = affine.apply [[MAP8]](%{{.*}}, %{{.*}})[%{{.*}}]
// CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP8]](%{{.*}}, %{{.*}})
// CHECK: [[I1:%[0-9]+]] = affine.apply [[MAP8a]](%{{.*}}, %{{.*}})
// CHECK-NEXT: [[L0:%[0-9]+]] = load %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
%v0 = load %0[%x40, %x41] : memref<16x32xf32>
@ -170,9 +172,9 @@ func @compose_affine_maps_dependent_loads() {
%x02 = affine.apply (d0, d1, d2)[s0, s1] -> (d2 * s0)
(%i0, %i1, %i2)[%c3, %c7]
// CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP9]](%{{.*}})[%{{.*}}]
// CHECK: [[I1:%[0-9]+]] = affine.apply [[MAP4]](%{{.*}})[%{{.*}}]
// CHECK: [[I2:%[0-9]+]] = affine.apply [[MAP10]](%{{.*}})[%{{.*}}]
// CHECK: [[I0:%[0-9]+]] = affine.apply [[MAP9]](%{{.*}})
// CHECK: [[I1:%[0-9]+]] = affine.apply [[MAP4b]](%{{.*}})
// CHECK: [[I2:%[0-9]+]] = affine.apply [[MAP10]](%{{.*}})
// CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I0]], [[I1]]{{\]}}
%v0 = load %0[%x00, %x01] : memref<16x32xf32>
@ -189,8 +191,8 @@ func @compose_affine_maps_dependent_loads() {
%x11 = affine.apply (d0, d1)[s0, s1] -> (d1 ceildiv s0)
(%x01, %x00)[%c7, %c3]
// CHECK-NEXT: [[I2A:%[0-9]+]] = affine.apply [[MAP12]](%{{.*}})[%{{.*}}]
// CHECK-NEXT: [[I2B:%[0-9]+]] = affine.apply [[MAP11]](%{{.*}})[%{{.*}}, %{{.*}}]
// CHECK-NEXT: [[I2A:%[0-9]+]] = affine.apply [[MAP12]](%{{.*}})
// CHECK-NEXT: [[I2B:%[0-9]+]] = affine.apply [[MAP11]](%{{.*}})
// CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I2A]], [[I2B]]{{\]}}
%v3 = load %0[%x10, %x11] : memref<16x32xf32>
}
@ -261,7 +263,7 @@ func @partial_fold_map(%arg0: memref<index>, %arg1: index, %arg2: index) {
%c42 = constant 42 : index
%2 = affine.apply (d0, d1) -> (d0 - d1) (%arg1, %c42)
store %2, %arg0[] : memref<index>
// CHECK: [[X:%[0-9]+]] = affine.apply [[MAP15]]()[%{{.*}}, %{{.*}}]
// CHECK: [[X:%[0-9]+]] = affine.apply [[MAP15]]()[%{{.*}}]
// CHECK-NEXT: store [[X]], %{{.*}}
return

12
mlir/test/EDSC/builder-api-test.cpp
View file

@ -594,10 +594,10 @@ TEST_FUNC(tile_2d) {
// CHECK: affine.for %{{.*}} = (d0) -> (d0)(%[[ZERO]]) to (d0) -> (d0)(%[[M]]) step 512 {
// CHECK-NEXT: affine.for %{{.*}} = (d0) -> (d0)(%[[ZERO]]) to (d0) -> (d0)(%[[N]]) step 1024 {
// CHECK-NEXT: affine.for %{{.*}} = (d0) -> (d0)(%[[ZERO]]) to (d0) -> (d0)(%[[P]]) {
// CHECK-NEXT: affine.for %{{.*}} = max (d0)[s0] -> (s0, d0)(%{{.*}})[%[[ZERO]]] to min (d0)[s0] -> (s0, d0 + 512)(%{{.*}})[%[[M]]] step 16 {
// CHECK-NEXT: affine.for %{{.*}} = max (d0)[s0] -> (s0, d0)(%{{.*}})[%[[ZERO]]] to min (d0)[s0] -> (s0, d0 + 1024)(%{{.*}})[%[[N]]] step 32 {
// CHECK-NEXT: affine.for %{{.*}} = max (d0, d1)[s0] -> (s0, d0, d1)(%{{.*}}, %{{.*}})[%[[ZERO]]] to min (d0, d1)[s0] -> (s0, d0 + 1024, d1 + 32)(%{{.*}}, %{{.*}})[%[[N]]] {
// CHECK-NEXT: affine.for %{{.*}} = max (d0, d1)[s0] -> (s0, d0, d1)(%{{.*}}, %{{.*}})[%[[ZERO]]] to min (d0, d1)[s0] -> (s0, d0 + 512, d1 + 16)(%{{.*}}, %{{.*}})[%[[M]]] {
// CHECK-NEXT: affine.for %{{.*}} = max (d0) -> (0, d0)(%{{.*}}) to min (d0)[s0] -> (s0, d0 + 512)(%{{.*}})[%[[M]]] step 16 {
// CHECK-NEXT: affine.for %{{.*}} = max (d0) -> (0, d0)(%{{.*}}) to min (d0)[s0] -> (s0, d0 + 1024)(%{{.*}})[%[[N]]] step 32 {
// CHECK-NEXT: affine.for %{{.*}} = max (d0, d1) -> (0, d0, d1)(%{{.*}}, %{{.*}}) to min (d0, d1)[s0] -> (s0, d0 + 1024, d1 + 32)(%{{.*}}, %{{.*}})[%[[N]]] {
// CHECK-NEXT: affine.for %{{.*}} = max (d0, d1) -> (0, d0, d1)(%{{.*}}, %{{.*}}) to min (d0, d1)[s0] -> (s0, d0 + 512, d1 + 16)(%{{.*}}, %{{.*}})[%[[M]]] {
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
// CHECK-NEXT: {{.*}} = addf {{.*}}, {{.*}} : f32
@ -608,8 +608,8 @@ TEST_FUNC(tile_2d) {
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %{{.*}} = (d0) -> (d0)(%[[ZERO]]) to (d0) -> (d0)(%[[P]]) {
// CHECK-NEXT: affine.for %{{.*}} = max (d0)[s0] -> (s0, d0)(%{{.*}})[%[[ZERO]]] to min (d0)[s0] -> (s0, d0 + 512)(%{{.*}})[%[[M]]] {
// CHECK-NEXT: affine.for %{{.*}} = max (d0)[s0] -> (s0, d0)(%{{.*}})[%[[ZERO]]] to min (d0)[s0] -> (s0, d0 + 1024)(%{{.*}})[%[[N]]] {
// CHECK-NEXT: affine.for %{{.*}} = max (d0) -> (0, d0)(%{{.*}}) to min (d0)[s0] -> (s0, d0 + 512)(%{{.*}})[%[[M]]] {
// CHECK-NEXT: affine.for %{{.*}} = max (d0) -> (0, d0)(%{{.*}}) to min (d0)[s0] -> (s0, d0 + 1024)(%{{.*}})[%[[N]]] {
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
// CHECK-NEXT: {{.*}}= addf {{.*}}, {{.*}} : f32

4
mlir/test/Linalg/tile_conv.mlir
View file

@ -3,7 +3,7 @@
// TILE-23004-DAG: #[[UB0:.*]] = (d0) -> (d0 + 2)
// TILE-23004-DAG: #[[UB1:.*]] = (d0) -> (d0 + 3)
// TILE-23004-DAG: #[[UB2:.*]] = (d0) -> (d0 + 4)
// TILE-23004-DAG: #[[D0x30pS0x10:.*]] = (d0)[s0] -> (d0 * 30 + s0 * 10)
// TILE-23004-DAG: #[[D0x30pS0x10:.*]] = (d0) -> (d0 * 30)
// TILE-23004-DAG: #[[D0x30pS0x10p90:.*]] = (d0)[s0] -> (d0 * 30 + s0 * 10 + 90)
func @conv(%arg0: !linalg.view<?x?x?x?xf32>, %arg1: !linalg.view<?x?x?x?xf32>, %arg2: !linalg.view<?x?x?x?xf32>) {
linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], strides = [30, 40]} : !linalg.view<?x?x?x?xf32>, !linalg.view<?x?x?x?xf32>, !linalg.view<?x?x?x?xf32>
@ -24,7 +24,7 @@ func @conv(%arg0: !linalg.view<?x?x?x?xf32>, %arg1: !linalg.view<?x?x?x?xf32>, %
// TILE-23004: %[[FilterView:.*]] = linalg.subview %{{.*}}[%{{.*}}, %[[Z0]], %{{.*}}, %{{.*}}, %[[Z1]], %{{.*}}, %{{.*}}, %[[I2p4]], %{{.*}}, %{{.*}}, %[[K]], %{{.*}}] : !linalg.view<?x?x?x?xf32>
//
// TILE-23004: %[[I0p3:.*]] = affine.apply #[[UB0]](%{{.*}})
// TILE-23004: %[[I1:.*]] = affine.apply #[[D0x30pS0x10]](%{{.*}})[%{{.*}}]
// TILE-23004: %[[I1:.*]] = affine.apply #[[D0x30pS0x10]](%{{.*}})
// TILE-23004: %[[I1pStep:.*]] = affine.apply #[[D0x30pS0x10p90]](%{{.*}})[%[[PaddedInput0]]]
// TILE-23004: %[[SZ2:.*]] = linalg.dim %{{.*}}, 2 : !linalg.view<?x?x?x?xf32>
// TILE-23004: %[[I2p2:.*]] = affine.apply #[[UB2]](%{{.*}})

16
mlir/test/Transforms/dma-generate.mlir
View file

@ -159,7 +159,7 @@ func @loop_nest_high_d(%A: memref<512 x 32 x f32>,
// CHECK-NEXT: %{{.*}} = alloc() : memref<1x2xf32, 2>
// CHECK-NEXT: %{{.*}} = alloc() : memref<1xi32>
// Composition of the affine map for '%{{.*}}' causes '%{{.*}}' to be added as a symbol.
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, symbol(%{{.*}})], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256x8xf32>, memref<1x2xf32, 2>, memref<1xi32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, 0], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256x8xf32>, memref<1x2xf32, 2>, memref<1xi32>
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
// CHECK-NEXT: affine.for %{{.*}} = 0 to 8 {
// ...
@ -222,7 +222,7 @@ func @dma_constant_dim_access(%A : memref<100x100xf32>) {
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
affine.for %i = 0 to 100 {
affine.for %j = 0 to ()[s0] -> (s0) ()[%N] {
// CHECK: %{{.*}} = affine.load %{{.*}}[symbol(%{{.*}}) - 1, %{{.*}}] : memref<1x100xf32, 2>
// CHECK: %{{.*}} = affine.load %{{.*}}[0, %{{.*}}] : memref<1x100xf32, 2>
affine.load %A[%one, %j] : memref<100 x 100 x f32>
}
}
@ -245,12 +245,12 @@ func @dma_with_symbolic_accesses(%A : memref<100x100xf32>, %M : index) {
return
// CHECK: %{{.*}} = alloc() : memref<100x100xf32, 2>
// CHECK-NEXT: %{{.*}} = alloc() : memref<1xi32>
// CHECK-NEXT: affine.dma_start %{{.*}}[symbol(%{{.*}}), symbol(%{{.*}}) + 9], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}
// CHECK-NEXT: affine.dma_start %{{.*}}[0, symbol(%{{.*}}) + 9], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}}
// CHECK-NEXT: affine.for %{{.*}} = 0 to 100 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 100 {
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP_SYM_SHIFT]](%{{.*}}, %{{.*}})[%{{.*}}, %{{.*}}]
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}} + symbol(%{{.*}}) - 9] : memref<100x100xf32, 2>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<100x100xf32, 2>
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK: return
@ -401,7 +401,7 @@ func @dma_loop_straightline_interspersed() {
// CHECK-NEXT: %{{.*}} = alloc() : memref<1xi32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32>, memref<1xf32, 2>, memref<1xi32>
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[symbol(%{{.*}})] : memref<1xf32, 2>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[0] : memref<1xf32, 2>
// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>
// CHECK-NEXT: dealloc %{{.*}} : memref<1xf32, 2>
// CHECK-NEXT: %{{.*}} = alloc() : memref<254xf32, 2>
@ -418,8 +418,8 @@ func @dma_loop_straightline_interspersed() {
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32>, memref<256xf32, 2>, memref<1xi32>
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
// CHECK-NEXT: %{{.*}} = alloc() : memref<1xi32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[symbol(%{{.*}})] : memref<256xf32, 2>
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[symbol(%{{.*}})] : memref<256xf32, 2>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[255] : memref<256xf32, 2>
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<256xf32, 2>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32, 2>, memref<256xf32>, memref<1xi32>
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>
@ -449,7 +449,7 @@ func @dma_mixed_loop_blocks() {
// CHECK: affine.dma_start [[MEM]][%{{.*}}, %{{.*}}], [[BUF]][%{{.*}}, %{{.*}}], [[TAG]][%{{.*}}], %{{.*}} : memref<256x256xvector<8xf32>>, memref<256x256xvector<8xf32>, 2>, memref<1xi32>
// CHECK-NEXT: affine.dma_wait [[TAG]][%{{.*}}], %{{.*}} : memref<1xi32>
// CHECK-NEXT: affine.for %{{.*}} = 0 to 256 {
// CHECK: %{{.*}} = affine.load [[BUF]][symbol(%{{.*}}), symbol(%{{.*}})] : memref<256x256xvector<8xf32>, 2>
// CHECK: %{{.*}} = affine.load [[BUF]][0, 0] : memref<256x256xvector<8xf32>, 2>
// CHECK: affine.for %{{.*}} = 0 to 256 {
// CHECK-NEXT: %{{.*}} = affine.load [[BUF]][%{{.*}}, %{{.*}}] : memref<256x256xvector<8xf32>, 2>

12
mlir/test/Transforms/loop-fusion-slice-computation.mlir
View file

@ -75,14 +75,14 @@ func @slice_depth2_loop_nest_two_loads() {
%c0 = constant 0 : index
%cst = constant 7.000000e+00 : f32
affine.for %i0 = 0 to 16 {
// expected-remark@-1 {{slice ( src loop: 1, dst loop: 0, depth: 1 : insert point: (1, 1) loop bounds: [(d0)[s0] -> (d0), (d0)[s0] -> (d0 + 1)] [(d0)[s0] -> (0), (d0)[s0] -> (8)] )}}
// expected-remark@-2 {{slice ( src loop: 1, dst loop: 0, depth: 2 : insert point: (2, 1) loop bounds: [(d0, d1)[s0] -> (d0), (d0, d1)[s0] -> (d0 + 1)] [(d0, d1)[s0] -> (0), (d0, d1)[s0] -> (8)] )}}
// expected-remark@-1 {{slice ( src loop: 1, dst loop: 0, depth: 1 : insert point: (1, 1) loop bounds: [(d0) -> (d0), (d0) -> (d0 + 1)] [(d0) -> (0), (d0) -> (8)] )}}
// expected-remark@-2 {{slice ( src loop: 1, dst loop: 0, depth: 2 : insert point: (2, 1) loop bounds: [(d0, d1) -> (d0), (d0, d1) -> (d0 + 1)] [(d0, d1) -> (0), (d0, d1) -> (8)] )}}
affine.for %i1 = 0 to 16 {
affine.store %cst, %0[%i0, %i1] : memref<100x100xf32>
}
}
affine.for %i2 = 0 to 10 {
// expected-remark@-1 {{slice ( src loop: 0, dst loop: 1, depth: 1 : insert point: (1, 0) loop bounds: [(d0)[s0] -> (d0), (d0)[s0] -> (d0 + 1)] [(d0)[s0] -> (0), (d0)[s0] -> (8)] )}}
// expected-remark@-1 {{slice ( src loop: 0, dst loop: 1, depth: 1 : insert point: (1, 0) loop bounds: [(d0) -> (d0), (d0) -> (d0 + 1)] [(d0) -> (0), (d0) -> (8)] )}}
affine.for %i3 = 0 to 8 {
%1 = affine.load %0[%i2, %i3] : memref<100x100xf32>
}
@ -103,15 +103,15 @@ func @slice_depth2_loop_nest_two_stores() {
%c0 = constant 0 : index
%cst = constant 7.000000e+00 : f32
affine.for %i0 = 0 to 16 {
// expected-remark@-1 {{slice ( src loop: 1, dst loop: 0, depth: 1 : insert point: (1, 2) loop bounds: [(d0)[s0] -> (d0), (d0)[s0] -> (d0 + 1)] [(d0)[s0] -> (0), (d0)[s0] -> (8)] )}}
// expected-remark@-1 {{slice ( src loop: 1, dst loop: 0, depth: 1 : insert point: (1, 2) loop bounds: [(d0) -> (d0), (d0) -> (d0 + 1)] [(d0) -> (0), (d0) -> (8)] )}}
affine.for %i1 = 0 to 16 {
affine.store %cst, %0[%i0, %i1] : memref<100x100xf32>
}
affine.store %cst, %0[%i0, %c0] : memref<100x100xf32>
}
affine.for %i2 = 0 to 10 {
// expected-remark@-1 {{slice ( src loop: 0, dst loop: 1, depth: 1 : insert point: (1, 0) loop bounds: [(d0)[s0] -> (d0), (d0)[s0] -> (d0 + 1)] [(d0)[s0] -> (0), (d0)[s0] -> (16)] )}}
// expected-remark@-2 {{slice ( src loop: 0, dst loop: 1, depth: 2 : insert point: (2, 0) loop bounds: [(d0, d1)[s0] -> (d0), (d0, d1)[s0] -> (d0 + 1)] [(d0, d1)[s0] -> (0), (d0, d1)[s0] -> (16)] )}}
// expected-remark@-1 {{slice ( src loop: 0, dst loop: 1, depth: 1 : insert point: (1, 0) loop bounds: [(d0) -> (d0), (d0) -> (d0 + 1)] [(d0) -> (0), (d0) -> (16)] )}}
// expected-remark@-2 {{slice ( src loop: 0, dst loop: 1, depth: 2 : insert point: (2, 0) loop bounds: [(d0, d1) -> (d0), (d0, d1) -> (d0 + 1)] [(d0, d1) -> (0), (d0, d1) -> (16)] )}}
affine.for %i3 = 0 to 8 {
%1 = affine.load %0[%i2, %i3] : memref<100x100xf32>
}

10
mlir/test/Transforms/loop-fusion.mlir
View file

@ -657,7 +657,7 @@ func @R6_to_R2_reshape_square() -> memref<64x9xi32> {
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP3]](%{{.*}}, %{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP4]](%{{.*}}, %{{.*}})
// CHECK-NEXT: %{{.*}} = "foo"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index, index, index, index) -> i32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, ((%{{.*}} * 9 + %{{.*}}) mod 288) floordiv 144, (((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) floordiv 48, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) floordiv 16, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) mod 16, symbol(%{{.*}})] : memref<1x2x3x3x16x1xi32>
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, ((%{{.*}} * 9 + %{{.*}}) mod 288) floordiv 144, (((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) floordiv 48, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) floordiv 16, ((((%{{.*}} * 9 + %{{.*}}) mod 288) mod 144) mod 48) mod 16, 0] : memref<1x2x3x3x16x1xi32>
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP11]](%{{.*}}, %{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP12]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP13]](%{{.*}})
@ -841,8 +841,8 @@ func @fusion_at_depth0_not_currently_supported() {
// nest, and make the store in the slice store to the same element.
// CHECK-DAG: %{{.*}} = alloc() : memref<1xf32>
// CHECK: affine.for %{{.*}} = 0 to 10 {
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[symbol(%{{.*}})] : memref<1xf32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[symbol(%{{.*}})] : memref<1xf32>
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[0] : memref<1xf32>
// CHECK-NEXT: }
// CHECK-NEXT: return
return
@ -1228,10 +1228,10 @@ func @R3_to_R2_reshape() {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 3 {
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP0]](%{{.*}}, %{{.*}})
// CHECK-NEXT: %{{.*}} = "foo"(%{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index) -> i32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, 0, symbol(%{{.*}})] : memref<1x1x1xi32>
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, 0, 0] : memref<1x1x1xi32>
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP1]](%{{.*}}, %{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP2]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[0, 0, symbol(%{{.*}})] : memref<1x1x1xi32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[0, 0, 0] : memref<1x1x1xi32>
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: return

24
mlir/test/Transforms/pipeline-data-transfer.mlir
View file

@ -31,13 +31,13 @@ func @loop_nest_dma() {
// CHECK: %{{.*}} = alloc() : memref<256xf32>
// CHECK: %{{.*}} = alloc() : memref<2x32xf32, 1>
// CHECK-NEXT: %{{.*}} = alloc() : memref<2x1xf32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}} mod 2, %{{.*}}], %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}} : memref<256xf32>, memref<2x32xf32, 1>, memref<2x1xf32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}} mod 2, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<256xf32>, memref<2x32xf32, 1>, memref<2x1xf32>
// CHECK-NEXT: affine.for %{{.*}} = 1 to 8 {
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}} mod 2, %{{.*}}], %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}} : memref<256xf32>, memref<2x32xf32, 1>, memref<2x1xf32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}} mod 2, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<256xf32>, memref<2x32xf32, 1>, memref<2x1xf32>
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP_MINUS_1]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MOD_2]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MOD_2]](%{{.*}})
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}} : memref<2x1xf32>
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xf32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
// CHECK-NEXT: %{{.*}} = "compute"(%{{.*}}) : (f32) -> f32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
@ -48,7 +48,7 @@ func @loop_nest_dma() {
// CHECK-NEXT: %{{.*}} = affine.apply [[MAP_MINUS_1]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MOD_2]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[MOD_2]](%{{.*}})
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}} : memref<2x1xf32>
// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xf32>
// CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
// CHECK-NEXT: %{{.*}} = "compute"(%{{.*}}) : (f32) -> f32
// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}} mod 2, %{{.*}}] : memref<2x32xf32, 1>
@ -81,17 +81,17 @@ func @loop_step(%arg0: memref<512xf32>,
return
}
// CHECK: [[TAG:%[0-9]+]] = alloc() : memref<2x1xi32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[(%{{.*}} floordiv 4) mod 2, symbol(%{{.*}})], [[TAG]][(%{{.*}} floordiv 4) mod 2, symbol(%{{.*}})], %{{.*}} : memref<512xf32>, memref<2x4xf32, 1>, memref<2x1xi32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[(%{{.*}} floordiv 4) mod 2, 0], [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<512xf32>, memref<2x4xf32, 1>, memref<2x1xi32>
// CHECK-NEXT: affine.for %{{.*}} = 4 to 512 step 4 {
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[(%{{.*}} floordiv 4) mod 2, symbol(%{{.*}})], [[TAG]][(%{{.*}} floordiv 4) mod 2, symbol(%{{.*}})], %{{.*}} : memref<512xf32>, memref<2x4xf32, 1>, memref<2x1xi32>
// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[(%{{.*}} floordiv 4) mod 2, 0], [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<512xf32>, memref<2x4xf32, 1>, memref<2x1xi32>
// CHECK-NEXT: %{{.*}} = affine.apply [[REMAP_SHIFT_MINUS_4]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[FLOOR_MOD_2]](%{{.*}})
// CHECK: affine.dma_wait [[TAG]][(%{{.*}} floordiv 4) mod 2, symbol(%{{.*}})], %{{.*}} : memref<2x1xi32>
// CHECK: affine.dma_wait [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<2x1xi32>
// CHECK-NEXT: "compute"(%{{.*}}) : (index) -> ()
// CHECK-NEXT: }
// CHECK-NEXT: [[SHIFTED:%[0-9]+]] = affine.apply [[REMAP_SHIFT_MINUS_4]](%{{.*}})
// CHECK-NEXT: %{{.*}} = affine.apply [[FLOOR_MOD_2]]([[SHIFTED]])
// CHECK: affine.dma_wait [[TAG]][(%{{.*}} floordiv 4) mod 2, symbol(%{{.*}})], %{{.*}} : memref<2x1xi32>
// CHECK: affine.dma_wait [[TAG]][(%{{.*}} floordiv 4) mod 2, 0], %{{.*}} : memref<2x1xi32>
// CHECK-NEXT: "compute"(%{{.*}}) : (index) -> ()
// CHECK: return
// CHECK-NEXT: }
@ -282,7 +282,7 @@ func @dynamic_shape_dma_buffer(%arg0: memref<512 x 32 x f32>) {
// CHECK-NEXT: %{{.*}} = dim %{{.*}}, 0 : memref<?x?xf32, 2>
// CHECK-NEXT: %{{.*}} = dim %{{.*}}, 1 : memref<?x?xf32, 2>
// CHECK-NEXT: %{{.*}} = alloc(%{{.*}}, %{{.*}}) : memref<2x?x?xf32, 2>
// CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}}), symbol(%{{.*}})], %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}}
// CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0, 0], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}}
affine.for %kTT = 0 to 16 {
affine.dma_start %arg0[%zero, %zero], %Av[%zero, %zero], %tag[%zero], %num_elt :
memref<512 x 32 x f32>,
@ -291,10 +291,10 @@ func @dynamic_shape_dma_buffer(%arg0: memref<512 x 32 x f32>) {
}
return
// CHECK-NEXT: affine.for %{{.*}} = 1 to 16 {
// CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}}), symbol(%{{.*}})], %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}}
// CHECK: affine.dma_wait %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}} : memref<2x1xi32>
// CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}} mod 2, 0, 0], %{{.*}}[%{{.*}} mod 2, 0], %{{.*}}
// CHECK: affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xi32>
// CHECK: }
// CHECK: affine.dma_wait %{{.*}}[%{{.*}} mod 2, symbol(%{{.*}})], %{{.*}} : memref<2x1xi32>
// CHECK: affine.dma_wait %{{.*}}[%{{.*}} mod 2, 0], %{{.*}} : memref<2x1xi32>
// CHECK: return
}