[Mlir-commits] [mlir] 5e0d291 - [MLIR] Fix bug in the method constructing semi affine expression from flattened form

[Uday Bondhugula]

Author: Arnab Dutta
Date: 2022-11-01T11:35:53+05:30
New Revision: 5e0d29103a3f9252b26669f6d4bdc9dc29cef1f1

URL: https://github.com/llvm/llvm-project/commit/5e0d29103a3f9252b26669f6d4bdc9dc29cef1f1
DIFF: https://github.com/llvm/llvm-project/commit/5e0d29103a3f9252b26669f6d4bdc9dc29cef1f1.diff

LOG: [MLIR] Fix bug in the method constructing semi affine expression from flattened form

Set proper offset to the second element of the index pair, so that
we do not have same index values for more than one local expression.

Reviewed By: springerm

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

Added: 
    

Modified: 
    mlir/lib/IR/AffineExpr.cpp
    mlir/test/Dialect/Affine/simplify-structures.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 991a644d9d8f1..e0f45470bf3bd 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -996,6 +996,8 @@ static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
     std::pair<unsigned, signed> indexEntry(j, -1);
     addEntry(indexEntry, flatExprs[j], getAffineDimExpr(j, context));
   }
+  // Ensure we do not have duplicate keys in `indexToExpr` map.
+  unsigned offset = 0;
   for (unsigned j = numDims; j < numDims + numSymbols; ++j) {
     if (flatExprs[j] == 0)
       continue;
@@ -1003,8 +1005,8 @@ static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
     // of the symbol, max(dimCount, symCount)> number,
     // as we want symbolic expressions with the same positional number to
     // appear after dimensional expressions having the same positional number.
-    std::pair<unsigned, signed> indexEntry(j - numDims,
-                                           std::max(numDims, numSymbols));
+    std::pair<unsigned, signed> indexEntry(
+        j - numDims, std::max(numDims, numSymbols) + offset++);
     addEntry(indexEntry, flatExprs[j],
              getAffineSymbolExpr(j - numDims, context));
   }
@@ -1041,8 +1043,8 @@ static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
                  expr);
       } else {
         lhsPos = lhs.cast<AffineSymbolExpr>().getPosition();
-        std::pair<unsigned, signed> indexEntry(lhsPos,
-                                               std::max(numDims, numSymbols));
+        std::pair<unsigned, signed> indexEntry(
+            lhsPos, std::max(numDims, numSymbols) + offset++);
         addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()],
                  expr);
       }
@@ -1063,7 +1065,8 @@ static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
       // the dimension and keyB is the position number of the symbol.
       lhsPos = lhs.cast<AffineSymbolExpr>().getPosition();
       rhsPos = rhs.cast<AffineSymbolExpr>().getPosition();
-      std::pair<unsigned, signed> indexEntry(lhsPos, rhsPos);
+      std::pair<unsigned, signed> indexEntry(
+          lhsPos, std::max(numDims, numSymbols) + offset++);
       addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr);
     }
     addedToMap[it.index()] = true;

diff  --git a/mlir/test/Dialect/Affine/simplify-structures.mlir b/mlir/test/Dialect/Affine/simplify-structures.mlir
index 6e7679c0652fd..903d11ea865fe 100644
--- a/mlir/test/Dialect/Affine/simplify-structures.mlir
+++ b/mlir/test/Dialect/Affine/simplify-structures.mlir
@@ -522,7 +522,7 @@ func.func @semiaffine_simplification_floordiv_and_ceildiv(%arg0: index, %arg1: i
 
 // Test simplification of product expressions.
 // CHECK-DAG:   #[[$PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 + (s0 - s1) * s2)>
-// CHECK-DAG:   #[[$SUM_OF_PRODUCTS:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s2 * s0 + s2 + s3 * s0 + s3 * s1 + s3 + s4 * s1 + s4)>
+// CHECK-DAG:   #[[$SUM_OF_PRODUCTS:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s2 + s2 * s0 + s3 + s3 * s0 + s3 * s1 + s4 + s4 * s1)>
 // CHECK-LABEL: func @semiaffine_simplification_product
 // CHECK-SAME:  (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
 func.func @semiaffine_simplification_product(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {
@@ -547,3 +547,13 @@ func.func @semi_affine_simplification_euclidean_lemma(%arg0: index, %arg1: index
 // CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : index
 // CHECK-NEXT: %[[RESULT:.*]] = affine.apply #[[$SIMPLIFIED_MAP]]()[%[[ARG2]], %[[ARG3]], %[[ARG0]], %[[ARG1]]]
 // CHECK-NEXT: return %[[ZERO]], %[[RESULT]]
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 mod 2 + (s0 floordiv 2) * s0)>
+// CHECK-LABEL: func @semiaffine_modulo
+func.func @semiaffine_modulo(%arg0: index) -> index {
+  %a = affine.apply affine_map<()[s0] -> (s0 mod 2 + (s0 floordiv 2) * s0)> ()[%arg0]
+  // CHECK: affine.apply #[[$MAP]]()[%{{.*}}]
+  return %a : index
+}