Uniformize composition of AffineApplyOp by construction

This CL is the 5th on the path to simplifying AffineMap composition.
This removes the distinction between normalized single-result AffineMap and
more general composed multi-result map.

One nice byproduct of making the implementation driven by single-result is
that the multi-result extension is a trivial change: the implementation is
still single-result and we just use:

```
unsigned idx = getIndexOf(...);
map.getResult(idx);
```

This CL also fixes an AffineNormalizer implementation issue related to symbols.
Namely it stops performing substitutions on symbols in AffineNormalizer and
instead concatenates them all to be consistent with the call to
`AffineMap::compose(AffineMap)`. This latter call to `compose` cannot perform
simplifications of symbols coming from different maps based on positions only:
i.e. dims are applied and renumbered but symbols must be concatenated.

The only way to determine whether symbols from different AffineApply are the
same is to look at the concrete values. The canonicalizeMapAndOperands is thus
extended with behavior to support replacing operands that appear multiple
times.

Lastly, this CL demonstrates that the implementation is correct by rewriting
ComposeAffineMaps using only `makeComposedAffineApply`. The implementation
uses a matcher because AffineApplyOp are introduced as composed operations on
the fly instead of iteratively forwardSubstituting. For this purpose, a walker
would revisit freshly introduced AffineApplyOp. Regardless, ComposeAffineMaps
is scheduled to disappear, this CL replaces the implementation based on
iterative `forwardSubstitute` by a composed-by-construction
`makeComposedAffineApply`.
Remaining calls to `forwardSubstitute` will be removed in the next CL.

PiperOrigin-RevId: 228830443

5 files changed, 67 insertions, 74 deletions

diff --git a/mlir/test/Transforms/Vectorize/materialize.mlir b/mlir/test/Transforms/Vectorize/materialize.mlir
index b1c56fe143c..80458c75333 100644
--- a/mlir/test/Transforms/Vectorize/materialize.mlir
+++ b/mlir/test/Transforms/Vectorize/materialize.mlir
@@ -1,13 +1,10 @@
 // RUN: mlir-opt %s -materialize-vectors -vector-size=4 -vector-size=4 | FileCheck %s
 
-// CHECK-DAG: #[[D0D1D2D3TOD0:map[0-9]+]] = (d0, d1, d2, d3) -> (d0)
-// CHECK-DAG: #[[D0D1D2D3TOD1:map[0-9]+]] = (d0, d1, d2, d3) -> (d1)
-// CHECK-DAG: #[[D0D1D2D3TOD2:map[0-9]+]] = (d0, d1, d2, d3) -> (d2)
-// CHECK-DAG: #[[D0D1D2D3TOD3:map[0-9]+]] = (d0, d1, d2, d3) -> (d3)
+// CHECK-DAG: #[[ID1:map[0-9]+]] = (d0) -> (d0)
 // CHECK-DAG: #[[D0D1D2D3TOD1D0:map[0-9]+]] = (d0, d1, d2, d3) -> (d1, d0)
