[Mlir-commits] [mlir] [MLIR][Affine] Fix index mapping for memref load/store in NormalizeMemRefs (PR #214994)

[llvmlistbot at llvm.org]

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

Fixes #111646.

### Description
When replacing memref uses in the `NormalizeMemRefs` pass, the `memref.load` and `memref.store` indices are updated using `AffineApplyOp::create`. However, because these newly created operations are not automatically folded by the pass, they leave unfolded `affine.apply` ops in the IR. Depending on later pass interactions (or lack of canonicalization), this results in improperly mapped indices compared to ops like `affine.load` which apply their mappings natively.

This patch replaces the explicit `AffineApplyOp::create` loop in `replaceAllMemRefUsesWith` with `affine::makeComposedFoldedAffineApply`. This ensures that the transformed indices are immediately folded into the correct `Value` expressions (e.g., swapping `%arg0` and `%arg1`) before being passed as operands to the updated memref access operations, bringing `memref.load/store` behavior in line with expected output.

**Note to reviewers:**
* AI was used to help identify the root cause and generate this patch. 


>From dd5804e17ec8fdeb3092ab0aa5da5d110921bbfd Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Sat, 8 Aug 2026 21:58:18 +0530
Subject: [PATCH] [MLIR][Affine] Fix index mapping for memref load/store in
 NormalizeMemRefs

---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7043083298615..2ad262185aeaa 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1277,12 +1277,17 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
     // 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.
+    SmallVector<OpFoldResult, 4> newMapOfrs;
+    for (Value v : newMapOperands)
+      newMapOfrs.push_back(v);
     for (unsigned i = 0; i < newMemRefRank; i++) {
-      state.operands.push_back(AffineApplyOp::create(
+      OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
           builder, op->getLoc(),
           AffineMap::get(newMap.getNumDims(), newMap.getNumSymbols(),
                          newMap.getResult(i)),
-          newMapOperands));
+          newMapOfrs);
+      state.operands.push_back(
+          getValueOrCreateConstantIndexOp(builder, op->getLoc(), ofr));
     }
   }