[Mlir-commits] [mlir] [MLIR][Affine] Fix affine.apply verifier and add functionality to demote invalid symbols to dims (PR #128289)

[Uday Bondhugula]

================
@@ -1359,13 +1368,62 @@ static void canonicalizePromotedSymbols(MapOrSet *mapOrSet,
 
   resultOperands.append(remappedSymbols.begin(), remappedSymbols.end());
   *operands = resultOperands;
-  *mapOrSet = mapOrSet->replaceDimsAndSymbols(dimRemapping, {}, nextDim,
-                                              oldNumSyms + nextSym);
+  *mapOrSet = mapOrSet->replaceDimsAndSymbols(
+      dimRemapping, /*symReplacements=*/{}, nextDim, oldNumSyms + nextSym);
 
   assert(mapOrSet->getNumInputs() == operands->size() &&
          "map/set inputs must match number of operands");
 }
 
+// A valid affine dimension may appear as a symbol in affine.apply operations.
+// This function canonicalizes symbols that are valid dims, but not valid
+// symbols into actual dims. Without such a legalization, the affine.apply will
+// be invalid. This method is the exact inverse of canonicalizePromotedSymbols.
+template <class MapOrSet>
+static void legalizeDemotedDims(MapOrSet *mapOrSet,
+                                SmallVectorImpl<Value> &operands) {
+  if (!mapOrSet || operands.empty())
+    return;
+
+  assert(mapOrSet->getNumInputs() == operands.size() &&
+         "map/set inputs must match number of operands");
+
+  auto *context = mapOrSet->getContext();
+  SmallVector<Value, 8> resultOperands;
+  resultOperands.reserve(operands.size());
+  SmallVector<Value, 8> remappedDims;
+  remappedDims.reserve(operands.size());
+  unsigned nextSym = 0;
+  unsigned nextDim = 0;
+  unsigned oldNumDims = mapOrSet->getNumDims();
+  SmallVector<AffineExpr, 8> symRemapping(mapOrSet->getNumSymbols());
+  for (unsigned i = 0, e = mapOrSet->getNumInputs(); i != e; ++i) {
----------------
bondhugula wrote:

Thanks, done.

https://github.com/llvm/llvm-project/pull/128289