-// CHECK-DAG: #[[D0D1D2D3TOD1P1:map[0-9]+]] = (d0, d1, d2, d3) -> (d1 + 1)
-// CHECK-DAG: #[[D0D1D2D3TOD1P2:map[0-9]+]] = (d0, d1, d2, d3) -> (d1 + 2)
-// CHECK-DAG: #[[D0D1D2D3TOD1P3:map[0-9]+]] = (d0, d1, d2, d3) -> (d1 + 3)
+// CHECK-DAG: #[[D0P1:map[0-9]+]] = (d0) -> (d0 + 1)
+// CHECK-DAG: #[[D0P2:map[0-9]+]] = (d0) -> (d0 + 2)
+// CHECK-DAG: #[[D0P3:map[0-9]+]] = (d0) -> (d0 + 3)
 
 // CHECK-LABEL: func @materialize
 func @materialize(%M : index, %N : index, %O : index, %P : index) {
@@ -17,16 +14,16 @@ func @materialize(%M : index, %N : index, %O : index, %P : index) {
   // CHECK-NEXT:    for %i1 = 0 to %arg1 step 4 {
   // CHECK-NEXT:      for %i2 = 0 to %arg2 {
   // CHECK-NEXT:        for %i3 = 0 to %arg3 step 4 {
-  // CHECK-NEXT:          %[[a:[0-9]+]] = {{.*}}[[D0D1D2D3TOD0]](%i0, %i1, %i2, %i3)
-  // CHECK-NEXT:          %[[b:[0-9]+]] = {{.*}}[[D0D1D2D3TOD1]](%i0, %i1, %i2, %i3)
-  // CHECK-NEXT:          %[[c:[0-9]+]] = {{.*}}[[D0D1D2D3TOD2]](%i0, %i1, %i2, %i3)
-  // CHECK-NEXT:          %[[d:[0-9]+]] = {{.*}}[[D0D1D2D3TOD3]](%i0, %i1, %i2, %i3)
+  // CHECK-NEXT:          %[[a:[0-9]+]] = {{.*}}[[ID1]](%i0)
+  // CHECK-NEXT:          %[[b:[0-9]+]] = {{.*}}[[ID1]](%i1)
+  // CHECK-NEXT:          %[[c:[0-9]+]] = {{.*}}[[ID1]](%i2)
+  // CHECK-NEXT:          %[[d:[0-9]+]] = {{.*}}[[ID1]](%i3)
   // CHECK-NEXT:          vector_transfer_write {{.*}}, %0, %[[a]], %[[b]], %[[c]], %[[d]] {permutation_map: #[[D0D1D2D3TOD1D0]]} : vector<4x4xf32>, memref<?x?x?x?xf32>, index, index, index, index
-  // CHECK:          %[[b1:[0-9]+]] = {{.*}}[[D0D1D2D3TOD1P1]](%i0, %i1, %i2, %i3)
+  // CHECK:          %[[b1:[0-9]+]] = {{.*}}[[D0P1]](%i1)
   // CHECK:          vector_transfer_write {{.*}}, %0, {{.*}}, %[[b1]], {{.*}} {permutation_map: #[[D0D1D2D3TOD1D0]]} : vector<4x4xf32>, memref<?x?x?x?xf32>, index, index, index, index
-  // CHECK:          %[[b2:[0-9]+]] = {{.*}}[[D0D1D2D3TOD1P2]](%i0, %i1, %i2, %i3)
+  // CHECK:          %[[b2:[0-9]+]] = {{.*}}[[D0P2]](%i1)
   // CHECK:          vector_transfer_write {{.*}}, %0, {{.*}}, %[[b2]], {{.*}} {permutation_map: #[[D0D1D2D3TOD1D0]]} : vector<4x4xf32>, memref<?x?x?x?xf32>, index, index, index, index
-  // CHECK:          %[[b3:[0-9]+]] = {{.*}}[[D0D1D2D3TOD1P3]](%i0, %i1, %i2, %i3)
+  // CHECK:          %[[b3:[0-9]+]] = {{.*}}[[D0P3]](%i1)
   // CHECK:          vector_transfer_write {{.*}}, %0, {{.*}}, %[[b3]], {{.*}} {permutation_map: #[[D0D1D2D3TOD1D0]]} : vector<4x4xf32>, memref<?x?x?x?xf32>, index, index, index, index
   for %i0 = 0 to %M step 4 {
     for %i1 = 0 to %N step 4 {
diff --git a/mlir/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir b/mlir/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir
index c2f4009fa06..100822ea455 100644
--- a/mlir/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir
+++ b/mlir/test/Transforms/Vectorize/materialize_vectors_1d_to_1d.mlir
@@ -1,11 +1,11 @@
 // RUN: mlir-opt %s -vectorize -virtual-vector-size 32 --test-fastest-varying=0 -materialize-vectors -vector-size=8 | FileCheck %s
 
 // vector<32xf32> -> vector<8xf32>
-// CHECK-DAG: [[D0D1TOD0:#.*]] = (d0, d1) -> (d0)
+// CHECK-DAG: [[ID1:#.*]] = (d0) -> (d0)
 // CHECK-DAG: [[D0D1TOD1:#.*]] = (d0, d1) -> (d1)
-// CHECK-DAG: [[D0D1TOD1P8:#.*]] = (d0, d1) -> (d1 + 8)
-// CHECK-DAG: [[D0D1TOD1P16:#.*]] = (d0, d1) -> (d1 + 16)
-// CHECK-DAG: [[D0D1TOD1P24:#.*]] = (d0, d1) -> (d1 + 24)
+// CHECK-DAG: [[D0P8:#.*]] = (d0) -> (d0 + 8)
+// CHECK-DAG: [[D0P16:#.*]] = (d0) -> (d0 + 16)
+// CHECK-DAG: [[D0P24:#.*]] = (d0) -> (d0 + 24)
 
 // CHECK-LABEL: func @vector_add_2d
 func @vector_add_2d(%M : index, %N : index) -> f32 {
@@ -21,17 +21,17 @@ func @vector_add_2d(%M : index, %N : index) -> f32 {
   // CHECK-NEXT:     [[CST1:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
   // CHECK-NEXT:     [[CST2:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
   // CHECK-NEXT:     [[CST3:%.*]] = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
-  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[D0D1TOD1]]{{.*}}
+  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[ID1]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST0]], {{.*}}, [[VAL00]], [[VAL01]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0D1TOD1P8]]{{.*}}
+  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0P8]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST1]], {{.*}}, [[VAL10]], [[VAL11]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL20:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL21:%.*]] = affine_apply [[D0D1TOD1P16]]{{.*}}
+  // CHECK-NEXT:     [[VAL20:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL21:%.*]] = affine_apply [[D0P16]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST2]], {{.*}}, [[VAL20]], [[VAL21]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL30:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL31:%.*]] = affine_apply [[D0D1TOD1P24]]{{.*}}
+  // CHECK-NEXT:     [[VAL30:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL31:%.*]] = affine_apply [[D0P24]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST3]], {{.*}}, [[VAL30]], [[VAL31]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
   //
   for %i0 = 0 to %M {
@@ -47,17 +47,17 @@ func @vector_add_2d(%M : index, %N : index) -> f32 {
   // CHECK-NEXT:     [[CST1:%.*]] = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
   // CHECK-NEXT:     [[CST2:%.*]] = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
   // CHECK-NEXT:     [[CST3:%.*]] = constant splat<vector<8xf32>, 2.000000e+00> : vector<8xf32>
-  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[D0D1TOD1]]{{.*}}
+  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[ID1]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST0]], {{.*}}, [[VAL00]], [[VAL01]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0D1TOD1P8]]{{.*}}
+  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0P8]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST1]], {{.*}}, [[VAL10]], [[VAL11]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL20:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL21:%.*]] = affine_apply [[D0D1TOD1P16]]{{.*}}
+  // CHECK-NEXT:     [[VAL20:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL21:%.*]] = affine_apply [[D0P16]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST2]], {{.*}}, [[VAL20]], [[VAL21]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL30:%.*]] = affine_apply [[D0D1TOD0]]{{.*}}
-  // CHECK-NEXT:     [[VAL31:%.*]] = affine_apply [[D0D1TOD1P24]]{{.*}}
+  // CHECK-NEXT:     [[VAL30:%.*]] = affine_apply [[ID1]]{{.*}}
+  // CHECK-NEXT:     [[VAL31:%.*]] = affine_apply [[D0P24]]{{.*}}
   // CHECK-NEXT:     vector_transfer_write [[CST3]], {{.*}}, [[VAL30]], [[VAL31]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
   //
   for %i2 = 0 to %M {
diff --git a/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir b/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir
index fdf7749b73e..cfabd5bd025 100644
--- a/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir
+++ b/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_1d.mlir
@@ -1,11 +1,11 @@
 // RUN: mlir-opt %s -vectorize -virtual-vector-size 3 -virtual-vector-size 16 --test-fastest-varying=1 --test-fastest-varying=0 -materialize-vectors -vector-size=8 | FileCheck %s
 
 // vector<3x16xf32> -> vector<8xf32>
-// CHECK-DAG: [[D0D1TOD0:#.*]] = (d0, d1) -> (d0)
+// CHECK-DAG: [[ID1:#.*]] = (d0) -> (d0)
 // CHECK-DAG: [[D0D1TOD1:#.*]] = (d0, d1) -> (d1)
-// CHECK-DAG: [[D0D1TOD1P8:#.*]] = (d0, d1) -> (d1 + 8)
-// CHECK-DAG: [[D0D1TOD0P1:#.*]] = (d0, d1) -> (d0 + 1)
-// CHECK-DAG: [[D0D1TOD0P2:#.*]] = (d0, d1) -> (d0 + 2)
+// CHECK-DAG: [[D0P8:#.*]] = (d0) -> (d0 + 8)
+// CHECK-DAG: [[D0P1:#.*]] = (d0) -> (d0 + 1)
+// CHECK-DAG: [[D0P2:#.*]] = (d0) -> (d0 + 2)
 
 // CHECK-LABEL: func @vector_add_2d
 func @vector_add_2d(%M : index, %N : index) -> f32 {
@@ -23,23 +23,23 @@ func @vector_add_2d(%M : index, %N : index) -> f32 {
   // CHECK-NEXT:     {{.*}} = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
   // CHECK-NEXT:     {{.*}} = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
   // CHECK-NEXT:     {{.*}} = constant splat<vector<8xf32>, 1.000000e+00> : vector<8xf32>
-  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[D0D1TOD0]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[D0D1TOD1]](%i0, %i1)
+  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[ID1]](%i0)
+  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[ID1]](%i1)
   // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL00]], [[VAL01]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[D0D1TOD0]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0D1TOD1P8]](%i0, %i1)
