[Mlir-commits] [mlir] [MLIR][Affine] Fix affine indices generation in normalize-memrefs (PR #214961)

[llvmlistbot at llvm.org]

https://github.com/RohithPariki created https://github.com/llvm/llvm-project/pull/214961

This fixes an issue in NormalizeMemRefs where dereferencing ops that don't implement AffineMapAccessInterface received incorrectly composed indices. The indices generated for non-AffineMapAccessInterface ops were incorrectly passed through composition with the new layout map which resulted in swapped or corrupted dimensions when the new map was the identity.

>From b6cfe1683da860cd82d5ab8388990e7369790f5f Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Sat, 8 Aug 2026 17:02:45 +0530
Subject: [PATCH] [MLIR][Affine] Fix affine indices generation in
 normalize-memrefs

This fixes an issue in NormalizeMemRefs where dereferencing ops that don't implement AffineMapAccessInterface received incorrectly composed indices. The indices generated for non-AffineMapAccessInterface ops were incorrectly passed through composition with the new layout map which resulted in swapped or corrupted dimensions when the new map was the identity.
---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp | 66 ++++++++++++-------------
 1 file changed, 31 insertions(+), 35 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7043083298615..846d38aa087ed 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1250,16 +1250,15 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
   // Append 'remapOutputs' to 'newMapOperands'.
   newMapOperands.append(remapOutputs.begin(), remapOutputs.end());
 
-  // Create new fully composed AffineMap for new op to be created.
-  assert(newMapOperands.size() == newMemRefRank);
-  auto newMap = builder.getMultiDimIdentityMap(newMemRefRank);
-  fullyComposeAffineMapAndOperands(&newMap, &newMapOperands);
-  newMap = simplifyAffineMap(newMap);
-  canonicalizeMapAndOperands(&newMap, &newMapOperands);
-  // Remove any affine.apply's that became dead as a result of composition.
-  for (Value value : affineApplyOps)
-    if (value.use_empty())
-      value.getDefiningOp()->erase();
+  AffineMap newMap;
+  if (affMapAccInterface) {
+    // Create new fully composed AffineMap for new op to be created.
+    assert(newMapOperands.size() == newMemRefRank);
+    newMap = builder.getMultiDimIdentityMap(newMemRefRank);
+    fullyComposeAffineMapAndOperands(&newMap, &newMapOperands);
+    newMap = simplifyAffineMap(newMap);
+    canonicalizeMapAndOperands(&newMap, &newMapOperands);
+  }
 
   OperationState state(op->getLoc(), op->getName());
   // Construct the new operation using this memref.
@@ -1271,21 +1270,7 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
   state.operands.push_back(newMemRef);
 
   // Insert the new memref map operands.
-  if (affMapAccInterface) {
-    state.operands.append(newMapOperands.begin(), newMapOperands.end());
-  } else {
-    // In the case of dereferencing ops not implementing
-    // AffineMapAccessInterface, we need to apply the values of `newMapOperands`
-    // to the `newMap` to get the correct indices.
-    for (unsigned i = 0; i < newMemRefRank; i++) {
-      state.operands.push_back(AffineApplyOp::create(
-          builder, op->getLoc(),
-          AffineMap::get(newMap.getNumDims(), newMap.getNumSymbols(),
-                         newMap.getResult(i)),
-          newMapOperands));
-    }
-  }
-
+  state.operands.append(newMapOperands.begin(), newMapOperands.end());
   // Insert the remaining operands unmodified.
   unsigned oldMapNumInputs = oldMapOperands.size();
   state.operands.append(op->operand_begin() + memRefOperandPos + 1 +
@@ -1296,22 +1281,33 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
   for (auto result : op->getResults())
     state.types.push_back(result.getType());
 
-  // Add attribute for 'newMap', other Attributes do not change.
-  auto newMapAttr = AffineMapAttr::get(newMap);
-  for (auto namedAttr : op->getAttrs()) {
-    if (affMapAccInterface &&
-        namedAttr.getName() ==
-            affMapAccInterface.getAffineMapAttrForMemRef(oldMemRef).getName())
-      state.attributes.push_back({namedAttr.getName(), newMapAttr});
-    else
+  if (affMapAccInterface) {
+    // Add attribute for 'newMap', other Attributes do not change.
+    auto newMapAttr = AffineMapAttr::get(newMap);
+    for (auto namedAttr : op->getAttrs()) {
+      if (namedAttr.getName() == affMapAccInterface.getMapAttrName()) {
+        state.attributes.push_back(
+            builder.getNamedAttr(namedAttr.getName(), newMapAttr));
+      } else {
+        state.attributes.push_back(namedAttr);
+      }
+    }
+  } else {
+    for (auto namedAttr : op->getAttrs()) {
       state.attributes.push_back(namedAttr);
+    }
   }
 
   // Create the new operation.
-  auto *repOp = builder.create(state);
-  op->replaceAllUsesWith(repOp);
+  Operation *newOp = builder.create(state);
+  op->replaceAllUsesWith(newOp);
   op->erase();
 
+  // Remove any affine.apply's that became dead as a result of composition.
+  for (Value value : affineApplyOps)
+    if (value.use_empty())
+      value.getDefiningOp()->erase();
+
   return success();
 }