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

[Hung-Kuan Tseng]

https://github.com/Tim096 updated https://github.com/llvm/llvm-project/pull/221314

>From c63ecc0da80d15d52cabdf2300a0add38431a448 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sat, 5 Sep 2026 02:33:33 +0800
Subject: [PATCH 1/2] [mlir][memref] Keep memref.load attributes when eliding
 reinterpret_cast

`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`:

  %rc = memref.reinterpret_cast %src ... : memref<1x8xf32> to memref<1x1x8xf32>
  %0 = memref.load %rc[%c0, %c0, %i] alignment(16) nontemporal(true) invariant(true) : memref<1x1x8xf32>
  // becomes
  %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. Forward them.
---
 .../MemRef/Transforms/ElideReinterpretCast.cpp    |  4 +++-
 .../MemRef/elide-reinterpret-cast-load.mlir       | 15 +++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

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..e389524c2c38f 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
@@ -252,6 +252,21 @@ func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
   return
 }
 
+// CHECK-LABEL: func.func private @expand_keeps_load_attrs(
+// CHECK-SAME:    %[[SRC:.*]]: memref<1x8xf32>, %[[I:.*]]: index) {
+func.func private @expand_keeps_load_attrs(%src : memref<1x8xf32>, %i : index) {
+  // CHECK:       %[[C0:.*]] = arith.constant 0 : index
+  %c0 = arith.constant 0 : index
+  // CHECK-NOT:   memref.reinterpret_cast
+  %reinterpret_cast = memref.reinterpret_cast %src
+    to offset: [0], sizes: [1, 1, 8], strides: [8, 8, 1]
+    : memref<1x8xf32> to memref<1x1x8xf32>
+  // CHECK:       memref.load %[[SRC]][%[[C0]], %[[I]]] alignment(16) nontemporal(true) invariant(true) : memref<1x8xf32>
+  %0 = memref.load %reinterpret_cast[%c0, %c0, %i]
+    alignment(16) nontemporal(true) invariant(true) : memref<1x1x8xf32>
+  return
+}
+
 //===----------------------------------------------------------------------===//
 // Negative tests (must NOT rewrite)
 //===----------------------------------------------------------------------===//

>From 8efe66193f278cff85065215144d5af82043fd74 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Fri, 11 Sep 2026 22:33:24 +0800
Subject: [PATCH 2/2] Move the attribute test next to @expand_1D

---
 .../MemRef/elide-reinterpret-cast-load.mlir   | 34 +++++++++++--------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
index e389524c2c38f..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(
@@ -252,21 +271,6 @@ func.func private @collapse_3d_moved_unit_dims(%i : index, %j : index,
   return
 }
 
-// CHECK-LABEL: func.func private @expand_keeps_load_attrs(
-// CHECK-SAME:    %[[SRC:.*]]: memref<1x8xf32>, %[[I:.*]]: index) {
-func.func private @expand_keeps_load_attrs(%src : memref<1x8xf32>, %i : index) {
-  // CHECK:       %[[C0:.*]] = arith.constant 0 : index
-  %c0 = arith.constant 0 : index
-  // CHECK-NOT:   memref.reinterpret_cast
-  %reinterpret_cast = memref.reinterpret_cast %src
-    to offset: [0], sizes: [1, 1, 8], strides: [8, 8, 1]
-    : memref<1x8xf32> to memref<1x1x8xf32>
-  // CHECK:       memref.load %[[SRC]][%[[C0]], %[[I]]] alignment(16) nontemporal(true) invariant(true) : memref<1x8xf32>
-  %0 = memref.load %reinterpret_cast[%c0, %c0, %i]
-    alignment(16) nontemporal(true) invariant(true) : memref<1x1x8xf32>
-  return
-}
-
 //===----------------------------------------------------------------------===//
 // Negative tests (must NOT rewrite)
 //===----------------------------------------------------------------------===//