+  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[ID1]](%i0)
+  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0P8]](%i1)
   // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL10]], [[VAL11]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL20:%.*]] = affine_apply [[D0D1TOD0P1]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL21:%.*]] = affine_apply [[D0D1TOD1]](%i0, %i1)
+  // CHECK-NEXT:     [[VAL20:%.*]] = affine_apply [[D0P1]](%i0)
+  // CHECK-NEXT:     [[VAL21:%.*]] = affine_apply [[ID1]](%i1)
   // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL20]], [[VAL21]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL30:%.*]] = affine_apply [[D0D1TOD0P1]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL31:%.*]] = affine_apply [[D0D1TOD1P8]](%i0, %i1)
+  // CHECK-NEXT:     [[VAL30:%.*]] = affine_apply [[D0P1]](%i0)
+  // CHECK-NEXT:     [[VAL31:%.*]] = affine_apply [[D0P8]](%i1)
   // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL30]], [[VAL31]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL40:%.*]] = affine_apply [[D0D1TOD0P2]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL41:%.*]] = affine_apply [[D0D1TOD1]](%i0, %i1)
+  // CHECK-NEXT:     [[VAL40:%.*]] = affine_apply [[D0P2]](%i0)
+  // CHECK-NEXT:     [[VAL41:%.*]] = affine_apply [[ID1]](%i1)
   // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL40]], [[VAL41]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
