[Mlir-commits] [mlir] 77a26fa - [mlir][memref] Keep memref.load attributes when eliding reinterpret_cast (#221314)

[llvmlistbot at llvm.org]

Author: Hung-Kuan Tseng
Date: 2026-09-14T13:27:03+02:00
New Revision: 77a26fadb88bdc523ef82bcaec8c7e6217701a57

URL: https://github.com/llvm/llvm-project/commit/77a26fadb88bdc523ef82bcaec8c7e6217701a57
DIFF: https://github.com/llvm/llvm-project/commit/77a26fadb88bdc523ef82bcaec8c7e6217701a57.diff

LOG: [mlir][memref] Keep memref.load attributes when eliding reinterpret_cast (#221314)

`RewriteLoadFromReinterpretCast` rebuilds the load on the source of the
`memref.reinterpret_cast` with the indices remapped to the non-unit
dimensions, but the rebuilt load carries none of `nontemporal`,
`alignment` or `invariant`:

```mlir
%rc = memref.reinterpret_cast %src to offset: [0], sizes: [1, 1, 8], strides: [8, 8, 1] : memref<1x8xf32> to memref<1x1x8xf32>
%0 = memref.load %rc[%c0, %c0, %i] alignment(16) nontemporal(true) invariant(true) : memref<1x1x8xf32>
// -memref-elide-reinterpret-cast today
%0 = memref.load %src[%c0, %i] : memref<1x8xf32>
```

The rewrite only changes how the same element is addressed, so the three
attributes describe the same access. This forwards them. Same shape of
fix as #221312 for `gpu-decompose-memrefs`.

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
    mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 49f30a594e0ee..67b53b78a7fef 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -619,7 +619,9 @@ struct RewriteLoadFromReinterpretCast
     // we can safely erase it.
     if (rc.getResult().hasOneUse())
       rewriter.eraseOp(rc);
-    rewriter.replaceOpWithNewOp<memref::LoadOp>(op, rcInput, rcInputIdxs);
+    rewriter.replaceOpWithNewOp<memref::LoadOp>(
+        op, rcInput, rcInputIdxs, op.getNontemporalAttr(),
+        op.getAlignmentAttr(), op.getInvariantAttr());
     return success();
   }
 };

diff  --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
index ce505485c5715..6db97cfcca0c1 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
@@ -54,6 +54,25 @@ func.func private @expand_1D(
   return
 }
 
+/// Same as @expand_1D, with attributes on the load.
+// CHECK-LABEL: func.func private @expand_1D_load_attrs(
+// CHECK-SAME:    %[[SRC:.*]]: memref<1x999xf32>) {
+func.func private @expand_1D_load_attrs(
+    %src : memref<1x999xf32>) {
+  // CHECK-DAG:   %[[IDX_1:.*]] = arith.constant 0 : index
+  // CHECK-DAG:   %[[IDX_2:.*]] = arith.constant 13 : index
+  %idx_1 = arith.constant 0 : index
+  %idx_2 = arith.constant 13 : index
+  // CHECK-NOT:   memref.reinterpret_cast
+  %reinterpret_cast = memref.reinterpret_cast %src
+    to offset: [0], sizes: [1, 1, 999], strides: [999, 999, 1]
+    : memref<1x999xf32> to memref<1x1x999xf32>
+  // CHECK:       %[[LOAD:.*]] = memref.load %[[SRC]][%[[IDX_1]], %[[IDX_2]]] alignment(16) nontemporal(true) invariant(true) : memref<1x999xf32>
+  %0 = memref.load %reinterpret_cast[%idx_1, %idx_1, %idx_2]
+    alignment(16) nontemporal(true) invariant(true) : memref<1x1x999xf32>
+  return
+}
+
 // CHECK-LABEL: func.func private @collapse_1D(
 // CHECK-SAME:    %[[SRC:.*]]: memref<1x1x999xf32>) {
 func.func private @collapse_1D(