[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 when either
lhs or rhs of a local expression is a dimensional identifier, so that
we do not have same index values for more than one local expression.

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

990–992

Please use LLVM/MLIR camel back style names.

[arnab-oss]

Changed variable names to camel back style

[hanchung]

thanks for the quick fix!

mlir/lib/IR/AffineExpr.cpp

auto addEntry = [&](std::pair<unsigned, signed> index, int64_t coefficient,
     coefficients.insert({index, coefficient});
     indexToExprMap.insert({index, expr});
   };
 
   // Design indices for dimensional or symbolic terms, and store the indices,
   // constant coefficient corresponding to the indices in `coefficients` map,
   // and affine expression corresponding to indices in `indexToExprMap` map.
 
-  for (unsigned j = 0; j < numDims; ++j) {
-    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;
+  unsigned offsetSym = 0;
+  signed offsetDim = -1;
   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) + offset++);
+        j - numDims, std::max(numDims, numSymbols) + offsetSym++);
     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);
for (const auto &it : llvm::enumerate(localExprs)) {
     }
     if (rhs.isa<AffineConstantExpr>()) {
       // For product/modulo/division expressions, when rhs of modulo/division
       // expression is constant, we put 0 in place of keyB, because we want
       // them to appear earlier in the semi-affine expression we are
       // 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);
+        std::pair<unsigned, signed> indexEntry(lhsPos, offsetDim--);
         addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()],
                  expr);
       } else {
         lhsPos = lhs.cast<AffineSymbolExpr>().getPosition();
         std::pair<unsigned, signed> indexEntry(
-            lhsPos, std::max(numDims, numSymbols) + offset++);
+            lhsPos, std::max(numDims, numSymbols) + offsetSym++);
         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, std::max(numDims, numSymbols) + offset++);
+          lhsPos, std::max(numDims, numSymbols) + offsetSym++);
       addEntry(indexEntry, flatExprs[numDims + numSymbols + it.index()], expr);
     }
     addedToMap[it.index()] = true;
   }
 
+  for (unsigned j = 0; j < numDims; ++j) {
+    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, offsetDim--);
+    addEntry(indexEntry, flatExprs[j], getAffineDimExpr(j, context));
+  }
+
   // 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.
   llvm::sort(indices);
   for (const std::pair<unsigned, unsigned> index : indices) {
     assert(indexToExprMap.lookup(index) &&
            "cannot find key in `indexToExprMap` map");
     expr = expr + indexToExprMap.lookup(index) * coefficients.lookup(index);

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

 
 // 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
 }
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0, s1, s2] -> (s2 mod 2 + (s1 floordiv 2) * 2 + ((s2 floordiv 2) * s0) * 2)>
+// CHECK-LABEL: func @semiaffine_modulo_dim
+func.func @semiaffine_modulo_dim(%arg0: index, %arg1: index, %arg2: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0, s1] -> (((d0 floordiv 2) * s0 + s1 floordiv 2) * 2 + d0 mod 2)> (%arg0)[%arg1, %arg2]
+  //CHECK: affine.apply #[[$MAP]]()[%{{.*}}, %{{.*}}, %{{.*}}]
+  return %a : index
+}