[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) {
+    if (i >= oldNumDims) {
+      if (operands[i] && isValidDim(operands[i]) &&
+          !isValidSymbol(operands[i])) {
+        // This is a valid dim that appears as a symbol, legalize it.
+        symRemapping[i - oldNumDims] =
+            getAffineDimExpr(oldNumDims + nextDim++, context);
+        remappedDims.push_back(operands[i]);
+      } else {
+        symRemapping[i - oldNumDims] = getAffineSymbolExpr(nextSym++, context);
+        resultOperands.push_back(operands[i]);
+      }
+    } else {
+      resultOperands.push_back(operands[i]);
+    }
+  }
+
+  resultOperands.insert(resultOperands.begin() + oldNumDims,
----------------
bondhugula wrote:

Good idea, thanks. Done.

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