-  // CHECK-NEXT:     [[VAL50:%.*]] = affine_apply [[D0D1TOD0P2]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL51:%.*]] = affine_apply [[D0D1TOD1P8]](%i0, %i1)
+  // CHECK-NEXT:     [[VAL50:%.*]] = affine_apply [[D0P2]](%i0)
+  // CHECK-NEXT:     [[VAL51:%.*]] = affine_apply [[D0P8]](%i1)
   // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL50]], [[VAL51]] {permutation_map: [[D0D1TOD1]]} : vector<8xf32>
   for %i0 = 0 to %M {
     for %i1 = 0 to %N {
diff --git a/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir b/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir
index 58865df7a9b..82cfd85d2dc 100644
--- a/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir
+++ b/mlir/test/Transforms/Vectorize/materialize_vectors_2d_to_2d.mlir
@@ -1,10 +1,9 @@
 // RUN: mlir-opt %s -vectorize -virtual-vector-size 3 -virtual-vector-size 32 --test-fastest-varying=1 --test-fastest-varying=0 -materialize-vectors -vector-size=3 -vector-size=16 | FileCheck %s 
 
 // vector<3x32xf32> -> vector<3x16xf32>
-// CHECK-DAG: [[D0D1TOD0:#.*]] = (d0, d1) -> (d0)
-// CHECK-DAG: [[D0D1TOD1:#.*]] = (d0, d1) -> (d1)
-// CHECK-DAG: [[D0D1TOD0D1:#.*]] = (d0, d1) -> (d0, d1)
-// CHECK-DAG: [[D0D1TOD1P16:#.*]] = (d0, d1) -> (d1 + 16)
+// CHECK-DAG: [[ID1:#.*]] = (d0) -> (d0)
+// CHECK-DAG: [[ID2:#.*]] = (d0, d1) -> (d0, d1)
+// CHECK-DAG: [[D0P16:#.*]] = (d0) -> (d0 + 16)
 
 // CHECK-LABEL: func @vector_add_2d
 func @vector_add_2d(%M : index, %N : index) -> f32 {
@@ -18,12 +17,12 @@ func @vector_add_2d(%M : index, %N : index) -> f32 {
   // CHECK-NEXT:   for %i1 = 0 to %arg1 step 32 {
   // CHECK-NEXT:     {{.*}} = constant splat<vector<3x16xf32>, 1.000000e+00> : vector<3x16xf32>
   // CHECK-NEXT:     {{.*}} = constant splat<vector<3x16xf32>, 1.000000e+00> : vector<3x16xf32>
-  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[D0D1TOD0]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[D0D1TOD1]](%i0, %i1)
-  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL00]], [[VAL01]] {permutation_map: [[D0D1TOD0D1]]} : vector<3x16xf32>
-  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[D0D1TOD0]](%i0, %i1)
-  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0D1TOD1P16]](%i0, %i1)
-  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL10]], [[VAL11]] {permutation_map: [[D0D1TOD0D1]]} : vector<3x16xf32>
+  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[ID1]](%i0)
+  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[ID1]](%i1)
+  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL00]], [[VAL01]] {permutation_map: [[ID2]]} : vector<3x16xf32>
+  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[ID1]](%i0)
+  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0P16]](%i1)
+  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL10]], [[VAL11]] {permutation_map: [[ID2]]} : vector<3x16xf32>
   //
   for %i0 = 0 to %M {
     for %i1 = 0 to %N {
@@ -36,12 +35,12 @@ func @vector_add_2d(%M : index, %N : index) -> f32 {
   // CHECK-NEXT:   for %i3 = 0 to %arg1 step 32 {
   // CHECK-NEXT:     {{.*}} = constant splat<vector<3x16xf32>, 2.000000e+00> : vector<3x16xf32>
   // CHECK-NEXT:     {{.*}} = constant splat<vector<3x16xf32>, 2.000000e+00> : vector<3x16xf32>
-  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[D0D1TOD0]](%i2, %i3)
-  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[D0D1TOD1]](%i2, %i3)
-  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL00]], [[VAL01]] {permutation_map: [[D0D1TOD0D1]]} : vector<3x16xf32>
-  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[D0D1TOD0]](%i2, %i3)
-  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0D1TOD1P16]](%i2, %i3)
-  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL10]], [[VAL11]] {permutation_map: [[D0D1TOD0D1]]} : vector<3x16xf32>
+  // CHECK-NEXT:     [[VAL00:%.*]] = affine_apply [[ID1]](%i2)
+  // CHECK-NEXT:     [[VAL01:%.*]] = affine_apply [[ID1]](%i3)
+  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL00]], [[VAL01]] {permutation_map: [[ID2]]} : vector<3x16xf32>
+  // CHECK-NEXT:     [[VAL10:%.*]] = affine_apply [[ID1]](%i2)
+  // CHECK-NEXT:     [[VAL11:%.*]] = affine_apply [[D0P16]](%i3)
+  // CHECK-NEXT:     vector_transfer_write {{.*}}, {{.*}}, [[VAL10]], [[VAL11]] {permutation_map: [[ID2]]} : vector<3x16xf32>
   //
   for %i2 = 0 to %M {
     for %i3 = 0 to %N {
diff --git a/mlir/test/Transforms/Vectorize/normalize_maps.mlir b/mlir/test/Transforms/Vectorize/normalize_maps.mlir
index 05ad1da55db..ff20cae7e44 100644
--- a/mlir/test/Transforms/Vectorize/normalize_maps.mlir
+++ b/mlir/test/Transforms/Vectorize/normalize_maps.mlir
@@ -1,14 +1,11 @@
 // RUN: mlir-opt %s -vectorizer-test -normalize-maps |  FileCheck %s
 
-// CHECK-DAG: #[[ZERO:[a-zA-Z0-9]+]] = (d0) -> (0)
+// CHECK-DAG: #[[ZERO:[a-zA-Z0-9]+]] = () -> (0)
 // CHECK-DAG: #[[ID1:[a-zA-Z0-9]+]] = (d0) -> (d0)
 // CHECK-DAG: #[[D0TIMES2:[a-zA-Z0-9]+]] = (d0) -> (d0 * 2)
 // CHECK-DAG: #[[D0PLUSD1:[a-zA-Z0-9]+]] = (d0, d1) -> (d0 + d1)
 // CHECK-DAG: #[[MINSD0PLUSD1:[a-zA-Z0-9]+]] = (d0, d1) -> (d0 * -1 + d1)
 // CHECK-DAG: #[[D0MINUSD1:[a-zA-Z0-9]+]] = (d0, d1) -> (d0 - d1)
-// CHECK-DAG: #[[D0D1D2TOD0:[a-zA-Z0-9]+]] = (d0, d1, d2) -> (d0)
-// CHECK-DAG: #[[D0D1D2TOD1:[a-zA-Z0-9]+]] = (d0, d1, d2) -> (d1)
-// CHECK-DAG: #[[D0D1D2TOD2:[a-zA-Z0-9]+]] = (d0, d1, d2) -> (d2)
 
 // CHECK-LABEL: func @simple()
 func @simple() {
@@ -21,7 +18,7 @@ func @simple() {
   // CHECK-NEXT: for %i0 = 0 to 7
   // CHECK-NEXT:   {{.*}} affine_apply #[[ID1]](%i0)
   // CHECK-NEXT:   {{.*}} affine_apply #[[D0TIMES2]](%i0)
-  // CHECK-NEXT:   {{.*}} affine_apply #[[ZERO]](%i0)
+  // CHECK-NEXT:   {{.*}} affine_apply #[[ZERO]]()
 
   for %i1 = 0 to 7 {
     for %i2 = 0 to 42 {
@@ -53,9 +50,9 @@ func @simple() {
   // CHECK:      for %i3 = 0 to 16
   // CHECK-NEXT:   for %i4 = 0 to 47 step 2
   // CHECK-NEXT:     for %i5 = 0 to 78 step 16
-  // CHECK-NEXT:       {{.*}} affine_apply #[[D0D1D2TOD0]](%i3, %i4, %i5)
-  // CHECK-NEXT:       {{.*}} affine_apply #[[D0D1D2TOD1]](%i3, %i4, %i5)
-  // CHECK-NEXT:       {{.*}} affine_apply #[[D0D1D2TOD2]](%i3, %i4, %i5)
+  // CHECK-NEXT:       {{.*}} affine_apply #[[ID1]](%i3)
+  // CHECK-NEXT:       {{.*}} affine_apply #[[ID1]](%i4)
+  // CHECK-NEXT:       {{.*}} affine_apply #[[ID1]](%i5)
 
   return
 }