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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: RohithPariki

<details>
<summary>Changes</summary>

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. 


---
Full diff: https://github.com/llvm/llvm-project/pull/214994.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Utils/Utils.cpp (+7-2) 


``````````diff
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));
     }
   }
 

``````````

</details>


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