[Mlir-commits] [mlir] [mlir][memref] Use access interfaces in address extraction (PR #198421)
Krzysztof Drewniak
llvmlistbot at llvm.org
Wed Jun 17 13:48:05 PDT 2026
https://github.com/krzysz00 updated https://github.com/llvm/llvm-project/pull/198421
>From 72c5bbccff944cbb60437cb0fbc6b88830133811 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 18 May 2026 16:13:57 +0000
Subject: [PATCH 1/3] [mlir][memref] Use access interfaces in address
extraction
Rework extract-address-computation patterns to use
IndexedAccessOpInterface for direct memref accesses and
VectorTransferOpInterface update hooks for transfer ops.
These rewrites are limited to operations that declare in-bounds
indices (so no vector.load currently) for now so that we always
create valid `memref.subview` opp.
As a consequence of this PR, the Memref dialect no longer depends on
the GPU and NVGPU dialects.
AI: Codex wrote the first drat, I silpfified it a bunch and made sure
the names of internal functions made sense.
Co-Authored-By: Codex <codex at openai.com>
---
.../Dialect/MemRef/Transforms/Transforms.h | 10 +-
.../Dialect/MemRef/Transforms/CMakeLists.txt | 3 -
.../Transforms/ExtractAddressComputations.cpp | 434 ++++++++----------
.../MemRef/extract-address-computations.mlir | 149 +++++-
mlir/test/lib/Dialect/Vector/CMakeLists.txt | 1 +
5 files changed, 328 insertions(+), 269 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h b/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
index 4d6c54d74d2a9..7d3e67f5ec29e 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
@@ -141,9 +141,13 @@ FailureOr<memref::AllocOp> multiBuffer(memref::AllocOp allocOp,
unsigned multiplier,
bool skipOverrideAnalysis = false);
-/// Appends patterns for extracting address computations from the instructions
-/// with memory accesses such that these memory accesses use only a base
-/// pointer.
+/// Appends patterns for extracting address computations from memory access
+/// operations such that these accesses use only a base pointer.
+///
+/// The patterns match memref::IndexedAccessOpInterface and
+/// VectorTransferOpInterface generically. Callers that rely on external models
+/// should register the appropriate dialect extensions, such as the NVGPU, GPU
+/// and Vector indexed-access models registered by RegisterAllDialects.
///
/// For instance,
/// ```mlir
diff --git a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
index 1c5e07f89b338..1d409b6f06242 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
@@ -31,16 +31,13 @@ add_mlir_dialect_library(MLIRMemRefTransforms
MLIRArithTransforms
MLIRDialectUtils
MLIRFuncDialect
- MLIRGPUDialect
MLIRInferTypeOpInterface
MLIRLoopLikeInterface
MLIRMemRefDialect
MLIRMemRefUtils
- MLIRNVGPUDialect
MLIRPass
MLIRTensorDialect
MLIRTransforms
MLIRValueBoundsOpInterface
MLIRVectorDialect
)
-
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp b/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
index 9c922c28d0f54..306b708b22cbf 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
@@ -6,275 +6,243 @@
//
//===----------------------------------------------------------------------===//
//
-/// This transformation pass rewrites loading/storing from/to a memref with
-/// offsets into loading/storing from/to a subview and without any offset on
-/// the instruction itself.
+/// This transformation pass rewrites memory access operations with offsets into
+/// accesses through a subview and without any offset on the access operation
+/// itself.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
-#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/PatternMatch.h"
-#include "llvm/ADT/Repeated.h"
using namespace mlir;
namespace {
//===----------------------------------------------------------------------===//
-// Helper functions for the `load base[off0...]`
-// => `load (subview base[off0...])[0...]` pattern.
+// Helper functions for the `access base[off0...]`
+// => `access (subview base[off0...])[0...]` pattern.
//===----------------------------------------------------------------------===//
-// Matches getFailureOrSrcMemRef specs for LoadOp.
-// \see LoadStoreLikeOpRewriter.
-static FailureOr<Value> getLoadOpSrcMemRef(memref::LoadOp loadOp) {
- return loadOp.getMemRef();
+/// Returns true if every index is zero.
+static bool hasAllZeroIndices(ValueRange indices) {
+ return llvm::all_of(getAsOpFoldResult(indices), isZeroInteger);
}
-// Matches rebuildOpFromAddressAndIndices specs for LoadOp.
-// \see LoadStoreLikeOpRewriter.
-static memref::LoadOp rebuildLoadOp(RewriterBase &rewriter,
- memref::LoadOp loadOp, Value srcMemRef,
- ValueRange indices) {
- Location loc = loadOp.getLoc();
- return memref::LoadOp::create(rewriter, loc, srcMemRef, indices,
- loadOp.getNontemporal());
-}
-
-// Matches getViewSizeForEachDim specs for LoadOp.
-// \see LoadStoreLikeOpRewriter.
-static SmallVector<OpFoldResult>
-getLoadOpViewSizeForEachDim(RewriterBase &rewriter, memref::LoadOp loadOp) {
- MemRefType ldTy = loadOp.getMemRefType();
- unsigned loadRank = ldTy.getRank();
- return SmallVector<OpFoldResult>(loadRank, rewriter.getIndexAttr(1));
-}
-
-//===----------------------------------------------------------------------===//
-// Helper functions for the `store val, base[off0...]`
-// => `store val, (subview base[off0...])[0...]` pattern.
-//===----------------------------------------------------------------------===//
+/// Get the remaining size in each dimension - that is, the size of the memref
+/// dimension minus the index. Used to preserve in_bounds behavior for
+/// transfer_read/write.
+static SmallVector<OpFoldResult> getRemainingSizes(RewriterBase &rewriter,
+ Location loc,
+ Value srcMemRef,
+ ValueRange indices) {
+ auto extractStridedMetadataOp =
+ memref::ExtractStridedMetadataOp::create(rewriter, loc, srcMemRef);
+ SmallVector<OpFoldResult> srcSizes =
+ extractStridedMetadataOp.getConstifiedMixedSizes();
+ SmallVector<OpFoldResult> mixedIndices = getAsOpFoldResult(indices);
+ SmallVector<OpFoldResult> finalSizes;
-// Matches getFailureOrSrcMemRef specs for StoreOp.
-// \see LoadStoreLikeOpRewriter.
-static FailureOr<Value> getStoreOpSrcMemRef(memref::StoreOp storeOp) {
- return storeOp.getMemRef();
-}
+ AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
+ AffineExpr s1 = rewriter.getAffineSymbolExpr(1);
-// Matches rebuildOpFromAddressAndIndices specs for StoreOp.
-// \see LoadStoreLikeOpRewriter.
-static memref::StoreOp rebuildStoreOp(RewriterBase &rewriter,
- memref::StoreOp storeOp, Value srcMemRef,
- ValueRange indices) {
- Location loc = storeOp.getLoc();
- return memref::StoreOp::create(rewriter, loc, storeOp.getValueToStore(),
- srcMemRef, indices, storeOp.getNontemporal());
+ for (auto [srcSize, index] : llvm::zip_equal(srcSizes, mixedIndices)) {
+ finalSizes.push_back(affine::makeComposedFoldedAffineApply(
+ rewriter, loc, s0 - s1, {srcSize, index}));
+ }
+ return finalSizes;
}
-// Matches getViewSizeForEachDim specs for StoreOp.
-// \see LoadStoreLikeOpRewriter.
+/// Get the sizes needed to create a valid subview for an indexed access.
+/// The trailing dimensions are sized using the accessed shape, taking
+/// the minimum of that shape's size and what's available along the relevant
+/// memref dimension as a dynamic value if the memref is dynamically shaped
+/// (so as to avoid subviews that exceed the bounds of the relevant memref
+/// dimension). If the operation accesses a dynamic number of elements along
+/// the dimension, the size of the subview will always be the remaining element
+/// count along the dimension.
static SmallVector<OpFoldResult>
-getStoreOpViewSizeForEachDim(RewriterBase &rewriter, memref::StoreOp storeOp) {
- MemRefType ldTy = storeOp.getMemRefType();
- unsigned loadRank = ldTy.getRank();
- return SmallVector<OpFoldResult>(loadRank, rewriter.getIndexAttr(1));
-}
+getIndexedAccessViewSizes(RewriterBase &rewriter,
+ memref::IndexedAccessOpInterface op) {
+ TypedValue<MemRefType> srcMemRef = op.getAccessedMemref();
+ assert(srcMemRef && "expected indexed access with a memref");
+
+ MemRefType srcType = srcMemRef.getType();
+ int64_t srcRank = srcType.getRank();
+ SmallVector<int64_t> accessedShape = op.getAccessedShape();
+ int64_t accessedRank = static_cast<int64_t>(accessedShape.size());
+ assert(accessedRank <= srcRank &&
+ "can't access more dimensions than a memref has");
+
+ SmallVector<OpFoldResult> indices = getAsOpFoldResult(op.getIndices());
+ int64_t firstAccessedDim = srcRank - accessedRank;
+
+ Location loc = op.getLoc();
+ SmallVector<OpFoldResult> viewSizes(srcRank, rewriter.getIndexAttr(1));
+ SmallVector<OpFoldResult> srcSizes;
+ AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
+ AffineExpr s1 = rewriter.getAffineSymbolExpr(1);
+ AffineExpr cst = rewriter.getAffineSymbolExpr(2);
-//===----------------------------------------------------------------------===//
-// Helper functions for the `ldmatrix base[off0...]`
-// => `ldmatrix (subview base[off0...])[0...]` pattern.
-//===----------------------------------------------------------------------===//
+ auto ensureSrcSizes = [&]() {
+ if (srcSizes.empty()) {
+ auto extractStridedMetadataOp =
+ memref::ExtractStridedMetadataOp::create(rewriter, loc, srcMemRef);
+ srcSizes = extractStridedMetadataOp.getConstifiedMixedSizes();
+ }
+ };
+
+ for (int64_t accessedDim : llvm::seq<int64_t>(0, accessedRank)) {
+ int64_t accessedSize = accessedShape[accessedDim];
+ int64_t dim = firstAccessedDim + accessedDim;
+ if (!ShapedType::isDynamic(accessedSize)) {
+ int64_t srcDimSize = srcType.getDimSize(dim);
+ if (!ShapedType::isDynamic(srcDimSize) || accessedSize == 1) {
+ viewSizes[dim] = rewriter.getIndexAttr(accessedSize);
+ continue;
+ }
+ ensureSrcSizes();
+ viewSizes[dim] = affine::makeComposedFoldedAffineMin(
+ rewriter, loc,
+ AffineMap::get(/*dimCount=*/0, /*symbolCount=*/3, {s0 - s1, cst},
+ rewriter.getContext()),
+ {srcSizes[dim], indices[dim], rewriter.getIndexAttr(accessedSize)});
+ } else {
+ ensureSrcSizes();
+ viewSizes[dim] = affine::makeComposedFoldedAffineApply(
+ rewriter, loc, s0 - s1, {srcSizes[dim], indices[dim]});
+ }
+ }
+ return viewSizes;
+}
-// Matches getFailureOrSrcMemRef specs for LdMatrixOp.
-// \see LoadStoreLikeOpRewriter.
-static FailureOr<Value> getLdMatrixOpSrcMemRef(nvgpu::LdMatrixOp ldMatrixOp) {
- return ldMatrixOp.getSrcMemref();
+static memref::SubViewOp createSubviewForAccess(RewriterBase &rewriter,
+ Location loc, Value srcMemRef,
+ ValueRange indices,
+ ArrayRef<OpFoldResult> sizes) {
+ int64_t rank = cast<MemRefType>(srcMemRef.getType()).getRank();
+ SmallVector<OpFoldResult> mixedIndices = getAsOpFoldResult(indices);
+ SmallVector<OpFoldResult> ones(rank, rewriter.getIndexAttr(1));
+
+ return memref::SubViewOp::create(rewriter, loc, /*source=*/srcMemRef,
+ /*offsets=*/mixedIndices,
+ /*sizes=*/sizes, /*strides=*/ones);
}
-// Matches rebuildOpFromAddressAndIndices specs for LdMatrixOp.
-// \see LoadStoreLikeOpRewriter.
-static nvgpu::LdMatrixOp rebuildLdMatrixOp(RewriterBase &rewriter,
- nvgpu::LdMatrixOp ldMatrixOp,
- Value srcMemRef,
- ValueRange indices) {
- Location loc = ldMatrixOp.getLoc();
- return nvgpu::LdMatrixOp::create(
- rewriter, loc, ldMatrixOp.getResult().getType(), srcMemRef, indices,
- ldMatrixOp.getTranspose(), ldMatrixOp.getNumTiles());
+static SmallVector<Value> getZeroIndices(RewriterBase &rewriter, Location loc,
+ int64_t rank) {
+ if (rank == 0)
+ return {};
+ Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ return SmallVector<Value>(rank, zero);
}
-//===----------------------------------------------------------------------===//
-// Helper functions for the `transfer_read base[off0...]`
-// => `transfer_read (subview base[off0...])[0...]` pattern.
-//===----------------------------------------------------------------------===//
+/// Rewrite an indexed access op so that all its indices are zeros.
+/// E.g., %res = indexed_access %base[%off0]...[%offN]
+/// =>
+/// %new_base = subview %base[%off0,.., %offN][1,..,1][1,..,1]
+/// %res = indexed_access %new_base[0,..,0] :
+/// memref<1x..x1xTy, strided<[1,..,1], offset: ?>>
+struct IndexedAccessOpRewriter final
+ : OpInterfaceRewritePattern<memref::IndexedAccessOpInterface> {
+ using Base::Base;
-// Matches getFailureOrSrcMemRef specs for TransferReadOp.
-// \see LoadStoreLikeOpRewriter.
-template <typename TransferLikeOp>
-static FailureOr<Value>
-getTransferLikeOpSrcMemRef(TransferLikeOp transferLikeOp) {
- Value src = transferLikeOp.getBase();
- if (isa<MemRefType>(src.getType()))
- return src;
- return failure();
-}
+ LogicalResult matchAndRewrite(memref::IndexedAccessOpInterface op,
+ PatternRewriter &rewriter) const override {
+ TypedValue<MemRefType> srcMemRef = op.getAccessedMemref();
+ if (!srcMemRef)
+ return rewriter.notifyMatchFailure(op, "source is not a memref");
-// Matches rebuildOpFromAddressAndIndices specs for TransferReadOp.
-// \see LoadStoreLikeOpRewriter.
-static vector::TransferReadOp
-rebuildTransferReadOp(RewriterBase &rewriter,
- vector::TransferReadOp transferReadOp, Value srcMemRef,
- ValueRange indices) {
- Location loc = transferReadOp.getLoc();
- return vector::TransferReadOp::create(
- rewriter, loc, transferReadOp.getResult().getType(), srcMemRef, indices,
- transferReadOp.getPermutationMap(), transferReadOp.getPadding(),
- transferReadOp.getMask(), transferReadOp.getInBoundsAttr());
-}
+ int64_t rank = srcMemRef.getType().getRank();
+ if (rank == 0)
+ return rewriter.notifyMatchFailure(op,
+ "0-D accesses don't need rewriting");
-//===----------------------------------------------------------------------===//
-// Helper functions for the `transfer_write base[off0...]`
-// => `transfer_write (subview base[off0...])[0...]` pattern.
-//===----------------------------------------------------------------------===//
+ if (static_cast<int64_t>(op.getAccessedShape().size()) > rank)
+ return rewriter.notifyMatchFailure(
+ op, "can't access more dimensions than a memref has");
-// Matches rebuildOpFromAddressAndIndices specs for TransferWriteOp.
-// \see LoadStoreLikeOpRewriter.
-static vector::TransferWriteOp
-rebuildTransferWriteOp(RewriterBase &rewriter,
- vector::TransferWriteOp transferWriteOp, Value srcMemRef,
- ValueRange indices) {
- Location loc = transferWriteOp.getLoc();
- return vector::TransferWriteOp::create(
- rewriter, loc, transferWriteOp.getValue(), srcMemRef, indices,
- transferWriteOp.getPermutationMapAttr(), transferWriteOp.getMask(),
- transferWriteOp.getInBoundsAttr());
-}
+ if (!op.hasInboundsIndices())
+ return rewriter.notifyMatchFailure(op, "indices may be out of bounds");
-//===----------------------------------------------------------------------===//
-// Generic helper functions used as default implementation in
-// LoadStoreLikeOpRewriter.
-//===----------------------------------------------------------------------===//
-/// Helper function to get the src memref.
-/// It uses the already defined getFailureOrSrcMemRef but asserts
-/// that the source is a memref.
-template <typename LoadStoreLikeOp,
- FailureOr<Value> (*getFailureOrSrcMemRef)(LoadStoreLikeOp)>
-static Value getSrcMemRef(LoadStoreLikeOp loadStoreLikeOp) {
- FailureOr<Value> failureOrSrcMemRef = getFailureOrSrcMemRef(loadStoreLikeOp);
- assert(!failed(failureOrSrcMemRef) && "Generic getSrcMemRef cannot be used");
- return *failureOrSrcMemRef;
-}
+ // If the access already has only zeros as indices there is nothing
+ // to do.
+ if (hasAllZeroIndices(op.getIndices()))
+ return rewriter.notifyMatchFailure(
+ op, "no computation to extract: offsets are 0s");
-/// Helper function to get the sizes of the resulting view.
-/// This function gets the sizes of the source memref then substracts the
-/// offsets used within \p loadStoreLikeOp. This gives the maximal (for
-/// inbound) sizes for the view.
-/// The source memref is retrieved using getSrcMemRef on \p loadStoreLikeOp.
-template <typename LoadStoreLikeOp, Value (*getSrcMemRef)(LoadStoreLikeOp)>
-static SmallVector<OpFoldResult>
-getGenericOpViewSizeForEachDim(RewriterBase &rewriter,
- LoadStoreLikeOp loadStoreLikeOp) {
- Location loc = loadStoreLikeOp.getLoc();
- auto extractStridedMetadataOp = memref::ExtractStridedMetadataOp::create(
- rewriter, loc, getSrcMemRef(loadStoreLikeOp));
- SmallVector<OpFoldResult> srcSizes =
- extractStridedMetadataOp.getConstifiedMixedSizes();
- SmallVector<OpFoldResult> indices =
- getAsOpFoldResult(loadStoreLikeOp.getIndices());
- SmallVector<OpFoldResult> finalSizes;
+ SmallVector<OpFoldResult> subviewSizes =
+ getIndexedAccessViewSizes(rewriter, op);
- AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
- AffineExpr s1 = rewriter.getAffineSymbolExpr(1);
+ Location loc = op.getLoc();
+ auto subview = createSubviewForAccess(rewriter, loc, srcMemRef,
+ op.getIndices(), subviewSizes);
+ SmallVector<Value> zeros = getZeroIndices(rewriter, loc, rank);
- for (auto [srcSize, indice] : llvm::zip(srcSizes, indices)) {
- finalSizes.push_back(affine::makeComposedFoldedAffineApply(
- rewriter, loc, s0 - s1, {srcSize, indice}));
+ std::optional<SmallVector<Value>> newValues =
+ op.updateMemrefAndIndices(rewriter, subview.getResult(), zeros);
+ if (newValues)
+ rewriter.replaceOp(op, *newValues);
+ return success();
}
- return finalSizes;
-}
+};
-/// Rewrite a store/load-like op so that all its indices are zeros.
-/// E.g., %ld = memref.load %base[%off0]...[%offN]
-/// =>
-/// %new_base = subview %base[%off0,.., %offN][1,..,1][1,..,1]
-/// %ld = memref.load %new_base[0,..,0] :
-/// memref<1x..x1xTy, strided<[1,..,1], offset: ?>>
-///
-/// `getSrcMemRef` returns the source memref for the given load-like operation.
-///
-/// `getViewSizeForEachDim` returns the sizes of view that is going to feed
-/// new operation. This must return one size per dimension of the view.
-/// The sizes of the view needs to be at least as big as what is actually
-/// going to be accessed. Use the provided `loadStoreOp` to get the right
-/// sizes.
-///
-/// Using the given rewriter, `rebuildOpFromAddressAndIndices` creates a new
-/// LoadStoreLikeOp that reads from srcMemRef[indices].
-/// The returned operation will be used to replace loadStoreOp.
-template <typename LoadStoreLikeOp,
- FailureOr<Value> (*getFailureOrSrcMemRef)(LoadStoreLikeOp),
- LoadStoreLikeOp (*rebuildOpFromAddressAndIndices)(
- RewriterBase & /*rewriter*/, LoadStoreLikeOp /*loadStoreOp*/,
- Value /*srcMemRef*/, ValueRange /*indices*/),
- SmallVector<OpFoldResult> (*getViewSizeForEachDim)(
- RewriterBase & /*rewriter*/, LoadStoreLikeOp /*loadStoreOp*/) =
- getGenericOpViewSizeForEachDim<
- LoadStoreLikeOp,
- getSrcMemRef<LoadStoreLikeOp, getFailureOrSrcMemRef>>>
-struct LoadStoreLikeOpRewriter : public OpRewritePattern<LoadStoreLikeOp> {
- using OpRewritePattern<LoadStoreLikeOp>::OpRewritePattern;
-
- LogicalResult matchAndRewrite(LoadStoreLikeOp loadStoreLikeOp,
+/// Rewrite a vector transfer op so that all its indices are zeros.
+struct TransferOpRewriter final
+ : OpInterfaceRewritePattern<VectorTransferOpInterface> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(VectorTransferOpInterface op,
PatternRewriter &rewriter) const override {
- FailureOr<Value> failureOrSrcMemRef =
- getFailureOrSrcMemRef(loadStoreLikeOp);
- if (failed(failureOrSrcMemRef))
- return rewriter.notifyMatchFailure(loadStoreLikeOp,
- "source is not a memref");
- Value srcMemRef = *failureOrSrcMemRef;
- auto ldStTy = cast<MemRefType>(srcMemRef.getType());
- unsigned loadStoreRank = ldStTy.getRank();
- // Don't waste compile time if there is nothing to rewrite.
- if (loadStoreRank == 0)
- return rewriter.notifyMatchFailure(loadStoreLikeOp,
+ Value srcMemRef = op.getBase();
+ auto srcType = dyn_cast<MemRefType>(srcMemRef.getType());
+ if (!srcType)
+ return rewriter.notifyMatchFailure(op, "source is not a memref");
+
+ int64_t rank = srcType.getRank();
+
+ if (rank == 0)
+ return rewriter.notifyMatchFailure(op,
"0-D accesses don't need rewriting");
- // If our load already has only zeros as indices there is nothing
- // to do.
- SmallVector<OpFoldResult> indices =
- getAsOpFoldResult(loadStoreLikeOp.getIndices());
- if (llvm::all_of(indices, isZeroInteger)) {
+ if (hasAllZeroIndices(op.getIndices()))
return rewriter.notifyMatchFailure(
- loadStoreLikeOp, "no computation to extract: offsets are 0s");
- }
+ op, "no computation to extract: offsets are 0s");
+
+ Location loc = op.getLoc();
+ SmallVector<OpFoldResult> offsets = getAsOpFoldResult(op.getIndices());
+ SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
+ // Approximate sizes needed so we can test the general case of the
+ // replacement we're planning to do - this can be tightened up later when
+ // this pattern is extended to reason about in_bounds, which dimensions are
+ // accessed, etc.
+ SmallVector<OpFoldResult> approximateSizes(
+ rank, rewriter.getIndexAttr(ShapedType::kDynamic));
+ MemRefType subviewType = memref::SubViewOp::inferResultType(
+ srcType, offsets, approximateSizes, strides);
+ if (!subviewType)
+ return rewriter.notifyMatchFailure(op, "failed to infer subview type");
+
+ AffineMap permutationMap = op.getPermutationMap();
+ if (failed(op.mayUpdateStartingPosition(subviewType, permutationMap)))
+ return rewriter.notifyMatchFailure(op,
+ "failed op-specific preconditions");
- // Create the array of ones of the right size.
- SmallVector<OpFoldResult> ones(loadStoreRank, rewriter.getIndexAttr(1));
SmallVector<OpFoldResult> sizes =
- getViewSizeForEachDim(rewriter, loadStoreLikeOp);
- assert(sizes.size() == loadStoreRank &&
- "Expected one size per load dimension");
- Location loc = loadStoreLikeOp.getLoc();
- // The subview inherits its strides from the original memref and will
- // apply them properly to the input indices.
- // Therefore the strides multipliers are simply ones.
- auto subview =
- memref::SubViewOp::create(rewriter, loc, /*source=*/srcMemRef,
- /*offsets=*/indices,
- /*sizes=*/sizes, /*strides=*/ones);
- // Rewrite the load/store with the subview as the base pointer.
- Repeated<Value> zeros(loadStoreRank,
- arith::ConstantIndexOp::create(rewriter, loc, 0));
- LoadStoreLikeOp newLoadStore = rebuildOpFromAddressAndIndices(
- rewriter, loadStoreLikeOp, subview.getResult(), zeros);
- rewriter.replaceOp(loadStoreLikeOp, newLoadStore->getResults());
+ getRemainingSizes(rewriter, loc, srcMemRef, op.getIndices());
+ auto subview = createSubviewForAccess(rewriter, loc, srcMemRef,
+ op.getIndices(), sizes);
+ SmallVector<Value> zeros = getZeroIndices(rewriter, loc, rank);
+
+ op.updateStartingPosition(rewriter, subview.getResult(), zeros,
+ AffineMapAttr::get(permutationMap));
return success();
}
};
@@ -282,28 +250,6 @@ struct LoadStoreLikeOpRewriter : public OpRewritePattern<LoadStoreLikeOp> {
void memref::populateExtractAddressComputationsPatterns(
RewritePatternSet &patterns) {
- patterns.add<
- LoadStoreLikeOpRewriter<
- memref::LoadOp,
- /*getSrcMemRef=*/getLoadOpSrcMemRef,
- /*rebuildOpFromAddressAndIndices=*/rebuildLoadOp,
- /*getViewSizeForEachDim=*/getLoadOpViewSizeForEachDim>,
- LoadStoreLikeOpRewriter<
- memref::StoreOp,
- /*getSrcMemRef=*/getStoreOpSrcMemRef,
- /*rebuildOpFromAddressAndIndices=*/rebuildStoreOp,
- /*getViewSizeForEachDim=*/getStoreOpViewSizeForEachDim>,
- LoadStoreLikeOpRewriter<
- nvgpu::LdMatrixOp,
- /*getSrcMemRef=*/getLdMatrixOpSrcMemRef,
- /*rebuildOpFromAddressAndIndices=*/rebuildLdMatrixOp>,
- LoadStoreLikeOpRewriter<
- vector::TransferReadOp,
- /*getSrcMemRef=*/getTransferLikeOpSrcMemRef<vector::TransferReadOp>,
- /*rebuildOpFromAddressAndIndices=*/rebuildTransferReadOp>,
- LoadStoreLikeOpRewriter<
- vector::TransferWriteOp,
- /*getSrcMemRef=*/getTransferLikeOpSrcMemRef<vector::TransferWriteOp>,
- /*rebuildOpFromAddressAndIndices=*/rebuildTransferWriteOp>>(
+ patterns.add<IndexedAccessOpRewriter, TransferOpRewriter>(
patterns.getContext());
}
diff --git a/mlir/test/Dialect/MemRef/extract-address-computations.mlir b/mlir/test/Dialect/MemRef/extract-address-computations.mlir
index eec3d5c62983b..44737fd23fce2 100644
--- a/mlir/test/Dialect/MemRef/extract-address-computations.mlir
+++ b/mlir/test/Dialect/MemRef/extract-address-computations.mlir
@@ -184,25 +184,137 @@ module attributes {transform.with_named_sequence} {
// -----
+// Check that generic IndexedAccessOpInterface users are supported.
+
+// CHECK-LABEL: @test_atomic_rmw(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref{{[^,]*}},
+// CHECK-SAME: %[[VALUE:[^:]*]]: f32,
+// CHECK-SAME: %[[DYN_OFFSET:.*]]: index)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET]], 1] [1, 1] [1, 1] : memref<8x8xf32> to memref<1x1xf32, strided<[8, 1], offset: ?>>
+// CHECK: %[[RES:.*]] = memref.atomic_rmw addf %[[VALUE]], %[[SUBVIEW]][%[[C0]], %[[C0]]] : (f32, memref<1x1xf32, strided<[8, 1], offset: ?>>) -> f32
+// CHECK: return %[[RES]] : f32
+func.func @test_atomic_rmw(%base : memref<8x8xf32>, %value : f32,
+ %offset : index) -> f32 {
+ %c1 = arith.constant 1 : index
+ %res = memref.atomic_rmw addf %value, %base[%offset, %c1] : (f32, memref<8x8xf32>) -> f32
+ return %res : f32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.memref.extract_address_computations
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Check that zero-result IndexedAccessOpInterface users are supported.
+
+// CHECK-LABEL: @test_prefetch(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref{{[^,]*}},
+// CHECK-SAME: %[[DYN_OFFSET:.*]]: index)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET]], 1] [1, 1] [1, 1] : memref<8x8xf32> to memref<1x1xf32, strided<[8, 1], offset: ?>>
+// CHECK: memref.prefetch %[[SUBVIEW]][%[[C0]], %[[C0]]], read, locality<3>, data : memref<1x1xf32, strided<[8, 1], offset: ?>>
+func.func @test_prefetch(%base : memref<8x8xf32>, %offset : index) {
+ %c1 = arith.constant 1 : index
+ memref.prefetch %base[%offset, %c1], read, locality<3>, data
+ : memref<8x8xf32>
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.memref.extract_address_computations
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Do not rewrite indexed access ops that do not guarantee in-bounds indices.
+
+// CHECK-LABEL: @negative_vector_load_without_inbounds(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref<{{[^,]*}}>,
+// CHECK-SAME: %[[DYN_OFFSET:.*]]: index)
+// CHECK-NOT: memref.subview
+// CHECK: %[[LOADED_VAL:.*]] = vector.load %[[BASE]][%[[DYN_OFFSET]], %{{.*}}] : memref<?x?xf32>, vector<4xf32>
+// CHECK: return %[[LOADED_VAL]] : vector<4xf32>
+func.func @negative_vector_load_without_inbounds(%base : memref<?x?xf32>,
+ %offset : index) -> vector<4xf32> {
+ %c0 = arith.constant 0 : index
+ %loaded_val = vector.load %base[%offset, %c0] : memref<?x?xf32>, vector<4xf32>
+ return %loaded_val : vector<4xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.memref.extract_address_computations
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Rewrite static accessed shapes against dynamic source dimensions using the
+// minimum of the remaining dimension size and the accessed shape size.
+
+// CHECK-DAG: #[[$MIN_496_MAP:.*]] = affine_map<()[s0] -> (s0, 496)>
+// CHECK-LABEL: @test_gpu_subgroup_mma_dynamic_source_shape(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref<?x?xf16>,
+// CHECK-SAME: %[[DYN_OFFSET:.*]]: index)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: {{.*}}, {{.*}}, %[[SIZES:.*]]:2, {{.*}} = memref.extract_strided_metadata %[[BASE]]
+// CHECK-DAG: %[[DYN_SIZE:.*]] = affine.min #[[$MIN_496_MAP]]()[%[[SIZES]]#1]
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET]], 0] [1, %[[DYN_SIZE]]] [1, 1] : memref<?x?xf16> to memref<1x?xf16, strided<[?, 1], offset: ?>>
+// CHECK: gpu.subgroup_mma_load_matrix %[[SUBVIEW]][%[[C0]], %[[C0]]] {leadDimension = 32 : index} : memref<1x?xf16, strided<[?, 1], offset: ?>> -> !gpu.mma_matrix<16x16xf16, "AOp">
+func.func @test_gpu_subgroup_mma_dynamic_source_shape(
+ %base : memref<?x?xf16>, %offset : index)
+ -> !gpu.mma_matrix<16x16xf16, "AOp"> {
+ %c0 = arith.constant 0 : index
+ %matrix = gpu.subgroup_mma_load_matrix %base[%offset, %c0]
+ {leadDimension = 32 : index}
+ : memref<?x?xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
+ return %matrix : !gpu.mma_matrix<16x16xf16, "AOp">
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.memref.extract_address_computations
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// Simple test: check that we extract the address computation of a ldmatrix into
// a dedicated subview.
// The resulting ldmatrix will loaded from with subview and have only indices set
// to zero.
-// Also the sizes of the view are adjusted to `original size - offset`.
+// The trailing subview dimensions are sized from getAccessedShape().
-// CHECK-DAG: #[[$FOUR_MINUS_OFF_MAP:.*]] = affine_map<()[s0] -> (-s0 + 4)>
-// CHECK-DAG: #[[$THIRTY_TWO_MINUS_OFF_MAP:.*]] = affine_map<()[s0] -> (-s0 + 32)>
// CHECK-LABEL: @test_ldmatrix(
// CHECK-SAME: %[[BASE:[^:]*]]: memref<{{[^,]*}}, 3>,
// CHECK-SAME: %[[DYN_OFFSET0:[^:]*]]: index,
// CHECK-SAME: %[[DYN_OFFSET1:[^:]*]]: index,
// CHECK-SAME: %[[DYN_OFFSET2:[^:]*]]: index)
-// CHECK-DAG: %[[DYN_SIZE0:.*]] = affine.apply #[[$FOUR_MINUS_OFF_MAP]]()[%[[DYN_OFFSET0]]]
-// CHECK-DAG: %[[DYN_SIZE1:.*]] = affine.apply #[[$THIRTY_TWO_MINUS_OFF_MAP]]()[%[[DYN_OFFSET1]]]
-// CHECK-DAG: %[[DYN_SIZE2:.*]] = affine.apply #[[$THIRTY_TWO_MINUS_OFF_MAP]]()[%[[DYN_OFFSET2]]]
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET0]], %[[DYN_OFFSET1]], %[[DYN_OFFSET2]]] [%[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_SIZE2]]] [1, 1, 1] : memref<4x32x32xf16, 3> to memref<?x?x?xf16, strided<[1024, 32, 1], offset: ?>, 3>
-// CHECK: %[[LOADED_VAL:.*]] = nvgpu.ldmatrix %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] {numTiles = 4 : i32, transpose = false} : memref<?x?x?xf16, strided<[1024, 32, 1], offset: ?>, 3> -> vector<4x2xf16>
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET0]], %[[DYN_OFFSET1]], %[[DYN_OFFSET2]]] [1, 1, 8] [1, 1, 1] : memref<4x32x32xf16, 3> to memref<1x1x8xf16, strided<[1024, 32, 1], offset: ?>, 3>
+// CHECK: %[[LOADED_VAL:.*]] = nvgpu.ldmatrix %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] {numTiles = 4 : i32, transpose = false} : memref<1x1x8xf16, strided<[1024, 32, 1], offset: ?>, 3> -> vector<4x2xf16>
// CHECK: return %[[LOADED_VAL]] : vector<4x2xf16>
func.func @test_ldmatrix(%base : memref<4x32x32xf16, 3>,
%offset0 : index, %offset1: index, %offset2: index)
@@ -226,23 +338,22 @@ module attributes {transform.with_named_sequence} {
// -----
-// Same as test_ldmatrix but with fully dynamic memref.
+// Same as test_ldmatrix but with fully dynamic memref. The accessed shape is
+// capped by the maximal in-bounds size for the dynamic source dimension.
-// CHECK-DAG: #[[$A_MINUS_B_MAP:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>
-// CHECK-LABEL: @test_ldmatrix(
+// CHECK-DAG: #[[$MIN_8_MAP:.*]] = affine_map<()[s0, s1] -> (s0 - s1, 8)>
+// CHECK-LABEL: @test_dynamic_ldmatrix(
// CHECK-SAME: %[[BASE:[^:]*]]: memref<{{[^,]*}}, 3>,
// CHECK-SAME: %[[DYN_OFFSET0:[^:]*]]: index,
// CHECK-SAME: %[[DYN_OFFSET1:[^:]*]]: index,
// CHECK-SAME: %[[DYN_OFFSET2:[^:]*]]: index)
-// CHECK-DAG: {{.*}}, {{.*}}, %[[DYN_SIZES:.*]]:3, {{.*}} = memref.extract_strided_metadata %[[BASE]]
-// CHECK-DAG: %[[DYN_SIZE0:.*]] = affine.apply #[[$A_MINUS_B_MAP]]()[%[[DYN_SIZES]]#0, %[[DYN_OFFSET0]]]
-// CHECK-DAG: %[[DYN_SIZE1:.*]] = affine.apply #[[$A_MINUS_B_MAP]]()[%[[DYN_SIZES]]#1, %[[DYN_OFFSET1]]]
-// CHECK-DAG: %[[DYN_SIZE2:.*]] = affine.apply #[[$A_MINUS_B_MAP]]()[%[[DYN_SIZES]]#2, %[[DYN_OFFSET2]]]
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET0]], %[[DYN_OFFSET1]], %[[DYN_OFFSET2]]] [%[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_SIZE2]]] [1, 1, 1] : memref<?x?x?xf16, 3> to memref<?x?x?xf16, strided<[?, ?, 1], offset: ?>, 3>
-// CHECK: %[[LOADED_VAL:.*]] = nvgpu.ldmatrix %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] {numTiles = 4 : i32, transpose = false} : memref<?x?x?xf16, strided<[?, ?, 1], offset: ?>, 3> -> vector<4x2xf16>
+// CHECK-DAG: {{.*}}, {{.*}}, %[[SIZES:.*]]:3, {{.*}} = memref.extract_strided_metadata %[[BASE]]
+// CHECK-DAG: %[[DYN_SIZE:.*]] = affine.min #[[$MIN_8_MAP]]()[%[[SIZES]]#2, %[[DYN_OFFSET2]]]
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET0]], %[[DYN_OFFSET1]], %[[DYN_OFFSET2]]] [1, 1, %[[DYN_SIZE]]] [1, 1, 1] : memref<?x?x?xf16, 3> to memref<1x1x?xf16, strided<[?, ?, 1], offset: ?>, 3>
+// CHECK: %[[LOADED_VAL:.*]] = nvgpu.ldmatrix %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] {numTiles = 4 : i32, transpose = false} : memref<1x1x?xf16, strided<[?, ?, 1], offset: ?>, 3> -> vector<4x2xf16>
// CHECK: return %[[LOADED_VAL]] : vector<4x2xf16>
-func.func @test_ldmatrix(%base : memref<?x?x?xf16, 3>,
+func.func @test_dynamic_ldmatrix(%base : memref<?x?x?xf16, 3>,
%offset0 : index, %offset1: index, %offset2: index)
-> vector<4x2xf16> {
%loaded_val = nvgpu.ldmatrix
@@ -301,6 +412,7 @@ module attributes {transform.with_named_sequence} {
}
}
+
// -----
// Same as test_transfer_read_op but with tensors.
@@ -440,4 +552,3 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
-
diff --git a/mlir/test/lib/Dialect/Vector/CMakeLists.txt b/mlir/test/lib/Dialect/Vector/CMakeLists.txt
index 2ba147941969c..385ef2f00ea5d 100644
--- a/mlir/test/lib/Dialect/Vector/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/Vector/CMakeLists.txt
@@ -14,6 +14,7 @@ mlir_target_link_libraries(MLIRVectorTestPasses PUBLIC
MLIRLinalgTransforms
MLIRLLVMDialect
MLIRMemRefDialect
+ MLIRNVGPUDialect
MLIRPass
MLIRSCFDialect
MLIRTensorDialect
>From 547e097f59584a8b619049b39a1bd19575d141fd Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 18 May 2026 23:50:02 +0000
Subject: [PATCH 2/3] clang-format
---
.../lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp b/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
index 306b708b22cbf..6c66f246331ad 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
@@ -171,7 +171,6 @@ struct IndexedAccessOpRewriter final
if (!op.hasInboundsIndices())
return rewriter.notifyMatchFailure(op, "indices may be out of bounds");
-
// If the access already has only zeros as indices there is nothing
// to do.
if (hasAllZeroIndices(op.getIndices()))
>From 206c23b305c9162cd0a440d442ed6955b52af7d1 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Wed, 17 Jun 2026 20:37:00 +0000
Subject: [PATCH 3/3] Remove blank line from tests
---
mlir/test/Dialect/MemRef/extract-address-computations.mlir | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/test/Dialect/MemRef/extract-address-computations.mlir b/mlir/test/Dialect/MemRef/extract-address-computations.mlir
index 44737fd23fce2..adf95bb1d7d54 100644
--- a/mlir/test/Dialect/MemRef/extract-address-computations.mlir
+++ b/mlir/test/Dialect/MemRef/extract-address-computations.mlir
@@ -412,7 +412,6 @@ module attributes {transform.with_named_sequence} {
}
}
-
// -----
// Same as test_transfer_read_op but with tensors.
More information about the Mlir-commits
mailing list