Extend map canonicalization to propagate constant operands

- 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

diff --git a/mlir/examples/Linalg/Linalg4/Example.cpp b/mlir/examples/Linalg/Linalg4/Example.cpp
index ed439f2..405eb68 100644 (file)
--- a/mlir/examples/Linalg/Linalg4/Example.cpp
+++ b/mlir/examples/Linalg/Linalg4/Example.cpp
@@ -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

diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
index 7db3fa0..9d1fd9e 100644 (file)
--- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp
+++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
@@ -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,16 +668,27 @@ 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]) {
-      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]));
-      } else {
-        symRemapping[i] = it->second;
-      }
+    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]));
+    } else {
+      symRemapping[i] = it->second;
     }
   }
   *map =

diff --git a/mlir/test/AffineOps/canonicalize.mlir b/mlir/test/AffineOps/canonicalize.mlir
index 2a93274..8b0c5e6 100644 (file)
--- a/mlir/test/AffineOps/canonicalize.mlir
+++ b/mlir/test/AffineOps/canonicalize.mlir
@@ -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
@@ -411,4 +413,4 @@ func @constant_fold_bounds(%N : index) {
     "foo"(%k, %c3) : (index, index) -> ()
   }
   return
-}
\ No newline at end of file
+}

diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp
index 529aaa0..367355c 100644 (file)
--- a/mlir/test/EDSC/builder-api-test.cpp
+++ b/mlir/test/EDSC/builder-api-test.cpp
@@ -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

diff --git a/mlir/test/Linalg/tile_conv.mlir b/mlir/test/Linalg/tile_conv.mlir
index 5fdb28d..128161e 100644 (file)
--- a/mlir/test/Linalg/tile_conv.mlir
+++ b/mlir/test/Linalg/tile_conv.mlir
@@ -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]](%{{.*}})

diff --git a/mlir/test/Transforms/dma-generate.mlir b/mlir/test/Transforms/dma-generate.mlir
index c5f2d4f..0ca3455 100644 (file)
--- a/mlir/test/Transforms/dma-generate.mlir
+++ b/mlir/test/Transforms/dma-generate.mlir
@@ -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>
 

diff --git a/mlir/test/Transforms/loop-fusion-slice-computation.mlir b/mlir/test/Transforms/loop-fusion-slice-computation.mlir
index 1e5e448..59845e9 100644 (file)
--- a/mlir/test/Transforms/loop-fusion-slice-computation.mlir
+++ b/mlir/test/Transforms/loop-fusion-slice-computation.mlir
@@ -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>
     }
@@ -142,4 +142,4 @@ func @slice_loop_nest_with_smaller_outer_trip_count() {
     }
   }
   return
-}
\ No newline at end of file
+}

diff --git a/mlir/test/Transforms/loop-fusion.mlir b/mlir/test/Transforms/loop-fusion.mlir
index 4fa8d51..c97e3df 100644 (file)
--- a/mlir/test/Transforms/loop-fusion.mlir
+++ b/mlir/test/Transforms/loop-fusion.mlir
@@ -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

diff --git a/mlir/test/Transforms/pipeline-data-transfer.mlir b/mlir/test/Transforms/pipeline-data-transfer.mlir
index ce266d5..88fc5ca 100644 (file)
--- a/mlir/test/Transforms/pipeline-data-transfer.mlir
+++ b/mlir/test/Transforms/pipeline-data-transfer.mlir
@@ -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
 }