[llvm-branch-commits] [mlir] 3bcca6b - [MLIR] Fix affine_map compose with multi-symbols

[Uday Bondhugula via llvm-branch-commits]

Author: Chengji Yao
Date: 2021-01-02T06:57:16+05:30
New Revision: 3bcca6b12d4cbe9d4571e799899e9b956e4711bf

URL: https://github.com/llvm/llvm-project/commit/3bcca6b12d4cbe9d4571e799899e9b956e4711bf
DIFF: https://github.com/llvm/llvm-project/commit/3bcca6b12d4cbe9d4571e799899e9b956e4711bf.diff

LOG: [MLIR] Fix affine_map compose with multi-symbols

Fix bug: https://bugs.llvm.org/show_bug.cgi?id=46845

Differential Revision: https://reviews.llvm.org/D93831

Added: 
    

Modified: 
    mlir/lib/IR/AffineMap.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 634917fe0faf..b51636091aad 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -319,7 +319,7 @@ AffineMap AffineMap::compose(AffineMap map) {
   for (unsigned idx = 0; idx < numDims; ++idx) {
     newDims[idx] = getAffineDimExpr(idx, getContext());
   }
-  SmallVector<AffineExpr, 8> newSymbols(numSymbols);
+  SmallVector<AffineExpr, 8> newSymbols(numSymbols - numSymbolsThisMap);
   for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
     newSymbols[idx - numSymbolsThisMap] =
         getAffineSymbolExpr(idx, getContext());

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 0e452b0c6338..b185e1c8b707 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -24,6 +24,9 @@
 // CHECK-DAG: [[$MAP13A:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
 // CHECK-DAG: [[$MAP13B:#map[0-9]+]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
 
+// Affine maps for test case: compose_affine_maps_multiple_symbols
+// CHECK-DAG: [[$MAP14:#map[0-9]+]] = affine_map<()[s0, s1] -> (((s1 + s0) * 4) floordiv s0)>
+
 // Affine maps for test case: partial_fold_map
 // CHECK-DAG: [[$MAP15:#map[0-9]+]] = affine_map<()[s0] -> (s0 - 42)>
 
@@ -218,6 +221,15 @@ func @compose_affine_maps_diamond_dependency(%arg0: f32, %arg1: memref<4x4xf32>)
   return
 }
 
+// CHECK-LABEL: func @compose_affine_maps_multiple_symbols
+func @compose_affine_maps_multiple_symbols(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] -> (s0 + d0)> (%arg0)[%arg1]
+  %c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
+  %e = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%c)[%arg1]
+  // CHECK: [[I0:%[0-9]+]] = affine.apply [[$MAP14]]()[%{{.*}}, %{{.*}}]
+  return %e : index
+}
+
 // CHECK-LABEL: func @arg_used_as_dim_and_symbol
 func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index, %arg2: f32) {
   %c9 = constant 9 : index