[Mlir-commits] [mlir] [mlir][memref] Define interfaces for ops that access memrefs at an index (PR #177013)
Krzysztof Drewniak
llvmlistbot at llvm.org
Wed Jan 21 10:53:50 PST 2026
https://github.com/krzysz00 updated https://github.com/llvm/llvm-project/pull/177013
>From f0492b2c0d836bd5881bd06985cd389cb4dd3c26 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Tue, 20 Jan 2026 17:39:08 +0000
Subject: [PATCH 1/2] [mlir][memref] Define interfaces for ops that access
memrefs at an index
This commit defines interfaces for operations that perform certain
kinds of indexed access on a memref. These interfaces are defined so
that passes like fold-memref-alias-ops and the memref flattener can be
made generic over operations that, informally, have the forms
`op ... %m[%i0, %i1, ...] ...` (an IndexedAccessOpInterface) or the
form `op %src[%s0, %s1, ...], %dst[%d0, %d1, ...] size ...` (an
IndexedMemCopyOpInterface).
These interfaces have been designed such that all the passes under
MemRef/Transforms that currently have a big switch-case on
memref.load, vector.load, nvgpu.ldmatrix, etc. can be migrated to use
them.
(This'll also let us get rid of the awkward fact that we have memref
transforms depending on the GPU and NVGPU dialects)
While the interface doesn't currently contemplate changing element
types (enabling, for example, writing a bf16 => u16 update to be done
in place), future extensions to allow such transformations could be
worth exploring.
This commit only defines the interfaces so that it'll be easier to
review the design - the implementation is in a future PR.
---
.../mlir/Dialect/MemRef/IR/CMakeLists.txt | 1 +
.../MemRef/IR/MemoryAccessOpInterfaces.h | 26 +++
.../MemRef/IR/MemoryAccessOpInterfaces.td | 188 ++++++++++++++++++
mlir/lib/Dialect/MemRef/IR/CMakeLists.txt | 2 +
.../MemRef/IR/MemoryAccessOpInterfaces.cpp | 22 ++
5 files changed, 239 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h
create mode 100644 mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td
create mode 100644 mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/MemRef/IR/CMakeLists.txt
index b7b12d49f9224..67e27141c18ab 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/MemRef/IR/CMakeLists.txt
@@ -1,2 +1,3 @@
add_mlir_dialect(MemRefOps memref)
+add_mlir_interface(MemoryAccessOpInterfaces)
add_mlir_doc(MemRefOps MemRefOps Dialects/ -gen-op-doc)
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h
new file mode 100644
index 0000000000000..39ee43ec7dcb7
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h
@@ -0,0 +1,26 @@
+//===- MemoryAccessOpInterfaces.h ---------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_MEMREF_IR_MEMORYACCESSOPINTERFACES_H
+#define MLIR_DIALECT_MEMREF_IR_MEMORYACCESSOPINTERFACES_H
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+class RewriterBase;
+} // namespace mlir
+//===----------------------------------------------------------------------===//
+// Memory Access Op Interfaces
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h.inc"
+
+#endif // MLIR_DIALECT_MemREF_IR_MEMORYACCESSOPINTERFACES_H
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td
new file mode 100644
index 0000000000000..4e7508266141f
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td
@@ -0,0 +1,188 @@
+//===-- MemoryAccessOpInterfaces.td ---------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MEMREF_MEMORY_ACCESS_OP_INTERFACES
+#define MEMREF_MEMORY_ACCESS_OP_INTERFACES
+
+include "mlir/IR/OpBase.td"
+
+def IndexedAccessOpInterface : OpInterface<"IndexedAccessOpInterface"> {
+ let description = [{
+ An interface for operations that operate on (by loading from or
+ storing to, atomically modifying, or otherwise) memory located at an
+ index within a memref whose semantics don't depend on the indexing scheme.
+
+ That is, a direct access op is one where, if `%b[%j0, %j1, ..., %jL]`
+ points to the same memory as `%a[%i0,%i1, ... %iK]`, it would be
+ trivial to replace `%a[%i0, ..., %iK]` with %b[%j0, ... %jL]`.
+
+ Operations may impose constaints on allowable reindexings.
+ Returning a non-empty result from `getAccessedShape()` imposes constraints
+ on the dimensions whose strides need to be preserved.
+
+ This interface is intended to enable transformations such as folding in
+ aliasing operations (like memref.subview or memref.collapse_shape) or
+ linearizing memrefs (making them 1-D) to be generic over in-tree and
+ out-of-tree operations.
+ }];
+ let cppNamespace = "::mlir::memref";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the memref that this operation accesses. If the operation
+ is still in tensor form, return the null value.
+ }],
+ /*retType=*/"::mlir::TypedValue<::mlir::MemRefType>",
+ /*methodName=*/"getMemref",
+ /*args=*/(ins)>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the indices that are used to access the memref returned by getMemref().
+
+ The size of this range must be equal to the rank of the memref returned by
+ getMemref().
+ }],
+ /*retType=*/"::mlir::Operation::operand_range",
+ /*methodName=*/"getIndices",
+ /*args=*/(ins)>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the shape of the portion of the memref that is being accessed by
+ this operation, if known, ignoring leading unit dimensions.
+ Reindexing transformations may not modify the *strides* of the tlaining
+ N dimensions, where N is the size returned value, and should ensure that
+ at least N indexing dimensions remain after the transformation.
+ }],
+ /*retType=*/"::llvm::SmallVector<int64_t>",
+ /*methodName=*/"getAccessedShape",
+ /*args=*/(ins),
+ /*methodBody=*/[{}],
+ /*defaultImplementation=*/[{
+ return ::llvm::SmallVector<int64_t>{};
+ }]>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Updates the memref being accessed to `newMemref` and the indices to
+ `newIndices`. If `std::nullopt` is returned, the operation was
+ updated in-place (the common case), while if a vector of values
+ is returned, they sohuld be used to replace the operation being
+ updated.
+
+ The implementor is responsible for rewriter notifications - that is,
+ using modifyOpInPlace().
+
+ The caller must ensure that the new memref/index pair points to the same
+ location in memory as the existing arguments.
+
+ The element types of the memref may not change.
+ }],
+ /*retType=*/"std::optional<llvm::SmallVector<mlir::Value>>",
+ /*methodName=*/"updateMemrefAndIndices",
+ /*args=*/(ins "::mlir::RewriterBase&":$rewriter, "::mlir::Value":$newMemref,
+ "::mlir::ValueRange":$newIndices)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return true if, either by definition or due to some attribute,
+ it's known that all indices are non-negative and less than the size
+ of the dimension they index.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"hasInboundsIndices",
+ /*args=*/(ins),
+ /*methodBody=*/[{}],
+ /*defaultImplementation=*/[{
+ return true;
+ }]>
+ ];
+}
+
+def IndexedMemCopyOpInterface : OpInterface<"IndexedMemCopyOpInterface"> {
+ let description = [{
+ This is an interface for operations that perform a copy of some number
+ of values from `%src[%srcIndices...]` (or some consistently related
+ location) to `%dst[%dstIndices...]` (or some consistently related location
+ - for example, the destination element may be offset by a lane ID in
+ a GPU subgroup).
+
+ The motivating examples for this interface are operations that perform
+ direct loads to workgroup memory on GPUs.
+
+ This allows patterns that reindex memrefs (like folding in subview operations)
+ to treach such operations as a class, just like `IndexedAccessOpInterface`.
+
+ Unlike `IndexedAccessOpInterface`, this interface assumes that the elements
+ being copied are contiguous in memory and that the produces of the
+ operation has ensured this. That is, if the source memref is a
+ `memref<8x3xf32, strided<[9, 1]>>`, it is presumed that a memcopy of
+ 4 floats starting at the indices `[%x, %y]` is intended to read into
+ the space between the length-three rows.
+ }];
+ let cppNamespace = "::mlir::memref";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the source memref for this copy operation.
+ }],
+ /*retType=*/"::mlir::TypedValue<::mlir::MemRefType>",
+ /*methodName=*/"getSrc",
+ /*args=*/(ins)>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the indices that are used to access the source memref.
+
+ The size of this range must be equal to the rank of the memref returned by
+ getSrc().
+ }],
+ /*retType=*/"::mlir::Operation::operand_range",
+ /*methodName=*/"getSrcIndices",
+ /*args=*/(ins)>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the destination memref for this copy operation.
+ }],
+ /*retType=*/"::mlir::TypedValue<::mlir::MemRefType>",
+ /*methodName=*/"getDst",
+ /*args=*/(ins)>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the indices that are used to access the destination memref.
+
+ The size of this range must be equal to the rank of the memref returned by
+ getDst().
+ }],
+ /*retType=*/"::mlir::Operation::operand_range",
+ /*methodName=*/"getDstIndices",
+ /*args=*/(ins)>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Update the operation with `newSrcMemref` as the new source memref,
+ `newSrcIndices` as the new source indices, `newDst` as the new destination
+ memref, and `newDstIndices` as the new destination indices, leaving
+ other properties unchanged.
+
+ Note that, unlike the method on indexed accessors, this method doesn't
+ currently allow for the possibility of cloning since DMA ops generally
+ do not return results. If needed, such suppport could be added in
+ the future.
+
+ The implementation is responsible for issuing rewriter notifications
+ vio modifyOpInPlace().
+
+ The caller must ensure that the new memref/index pairs point to the same
+ location in memory.
+ }],
+ /*retType=*/"void",
+ /*methodName=*/"setMemrefsAndIndices",
+ /*args=*/(ins "::mlir::RewriterBase&":$rewriter,
+ "::mlir::Value":$newSrc, "::mlir::ValueRange":$newSrcIndices,
+ "::mlir::Value":$newDst, "::mlir::ValueRange":$newDstIndices)
+ >
+ ];
+}
+#endif // MEMREF_MEMORY_ACCESS_OP_INTERFACES
diff --git a/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt b/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt
index d358362f1984b..56d98dd431c8e 100644
--- a/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/MemRef/IR/CMakeLists.txt
@@ -1,5 +1,6 @@
add_mlir_dialect_library(MLIRMemRefDialect
MemRefDialect.cpp
+ MemoryAccessOpInterfaces.cpp
MemRefMemorySlot.cpp
MemRefOps.cpp
ValueBoundsOpInterfaceImpl.cpp
@@ -9,6 +10,7 @@ add_mlir_dialect_library(MLIRMemRefDialect
DEPENDS
MLIRMemRefOpsIncGen
+ MLIRMemoryAccessOpInterfacesIncGen
LINK_LIBS PUBLIC
MLIRArithDialect
diff --git a/mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp b/mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp
new file mode 100644
index 0000000000000..e82b036c6f89b
--- /dev/null
+++ b/mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp
@@ -0,0 +1,22 @@
+//===- MemoryAccessOpInterfaces.cpp ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+
+//===----------------------------------------------------------------------===//
+// IndexedAccessOpInterface and IndexedMemCpyOpInterface
+//===----------------------------------------------------------------------===//
+
+namespace mlir::memref {
+
+#include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp.inc"
+
+} // namespace mlir::memref
>From 7aa8d96c7022c012886a0efb746a28dd919b1fe1 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Tue, 20 Jan 2026 23:38:56 +0000
Subject: [PATCH 2/2] Update header comments to align, interface def part
---
mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h | 2 +-
mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td | 2 +-
mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp | 2 +-
3 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h
index 39ee43ec7dcb7..d241c19e573ac 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h
@@ -1,4 +1,4 @@
-//===- MemoryAccessOpInterfaces.h ---------------------*- C++ -*-===//
+//===- MemoryAccessOpInterfaces.h -------------------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td
index 4e7508266141f..8df22cad52205 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.td
@@ -1,4 +1,4 @@
-//===-- MemoryAccessOpInterfaces.td ---------------*- tablegen -*-===//
+//===-- MemoryAccessOpInterfaces.td ------------------------*- tablegen -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
diff --git a/mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp b/mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp
index e82b036c6f89b..254e5e7d448f2 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp
@@ -1,4 +1,4 @@
-//===- MemoryAccessOpInterfaces.cpp ----------------------------===//
+//===- MemoryAccessOpInterfaces.cpp ---------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
More information about the Mlir-commits
mailing list