[Mlir-commits] [mlir] [mlir][memref] Drop innaccurate bounds on collapse_shape delinearize (PR #197041)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon May 11 14:52:13 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-memref
Author: Krzysztof Drewniak (krzysz00)
<details>
<summary>Changes</summary>
If (as with something like `vector.load`) we know that the indices that are indexnig a `collapse_shape`'d memref are not necessarily within the bounds of the product of tose collapsed dimensions, we shouldn't be setting the outrmost bound on the
`affine.delinearize_index` we use to split up those indices, as that would incorrectly assert in-bounds-ness per the semantics of `affine.delinearize_index`.
Resolve this by giving `resolveSourceIndicesCollapseShape` a `startsInbounds` parameter by analogy to the one for expand shape, and conservatively set that to false for AMDGPU ops for now.
AI: An LLM spotted this isssue while I was implementing the indexed access ops for AMDGPU operations, but I made the changes by hand.
---
Full diff: https://github.com/llvm/llvm-project/pull/197041.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h (+6-1)
- (modified) mlir/lib/Dialect/AMDGPU/Transforms/FoldMemRefsOps.cpp (+4-1)
- (modified) mlir/lib/Dialect/Affine/Transforms/FoldMemRefAliasOps.cpp (+4-2)
- (modified) mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp (+8-4)
- (modified) mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp (+10-4)
- (modified) mlir/test/Dialect/AMDGPU/amdgpu-fold-memrefs.mlir (+5-5)
- (modified) mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir (+34-4)
``````````diff
diff --git a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
index 9af0f301d763c..154e54e3939c5 100644
--- a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
+++ b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
@@ -152,10 +152,15 @@ void resolveSourceIndicesExpandShape(Location loc, PatternRewriter &rewriter,
///
/// %2 = load %0[%i1 / 6, %i1 % 6, %i2] :
/// memref<2x6x42xf32>
+///
+/// If `startsInbounds` is true, optimizations that rely on all indices being
+/// non-negative and less than the corresponding memref dimension may be
+/// performed.
void resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,
memref::CollapseShapeOp collapseShapeOp,
ValueRange indices,
- SmallVectorImpl<Value> &sourceIndices);
+ SmallVectorImpl<Value> &sourceIndices,
+ bool startsInbounds);
/// Given the 'indices' of a load/store operation where the memref is a result
/// of a rank-reducing full subview op, returns the indices w.r.t to the source
diff --git a/mlir/lib/Dialect/AMDGPU/Transforms/FoldMemRefsOps.cpp b/mlir/lib/Dialect/AMDGPU/Transforms/FoldMemRefsOps.cpp
index 53cb673ced999..24c30525957c7 100644
--- a/mlir/lib/Dialect/AMDGPU/Transforms/FoldMemRefsOps.cpp
+++ b/mlir/lib/Dialect/AMDGPU/Transforms/FoldMemRefsOps.cpp
@@ -46,14 +46,17 @@ static LogicalResult foldMemrefViewOp(PatternRewriter &rewriter, Location loc,
return success();
})
.Case([&](memref::ExpandShapeOp expandShapeOp) {
+ // The lack of inbounds is conservative and will be fixed.
mlir::memref::resolveSourceIndicesExpandShape(
loc, rewriter, expandShapeOp, indices, resolvedIndices, false);
memrefBase = expandShapeOp.getViewSource();
return success();
})
.Case([&](memref::CollapseShapeOp collapseShapeOp) {
+ // The collapse shape in-bounds-ness is defaulted to false
+ // conservatively.
mlir::memref::resolveSourceIndicesCollapseShape(
- loc, rewriter, collapseShapeOp, indices, resolvedIndices);
+ loc, rewriter, collapseShapeOp, indices, resolvedIndices, false);
memrefBase = collapseShapeOp.getViewSource();
return success();
})
diff --git a/mlir/lib/Dialect/Affine/Transforms/FoldMemRefAliasOps.cpp b/mlir/lib/Dialect/Affine/Transforms/FoldMemRefAliasOps.cpp
index 6f6e40f586fc8..bb7b231ea02c2 100644
--- a/mlir/lib/Dialect/Affine/Transforms/FoldMemRefAliasOps.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/FoldMemRefAliasOps.cpp
@@ -128,7 +128,8 @@ struct AffineLoadOpOfCollapseShapeOpFolder final
SmallVector<Value> sourceIndices;
memref::resolveSourceIndicesCollapseShape(
- loadOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices);
+ loadOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices,
+ /*startsInbounds=*/true);
rewriter.replaceOpWithNewOp<AffineLoadOp>(
loadOp, collapseShapeOp.getViewSource(), sourceIndices);
@@ -212,7 +213,8 @@ struct AffineStoreOpOfCollapseShapeOpFolder final
SmallVector<Value> sourceIndices;
memref::resolveSourceIndicesCollapseShape(
- storeOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices);
+ storeOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices,
+ /*startsInbounds=*/true);
rewriter.replaceOpWithNewOp<AffineStoreOp>(
storeOp, storeOp.getValueToStore(), collapseShapeOp.getViewSource(),
diff --git a/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp b/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
index de7662753d142..f5c5a48e7f543 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
@@ -316,7 +316,8 @@ LogicalResult AccessOpOfCollapseShapeOpFolder::matchAndRewrite(
SmallVector<Value> sourceIndices;
memref::resolveSourceIndicesCollapseShape(op.getLoc(), rewriter, collapse,
- op.getIndices(), sourceIndices);
+ op.getIndices(), sourceIndices,
+ op.hasInboundsIndices());
std::optional<SmallVector<Value>> newValues = op.updateMemrefAndIndices(
rewriter, collapse.getViewSource(), sourceIndices);
@@ -405,13 +406,15 @@ LogicalResult IndexedMemCopyOpOfCollapseShapeOpFolder::matchAndRewrite(
newSrc = srcCollapse.getViewSource();
newSrcIndices.clear();
memref::resolveSourceIndicesCollapseShape(
- op.getLoc(), rewriter, srcCollapse, op.getSrcIndices(), newSrcIndices);
+ op.getLoc(), rewriter, srcCollapse, op.getSrcIndices(), newSrcIndices,
+ /*startsInbounds=*/true);
}
if (dstCollapse) {
newDst = dstCollapse.getViewSource();
newDstIndices.clear();
memref::resolveSourceIndicesCollapseShape(
- op.getLoc(), rewriter, dstCollapse, op.getDstIndices(), newDstIndices);
+ op.getLoc(), rewriter, dstCollapse, op.getDstIndices(), newDstIndices,
+ /*startsInbounds=*/true);
}
op.setMemrefsAndIndices(rewriter, newSrc, newSrcIndices, newDst,
newDstIndices);
@@ -550,7 +553,8 @@ LogicalResult TransferOpOfCollapseShapeOpFolder::matchAndRewrite(
SmallVector<Value> newIndices;
memref::resolveSourceIndicesCollapseShape(op.getLoc(), rewriter, collapse,
- op.getIndices(), newIndices);
+ op.getIndices(), newIndices,
+ /*startsInbounds=*/!op.getMask());
op.updateStartingPosition(rewriter, collapse.getViewSource(), newIndices,
AffineMapAttr::get(newPerm));
diff --git a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
index cf126cd85ddce..54fb02d6d8ba8 100644
--- a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
+++ b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
@@ -252,7 +252,8 @@ void resolveSourceIndicesExpandShape(Location loc, PatternRewriter &rewriter,
void resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,
memref::CollapseShapeOp collapseShapeOp,
ValueRange indices,
- SmallVectorImpl<Value> &sourceIndices) {
+ SmallVectorImpl<Value> &sourceIndices,
+ bool startsInbounds) {
// Note: collapse_shape requires a strided memref, we can do this.
auto metadata = memref::ExtractStridedMetadataOp::create(
rewriter, loc, collapseShapeOp.getSrc());
@@ -267,10 +268,15 @@ void resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,
continue;
}
- SmallVector<OpFoldResult> basis =
- llvm::map_to_vector(group, [&](int64_t d) { return sourceSizes[d]; });
+ // If we don't know that this value is in-bounds, the largest return value
+ // of the delinearization may exceed `sourceSizes[d]`, so we drop that first
+ // group entry in order to maintain soundness.
+ auto trimmedGroup =
+ ArrayRef<int64_t>(group).drop_front(startsInbounds ? 0 : 1);
+ SmallVector<OpFoldResult> basis = llvm::map_to_vector(
+ trimmedGroup, [&](int64_t d) { return sourceSizes[d]; });
auto delinearize = affine::AffineDelinearizeIndexOp::create(
- rewriter, loc, index, basis, /*hasOuterBound=*/true);
+ rewriter, loc, index, basis, /*hasOuterBound=*/startsInbounds);
llvm::append_range(sourceIndices, delinearize.getResults());
}
if (collapseShapeOp.getReassociationIndices().empty()) {
diff --git a/mlir/test/Dialect/AMDGPU/amdgpu-fold-memrefs.mlir b/mlir/test/Dialect/AMDGPU/amdgpu-fold-memrefs.mlir
index 1274fe59f8be5..82f76953e2522 100644
--- a/mlir/test/Dialect/AMDGPU/amdgpu-fold-memrefs.mlir
+++ b/mlir/test/Dialect/AMDGPU/amdgpu-fold-memrefs.mlir
@@ -82,8 +82,8 @@ func.func @test_collapse_shape(%offset_i: index, %offset_j: index) {
// CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, #gpu.address_space<workgroup>>
// CHECK: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>
// CHECK: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[INDICES_MEM:.*]]:2 = affine.delinearize_index %[[ARG0]] into (64, 128) : index, index
- // CHECK: %[[INDICES_LDS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (64, 64) : index, index
+ // CHECK: %[[INDICES_MEM:.*]]:2 = affine.delinearize_index %[[ARG0]] into (128) : index, index
+ // CHECK: %[[INDICES_LDS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (64) : index, index
// CHECK: amdgpu.gather_to_lds %[[MEM]][%[[INDICES_MEM]]#0, %[[INDICES_MEM]]#1], %[[LOCAL]][%[[INDICES_LDS]]#0, %[[INDICES_LDS]]#1]
// CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, #gpu.address_space<workgroup>>
@@ -256,7 +256,7 @@ func.func @test_transpose_load_expand_shape(%offset_i: index, %offset_j: index)
// CHECK-SAME: %[[ARG0:.*]]: index
func.func @test_transpose_load_collapse_shape(%offset_i: index) -> vector<4xf16> {
// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<32x128xf16, #gpu.address_space<workgroup>>
- // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[ARG0]] into (32, 128) : index, index
+ // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[ARG0]] into (128) : index, index
// CHECK: amdgpu.transpose_load %[[ALLOC]][%[[INDICES]]#0, %[[INDICES]]#1]
// CHECK-SAME: memref<32x128xf16, #gpu.address_space<workgroup>> -> vector<4xf16>
@@ -347,7 +347,7 @@ func.func @test_make_gather_dma_base_subview(%mem: memref<64x128xf16, #gpu_globa
// CHECK: func @test_make_gather_dma_base_collapse_shape
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_IDX:.*]]: index
func.func @test_make_gather_dma_base_collapse_shape(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<64x64xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_idx: index) {
- // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[LDS_IDX]] into (64, 64) : index, index
+ // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[LDS_IDX]] into (64) : index, index
// CHECK: amdgpu.make_gather_dma_base %[[MEM]][%[[GLOBAL_I]], %[[GLOBAL_J]]], %[[LDS]][%[[INDICES]]#0, %[[INDICES]]#1]
// CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<64x64xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_gather_base<f16, i16>
@@ -456,7 +456,7 @@ func.func @test_global_load_async_to_lds_both_fold_masked(%src: memref<64x128xf3
// CHECK: func @test_global_load_async_to_lds_no_mask_dst_collapse
// CHECK-SAME: %[[SRC:.*]]: memref<8192xi32, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xi32, #gpu.address_space<workgroup>>, %[[SRC_IDX:.*]]: index, %[[DST_IDX:.*]]: index
func.func @test_global_load_async_to_lds_no_mask_dst_collapse(%src: memref<8192xi32, #gpu_global_addrspace>, %lds: memref<64x64xi32, #gpu_lds_addrspace>, %src_idx: index, %dst_idx: index) {
- // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[DST_IDX]] into (64, 64) : index, index
+ // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[DST_IDX]] into (64) : index, index
// CHECK: amdgpu.global_load_async_to_lds %[[SRC]][%[[SRC_IDX]]], %[[LDS]][%[[INDICES]]#0, %[[INDICES]]#1] :
// CHECK-SAME: i32, memref<8192xi32, #gpu.address_space<global>>, memref<64x64xi32, #gpu.address_space<workgroup>>
diff --git a/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir b/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
index 6e2702d936ee0..48241a3e5d9df 100644
--- a/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
+++ b/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
@@ -988,6 +988,21 @@ func.func @fold_vector_transfer_read_rank_mismatch(
// -----
+func.func @fold_memref_load_collapse_shape(
+ %arg0 : memref<4x8xf32>, %arg1 : index) -> f32 {
+ %0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>
+ %1 = memref.load %0[%arg1] {nontemporal = true} : memref<32xf32>
+ return %1 : f32
+}
+
+// CHECK-LABEL: func @fold_memref_load_collapse_shape
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
+// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)
+// CHECK: memref.load %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1] {nontemporal = true}
+
+// -----
+
func.func @fold_vector_load_collapse_shape(
%arg0 : memref<4x8xf32>, %arg1 : index) -> vector<8xf32> {
%0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>
@@ -998,7 +1013,7 @@ func.func @fold_vector_load_collapse_shape(
// CHECK-LABEL: func @fold_vector_load_collapse_shape
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
-// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)
+// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (8)
// CHECK: vector.load %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1] {nontemporal = true}
// -----
@@ -1015,11 +1030,26 @@ func.func @fold_vector_maskedload_collapse_shape(
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<8xi1>
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<8xf32>
-// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)
+// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (8)
// CHECK: vector.maskedload %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1], %[[ARG3]], %[[ARG4]]
// -----
+func.func @fold_memref_store_collapse_shape(
+ %arg0 : memref<4x8xf32>, %arg1 : index, %val : f32) {
+ %0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>
+ memref.store %val, %0[%arg1] {nontemporal = true} : memref<32xf32>
+ return
+}
+
+// CHECK-LABEL: func @fold_memref_store_collapse_shape
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
+// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)
+// CHECK: memref.store %{{.*}}, %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1] {nontemporal = true}
+
+// -----
+
func.func @fold_vector_store_collapse_shape(
%arg0 : memref<4x8xf32>, %arg1 : index, %val : vector<8xf32>) {
%0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>
@@ -1030,7 +1060,7 @@ func.func @fold_vector_store_collapse_shape(
// CHECK-LABEL: func @fold_vector_store_collapse_shape
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
-// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)
+// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (8)
// CHECK: vector.store %{{.*}}, %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1] {nontemporal = true}
// -----
@@ -1047,7 +1077,7 @@ func.func @fold_vector_maskedstore_collapse_shape(
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<8xi1>
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<8xf32>
-// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)
+// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (8)
// CHECK: vector.maskedstore %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1], %[[ARG3]], %[[ARG4]]
// -----
``````````
</details>
https://github.com/llvm/llvm-project/pull/197041
More information about the Mlir-commits
mailing list