[PATCH] D75247: [mlir] fix wrong symbol order in AffineApplyNormalizer

[Alex Zinenko via Phabricator via llvm-commits]

This revision was automatically updated to reflect the committed changes.
Closed by commit rG54e5600e4d28: [mlir] fix wrong symbol order in AffineApplyNormalizer (authored by ftynse).

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D75247/new/

https://reviews.llvm.org/D75247

Files:
  mlir/include/mlir/Dialect/AffineOps/AffineOps.h
  mlir/lib/Dialect/AffineOps/AffineOps.cpp
  mlir/test/Dialect/AffineOps/canonicalize.mlir


Index: mlir/test/Dialect/AffineOps/canonicalize.mlir
===================================================================
--- mlir/test/Dialect/AffineOps/canonicalize.mlir
+++ mlir/test/Dialect/AffineOps/canonicalize.mlir
@@ -17,7 +17,7 @@
 // Affine maps for test case: compose_affine_maps_dependent_loads
 // CHECK-DAG: [[MAP9:#map[0-9]+]] = affine_map<(d0) -> (d0 + 3)>
 // CHECK-DAG: [[MAP10:#map[0-9]+]] = affine_map<(d0) -> (d0 * 3)>
-// CHECK-DAG: [[MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 7) ceildiv 3)>
+// CHECK-DAG: [[MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 3) ceildiv 3)>
 // CHECK-DAG: [[MAP12:#map[0-9]+]] = affine_map<(d0) -> (d0 * 7 - 49)>
 
 // Affine maps for test case: compose_affine_maps_diamond_dependency
@@ -187,9 +187,9 @@
 
         // Swizzle %x00, %x01 and %c3, %c7
         %x10 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 * s1)>
-           (%x01, %x00)[%c7, %c3]
+           (%x01, %x00)[%c3, %c7]
         %x11 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 ceildiv s0)>
-           (%x01, %x00)[%c7, %c3]
+           (%x01, %x00)[%c3, %c7]
 
         // CHECK-NEXT: [[I2A:%[0-9]+]] = affine.apply [[MAP12]](%{{.*}})
         // CHECK-NEXT: [[I2B:%[0-9]+]] = affine.apply [[MAP11]](%{{.*}})
@@ -568,3 +568,29 @@
   }
   return
 }
+
+// -----
+
+// Reproducer for PR45031. This used to fold into an incorrect map because
+// symbols were concatenated in the wrong order during map folding. Map
+// composition places the symbols of the original map before those of the map
+// it is composed with, e.g. A.compose(B) will first have all symbols of A,
+// then all symbols of B.
+
+#map1 = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
+#map2 = affine_map<(d0)[s0] -> (1024, -d0 + s0)>
+
+// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (1024, s1 * -1024 + s0)>
+
+// CHECK: func @rep(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index)
+func @rep(%arg0 : index, %arg1 : index) -> index {
+  // CHECK-NOT: constant
+  %c0 = constant 0 : index
+  %c1024 = constant 1024 : index
+  // CHECK-NOT: affine.apply
+  %0 = affine.apply #map1(%arg0)[%c1024, %c0]
+
+  // CHECK: affine.min #[[MAP]]()[%[[ARG1]], %[[ARG0]]]
+  %1 = affine.min #map2(%0)[%arg1]
+  return %1 : index
+}
Index: mlir/lib/Dialect/AffineOps/AffineOps.cpp
===================================================================
--- mlir/lib/Dialect/AffineOps/AffineOps.cpp
+++ mlir/lib/Dialect/AffineOps/AffineOps.cpp
@@ -536,8 +536,12 @@
           auxiliaryExprs.push_back(renumberOneDim(t));
         } else {
           // c. The mathematical composition of AffineMap concatenates symbols.
-          //    We do the same for symbol operands.
-          concatenatedSymbols.push_back(t);
+          //    Note that the map composition will put symbols already present
+          //    in the map before any symbols coming from the auxiliary map, so
+          //    we insert them before any symbols that are due to renumbering,
+          //    and after the proper symbols we have seen already.
+          concatenatedSymbols.insert(
+              std::next(concatenatedSymbols.begin(), numProperSymbols++), t);
         }
       }
     }
Index: mlir/include/mlir/Dialect/AffineOps/AffineOps.h
===================================================================
--- mlir/include/mlir/Dialect/AffineOps/AffineOps.h
+++ mlir/include/mlir/Dialect/AffineOps/AffineOps.h
@@ -654,6 +654,10 @@
   SmallVector<Value, 8> reorderedDims;
   SmallVector<Value, 8> concatenatedSymbols;
 
+  /// The number of symbols in concatenated symbols that belong to the original
+  /// map as opposed to those concatendated during map composition.
+  unsigned numProperSymbols;
+
   AffineMap affineMap;
 
   /// Used with RAII to control the depth at which AffineApply are composed
@@ -667,7 +671,7 @@
   }
   static constexpr unsigned kMaxAffineApplyDepth = 1;
 
-  AffineApplyNormalizer() { affineApplyDepth()++; }
+  AffineApplyNormalizer() : numProperSymbols(0) { affineApplyDepth()++; }
 
 public:
   ~AffineApplyNormalizer() { affineApplyDepth()--; }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D75247.247004.patch
Type: text/x-patch
Size: 4042 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200227/6159428a/attachment.bin>