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

[arnab-oss]

Summary

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.

[hanchung]

Thanks for fixing it!

mlir/lib/IR/AffineExpr.cpp

if (flatExprs[j] == 0)
       continue;
     // For dimensional expressions we set the index as <position number of the
     // dimension, 0>, as we want dimensional expressions to appear before
     // symbolic ones and products of dimensional and symbolic expressions
     // having the dimension with the same position number.
     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;
     // For symbolic expression set the index as <position number
     // 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));
   }
 
   // Denotes semi-affine product, modulo or division terms, which has been added
   // to the `indexToExpr` map.
   SmallVector<bool, 4> addedToMap(flatExprs.size() - numDims - numSymbols - 1,
                                   false);
if (rhs.isa<AffineConstantExpr>()) {
       // constructing. When rhs is constant, we place 0 in place of keyB.
       if (lhs.isa<AffineDimExpr>()) {
         lhsPos = lhs.cast<AffineDimExpr>().getPosition();
         std::pair<unsigned, signed> indexEntry(lhsPos, -1);
         addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()],
                  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);
       }
     } else if (lhs.isa<AffineDimExpr>()) {
       // For product/modulo/division expressions having lhs as dimension and rhs
       // as symbol, we order the terms in the semi-affine expression based on
       // the pair: <keyA, keyB> for expressions of the form dimension * symbol,
       // where keyA is the position number of the dimension and keyB is the
       // position number of the symbol.
       lhsPos = lhs.cast<AffineDimExpr>().getPosition();
       rhsPos = rhs.cast<AffineSymbolExpr>().getPosition();
       std::pair<unsigned, signed> indexEntry(lhsPos, rhsPos);
       addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr);
     } else {
       // For product/modulo/division expressions having both lhs and rhs as
       // symbol, we design indices as a pair: <keyA, keyB> for expressions
       // of the form dimension * symbol, where keyA is the position number of
       // 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;
   }
 
   // Constructing the simplified semi-affine sum of product/division/mod
   // expression from the flattened form in the desired sorted order of indices
   // of the various individual product/division/mod expressions.

mlir/test/Dialect/Affine/simplify-structures.mlir

 // CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$FLOORDIV]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]]
 // CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$SIMPLIFIED_CEILDIV_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]
 // CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]]
 
 // -----
 
 // 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) {
   %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> ((-(s0 - (s0 - s1) * s2) + s3) + (d0 + s0))>(%arg0)[%arg1, %arg2, %arg3, %arg4]
   %b = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + d1 * s1 + d1 + d0 * s0 + d1 * s0 + d2 * s1 + d2)>(%arg0, %arg1, %arg2)[%arg3, %arg4]
   return %a, %b : index, index
 }
 // CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG0]]]
 // CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SUM_OF_PRODUCTS]]()[%[[ARG3]], %[[ARG4]], %[[ARG0]], %[[ARG1]], %[[ARG2]]]
 // CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]]
 
 // -----
 
 // CHECK-DAG: #[[$SIMPLIFIED_MAP:.*]] = affine_map<()[s0, s1, s2, s3] -> ((-s0 + s2 + s3) mod (s0 + s1))>
 // CHECK-LABEL: func @semi_affine_simplification_euclidean_lemma
 // CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
 func.func @semi_affine_simplification_euclidean_lemma(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {
   %a = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1) - ((d0 + d1) floordiv (s0 - s1)) * (s0 - s1) - (d0 + d1) mod (s0 - s1))>(%arg0, %arg1)[%arg2, %arg3]
   %b = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1 - s0) - ((d0 + d1 - s0) floordiv (s0 + s1)) * (s0 + s1))>(%arg0, %arg1)[%arg2, %arg3]
   return %a, %b : index, 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
+}