[Mlir-commits] [mlir] [mlir][spirv] Add in-bounds access chain conversion (PR #216096)
Hsiangkai Wang
llvmlistbot at llvm.org
Thu Aug 13 08:59:28 PDT 2026
https://github.com/Hsiangkai created https://github.com/llvm/llvm-project/pull/216096
Add spirv.InBoundsAccessChain and use it for static StorageBuffer
accesses whose linearized range fits the declared SPIR-V object.
Keep plain access chains for dynamic layouts and packed sub-16-bit storage.
>From 29e8add1cd3723d3c244ccf783344eb48ae147a5 Mon Sep 17 00:00:00 2001
From: Hsiangkai Wang <hsiangkai.wang at arm.com>
Date: Wed, 12 Aug 2026 14:46:16 +0100
Subject: [PATCH 1/2] [mlir][spirv] Mark in-bounds linearized indices no-wrap
Preserve signed and unsigned no-wrap guarantees for statically bounded
SPIR-V index linearization, allowing buffer address arithmetic to be
reassociated safely.
---
.../SPIRV/Transforms/SPIRVConversion.h | 17 ++++-
.../TensorToSPIRV/TensorToSPIRV.cpp | 8 +-
.../SPIRV/Transforms/SPIRVConversion.cpp | 74 +++++++++++++++++--
.../MemRefToSPIRV/memref-to-spirv.mlir | 50 +++++++++++++
4 files changed, 142 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
index 03ae54a8ae30a..32fe0062cf168 100644
--- a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
+++ b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
@@ -168,11 +168,26 @@ Value getPushConstantValue(Operation *op, unsigned elementCount,
unsigned offset, Type integerType,
OpBuilder &builder);
+/// No-wrap guarantees proven for a linearized index calculation.
+struct LinearizedIndexNoWrapFlags {
+ bool noSignedWrap = false;
+ bool noUnsignedWrap = false;
+};
+
/// Generates IR to perform index linearization with the given `indices` and
/// their corresponding `strides`, adding an initial `offset`.
Value linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
int64_t offset, Type integerType, Location loc,
- OpBuilder &builder);
+ OpBuilder &builder,
+ LinearizedIndexNoWrapFlags noWrapFlags = {});
+
+/// Returns no-wrap guarantees for an in-bounds index into the static layout
+/// described by `shape`, `strides`, and `offset` when linearized as
+/// `integerType`.
+LinearizedIndexNoWrapFlags
+getLinearizedIndexNoWrapFlags(ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> strides, int64_t offset,
+ Type integerType);
/// Performs the index computation to get to the element at `indices` of the
/// memory pointed to by `basePtr`, using the layout map of `baseType`.
diff --git a/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp b/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
index f24972f6b6ee1..a5e5b7c3c658f 100644
--- a/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
+++ b/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
@@ -81,8 +81,14 @@ class TensorExtractPattern final
auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
auto indexType = typeConverter.getIndexType();
+ spirv::LinearizedIndexNoWrapFlags noWrapFlags;
+ if (typeConverter.getTargetEnv().allows(
+ spirv::Extension::SPV_KHR_no_integer_wrap_decoration))
+ noWrapFlags = spirv::getLinearizedIndexNoWrapFlags(
+ tensorType.getShape(), strides, /*offset=*/0, indexType);
Value index = spirv::linearizeIndex(adaptor.getIndices(), strides,
- /*offset=*/0, indexType, loc, rewriter);
+ /*offset=*/0, indexType, loc, rewriter,
+ noWrapFlags);
auto acOp = spirv::AccessChainOp::create(rewriter, loc, varOp, index);
rewriter.replaceOpWithNewOp<spirv::LoadOp>(extractOp, acOp);
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index ef4d79c827bb2..0f823b9505c2a 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -28,9 +28,11 @@
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/APInt.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/CheckedArithmetic.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/MathExtras.h"
@@ -1283,6 +1285,28 @@ struct ReturnOpVectorUnroll final : OpRewritePattern<func::ReturnOp> {
}
};
+static void addNoWrapDecorations(Operation *op,
+ spirv::LinearizedIndexNoWrapFlags flags,
+ OpBuilder &builder) {
+ if (flags.noSignedWrap)
+ op->setAttr(spirv::getDecorationString(spirv::Decoration::NoSignedWrap),
+ builder.getUnitAttr());
+ if (flags.noUnsignedWrap)
+ op->setAttr(spirv::getDecorationString(spirv::Decoration::NoUnsignedWrap),
+ builder.getUnitAttr());
+}
+
+static spirv::LinearizedIndexNoWrapFlags
+shouldEmitNoWrapDecorations(const SPIRVTypeConverter &typeConverter,
+ MemRefType baseType, ArrayRef<int64_t> strides,
+ int64_t offset, Type indexType) {
+ if (!typeConverter.getTargetEnv().allows(
+ spirv::Extension::SPV_KHR_no_integer_wrap_decoration))
+ return {};
+ return spirv::getLinearizedIndexNoWrapFlags(baseType.getShape(), strides,
+ offset, indexType);
+}
+
} // namespace
//===----------------------------------------------------------------------===//
@@ -1336,9 +1360,38 @@ Value spirv::getPushConstantValue(Operation *op, unsigned elementCount,
// Public functions for index calculation
//===----------------------------------------------------------------------===//
+mlir::spirv::LinearizedIndexNoWrapFlags
+mlir::spirv::getLinearizedIndexNoWrapFlags(ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> strides,
+ int64_t offset, Type integerType) {
+ LinearizedIndexNoWrapFlags flags;
+ auto integer = dyn_cast<IntegerType>(integerType);
+ if (!integer || shape.size() != strides.size() || offset < 0)
+ return flags;
+
+ std::optional<uint64_t> maxLinearIndex = static_cast<uint64_t>(offset);
+ for (auto [dimension, stride] : llvm::zip(shape, strides)) {
+ if (dimension <= 0 || stride < 0)
+ return flags;
+ maxLinearIndex = llvm::checkedMulAddUnsigned(
+ static_cast<uint64_t>(dimension - 1), static_cast<uint64_t>(stride),
+ *maxLinearIndex);
+ if (!maxLinearIndex)
+ return flags;
+ }
+
+ flags.noSignedWrap =
+ *maxLinearIndex <=
+ APInt::getSignedMaxValue(integer.getWidth()).getZExtValue();
+ flags.noUnsignedWrap =
+ *maxLinearIndex <= APInt::getMaxValue(integer.getWidth()).getZExtValue();
+ return flags;
+}
+
Value mlir::spirv::linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
int64_t offset, Type integerType,
- Location loc, OpBuilder &builder) {
+ Location loc, OpBuilder &builder,
+ LinearizedIndexNoWrapFlags noWrapFlags) {
assert(indices.size() == strides.size() &&
"must provide indices for all dimensions");
@@ -1355,8 +1408,15 @@ Value mlir::spirv::linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
IntegerAttr::get(integerType, strides[index.index()]));
Value update =
builder.createOrFold<spirv::IMulOp>(loc, index.value(), strideVal);
+ if (noWrapFlags.noSignedWrap || noWrapFlags.noUnsignedWrap)
+ if (auto mul = update.getDefiningOp<spirv::IMulOp>())
+ addNoWrapDecorations(mul, noWrapFlags, builder);
+
linearizedIndex =
builder.createOrFold<spirv::IAddOp>(loc, update, linearizedIndex);
+ if (noWrapFlags.noSignedWrap || noWrapFlags.noUnsignedWrap)
+ if (auto add = linearizedIndex.getDefiningOp<spirv::IAddOp>())
+ addNoWrapDecorations(add, noWrapFlags, builder);
}
return linearizedIndex;
}
@@ -1376,6 +1436,8 @@ Value mlir::spirv::getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
}
auto indexType = typeConverter.getIndexType();
+ LinearizedIndexNoWrapFlags noWrapFlags = shouldEmitNoWrapDecorations(
+ typeConverter, baseType, strides, offset, indexType);
SmallVector<Value, 2> linearizedIndices;
auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
@@ -1383,8 +1445,8 @@ Value mlir::spirv::getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
if (baseType.getRank() == 0) {
linearizedIndices.push_back(zero);
} else {
- linearizedIndices.push_back(
- linearizeIndex(indices, strides, offset, indexType, loc, builder));
+ linearizedIndices.push_back(linearizeIndex(
+ indices, strides, offset, indexType, loc, builder, noWrapFlags));
}
const Type pointeeType =
@@ -1410,14 +1472,16 @@ Value mlir::spirv::getOpenCLElementPtr(const SPIRVTypeConverter &typeConverter,
}
auto indexType = typeConverter.getIndexType();
+ LinearizedIndexNoWrapFlags noWrapFlags = shouldEmitNoWrapDecorations(
+ typeConverter, baseType, strides, offset, indexType);
SmallVector<Value, 2> linearizedIndices;
Value linearIndex;
if (baseType.getRank() == 0) {
linearIndex = spirv::ConstantOp::getZero(indexType, loc, builder);
} else {
- linearIndex =
- linearizeIndex(indices, strides, offset, indexType, loc, builder);
+ linearIndex = linearizeIndex(indices, strides, offset, indexType, loc,
+ builder, noWrapFlags);
}
Type pointeeType =
cast<spirv::PointerType>(basePtr.getType()).getPointeeType();
diff --git a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
index 5163120a8339e..5550c778b6f31 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -220,6 +220,56 @@ func.func @load_store_vec4f32_dynamic_physical(%arg0: memref<?xvector<4xf32>, #s
// -----
+module attributes {
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.1, [Shader], [SPV_KHR_storage_buffer_storage_class, SPV_KHR_no_integer_wrap_decoration]>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @static_linearized_index
+func.func @static_linearized_index(
+ %arg0: memref<2x4xf32, #spirv.storage_class<StorageBuffer>>,
+ %row: index, %column: index) -> f32 {
+ // CHECK: %[[STRIDE:.+]] = spirv.Constant 4 : i32
+ // CHECK: %[[OFFSET:.+]] = spirv.IMul %{{.*}}, %[[STRIDE]] {no_signed_wrap, no_unsigned_wrap} : i32
+ // CHECK: %[[LINEAR:.+]] = spirv.IAdd %{{.*}}, %[[OFFSET]] {no_signed_wrap, no_unsigned_wrap} : i32
+ // CHECK: spirv.AccessChain {{.*}}[%{{.*}}, %[[LINEAR]]]
+ %0 = memref.load %arg0[%row, %column] : memref<2x4xf32, #spirv.storage_class<StorageBuffer>>
+ return %0 : f32
+}
+
+// The maximum linearized index is 2147483649: too large for a signed i32,
+// but representable in an unsigned i32.
+// CHECK-LABEL: @unsigned_only_linearized_index
+func.func @unsigned_only_linearized_index(
+ %arg0: memref<2x1073741825xf32, #spirv.storage_class<StorageBuffer>>,
+ %row: index, %column: index) -> f32 {
+ // CHECK: %[[STRIDE:.+]] = spirv.Constant 1073741825 : i32
+ // CHECK: %[[OFFSET:.+]] = spirv.IMul %{{.*}}, %[[STRIDE]] {no_unsigned_wrap} : i32
+ // CHECK-NOT: no_signed_wrap
+ // CHECK: %[[LINEAR:.+]] = spirv.IAdd %{{.*}}, %[[OFFSET]] {no_unsigned_wrap} : i32
+ // CHECK-NOT: no_signed_wrap
+ // CHECK: spirv.AccessChain {{.*}}[%{{.*}}, %[[LINEAR]]]
+ %0 = memref.load %arg0[%row, %column] : memref<2x1073741825xf32, #spirv.storage_class<StorageBuffer>>
+ return %0 : f32
+}
+
+// Do not decorate index arithmetic whose layout bounds cannot be proven.
+// CHECK-LABEL: @dynamic_linearized_index
+func.func @dynamic_linearized_index(
+ %arg0: memref<?x4xf32, #spirv.storage_class<StorageBuffer>>,
+ %row: index, %column: index) -> f32 {
+ // CHECK: %[[STRIDE:.+]] = spirv.Constant 4 : i32
+ // CHECK: %[[OFFSET:.+]] = spirv.IMul {{.*}}, %[[STRIDE]] : i32
+ // CHECK-NOT: no_signed_wrap
+ // CHECK-NOT: no_unsigned_wrap
+ // CHECK: spirv.IAdd {{.*}}, %[[OFFSET]] : i32
+ %0 = memref.load %arg0[%row, %column] : memref<?x4xf32, #spirv.storage_class<StorageBuffer>>
+ return %0 : f32
+}
+
+}
+
+// -----
+
// Check for Kernel capability, that with proper compute and storage extensions, we don't need to
// perform special tricks.
>From b2ea368610ac978e547a723ad7cb1445fbe090dc Mon Sep 17 00:00:00 2001
From: Hsiangkai Wang <hsiangkai.wang at arm.com>
Date: Thu, 13 Aug 2026 16:11:57 +0100
Subject: [PATCH 2/2] [mlir][spirv] Add in-bounds access chain conversion
Add spirv.InBoundsAccessChain and use it for static StorageBuffer
accesses whose linearized range fits the declared SPIR-V object.
Keep plain access chains for dynamic layouts and packed sub-16-bit storage.
---
.../mlir/Dialect/SPIRV/IR/SPIRVBase.td | 3 +-
.../mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td | 34 ++++++
.../SPIRV/Transforms/SPIRVConversion.h | 14 +++
.../VectorToSPIRV/VectorToSPIRV.cpp | 12 +-
mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp | 15 +++
.../SPIRV/IR/SPIRVCanonicalization.cpp | 34 ++++++
.../SPIRV/Transforms/SPIRVConversion.cpp | 107 +++++++++++++++---
.../MemRefToSPIRV/memref-to-spirv.mlir | 12 +-
.../VectorToSPIRV/vector-to-spirv.mlir | 23 +++-
mlir/test/Dialect/SPIRV/IR/memory-ops.mlir | 11 ++
.../SPIRV/Transforms/canonicalize.mlir | 19 ++++
mlir/test/Target/SPIRV/memory-ops.mlir | 12 +-
12 files changed, 262 insertions(+), 34 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index 68a0ee470709d..34d0d0b2b4d3b 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4492,6 +4492,7 @@ def SPIRV_OC_OpLoad : I32EnumAttrCase<"OpLoad", 61>;
def SPIRV_OC_OpStore : I32EnumAttrCase<"OpStore", 62>;
def SPIRV_OC_OpCopyMemory : I32EnumAttrCase<"OpCopyMemory", 63>;
def SPIRV_OC_OpAccessChain : I32EnumAttrCase<"OpAccessChain", 65>;
+def SPIRV_OC_OpInBoundsAccessChain : I32EnumAttrCase<"OpInBoundsAccessChain", 66>;
def SPIRV_OC_OpPtrAccessChain : I32EnumAttrCase<"OpPtrAccessChain", 67>;
def SPIRV_OC_OpInBoundsPtrAccessChain : I32EnumAttrCase<"OpInBoundsPtrAccessChain", 70>;
def SPIRV_OC_OpDecorate : I32EnumAttrCase<"OpDecorate", 71>;
@@ -4743,7 +4744,7 @@ def SPIRV_OpcodeAttr :
SPIRV_OC_OpSpecConstantOp, SPIRV_OC_OpFunction, SPIRV_OC_OpFunctionParameter,
SPIRV_OC_OpFunctionEnd, SPIRV_OC_OpFunctionCall, SPIRV_OC_OpVariable,
SPIRV_OC_OpLoad, SPIRV_OC_OpStore, SPIRV_OC_OpCopyMemory,
- SPIRV_OC_OpAccessChain, SPIRV_OC_OpPtrAccessChain,
+ SPIRV_OC_OpAccessChain, SPIRV_OC_OpInBoundsAccessChain, SPIRV_OC_OpPtrAccessChain,
SPIRV_OC_OpInBoundsPtrAccessChain, SPIRV_OC_OpDecorate,
SPIRV_OC_OpMemberDecorate, SPIRV_OC_OpVectorExtractDynamic,
SPIRV_OC_OpVectorInsertDynamic, SPIRV_OC_OpVectorShuffle,
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td
index e13909e8eeeae..5ae6c8af83003 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td
@@ -81,6 +81,40 @@ def SPIRV_AccessChainOp : SPIRV_Op<"AccessChain", [Pure]> {
// -----
+def SPIRV_InBoundsAccessChainOp : SPIRV_Op<"InBoundsAccessChain", [Pure]> {
+ let summary = [{
+ Create a pointer into a composite object that is known to stay within the
+ base object.
+ }];
+
+ let description = [{
+ Has the same operands, result, and type rules as `spirv.AccessChain`, with
+ the additional contract that the resulting pointer points within the base
+ object.
+ }];
+
+ let arguments = (ins
+ SPIRV_AnyPtr:$base_ptr,
+ Variadic<SPIRV_Integer>:$indices
+ );
+
+ let results = (outs
+ SPIRV_AnyPtr:$component_ptr
+ );
+
+ let builders = [OpBuilder<(ins "Value":$basePtr, "ValueRange":$indices)>];
+
+ let hasCanonicalizer = 1;
+
+ let hasCustomAssemblyFormat = 0;
+
+ let assemblyFormat = [{
+ $base_ptr `[` $indices `]` attr-dict `:` type($base_ptr) `,` type($indices) `->` type(results)
+ }];
+}
+
+// -----
+
def SPIRV_CopyMemoryOp : SPIRV_Op<"CopyMemory", [DeclareOpInterfaceMethods<AlignmentAttrOpInterface>]> {
let summary = [{
Copy from the memory pointed to by Source to the memory pointed to by
diff --git a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
index 32fe0062cf168..cfd0340fbe510 100644
--- a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
+++ b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
@@ -199,6 +199,13 @@ Value getElementPtr(const SPIRVTypeConverter &typeConverter,
MemRefType baseType, Value basePtr, ValueRange indices,
Location loc, OpBuilder &builder);
+/// As above, with the number of contiguous memref elements accessed through
+/// the pointer. This lets vector conversions retain their full access range.
+Value getElementPtr(const SPIRVTypeConverter &typeConverter,
+ MemRefType baseType, Value basePtr, ValueRange indices,
+ Location loc, OpBuilder &builder,
+ uint64_t accessElementCount);
+
// GetElementPtr implementation for Kernel/OpenCL flavored SPIR-V.
Value getOpenCLElementPtr(const SPIRVTypeConverter &typeConverter,
MemRefType baseType, Value basePtr,
@@ -209,6 +216,13 @@ Value getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
MemRefType baseType, Value basePtr,
ValueRange indices, Location loc, OpBuilder &builder);
+/// As above, with the number of contiguous memref elements accessed through
+/// the pointer.
+Value getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
+ MemRefType baseType, Value basePtr,
+ ValueRange indices, Location loc, OpBuilder &builder,
+ uint64_t accessElementCount);
+
// Find the largest factor of size among {2,3,4} for the lowest dimension of
// the target shape.
int getComputeVectorSize(int64_t size);
diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 78693e924c4d9..0808d13620be0 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
@@ -743,9 +743,9 @@ struct VectorLoadOpConverter final
const auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
auto loc = loadOp.getLoc();
- Value accessChain =
- spirv::getElementPtr(typeConverter, memrefType, adaptor.getBase(),
- adaptor.getIndices(), loc, rewriter);
+ Value accessChain = spirv::getElementPtr(
+ typeConverter, memrefType, adaptor.getBase(), adaptor.getIndices(), loc,
+ rewriter, loadOp.getVectorType().getNumElements());
if (!accessChain)
return rewriter.notifyMatchFailure(
loadOp, "failed to get memref element pointer");
@@ -809,9 +809,9 @@ struct VectorStoreOpConverter final
const auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
auto loc = storeOp.getLoc();
- Value accessChain =
- spirv::getElementPtr(typeConverter, memrefType, adaptor.getBase(),
- adaptor.getIndices(), loc, rewriter);
+ Value accessChain = spirv::getElementPtr(
+ typeConverter, memrefType, adaptor.getBase(), adaptor.getIndices(), loc,
+ rewriter, storeOp.getVectorType().getNumElements());
if (!accessChain)
return rewriter.notifyMatchFailure(
storeOp, "failed to get memref element pointer");
diff --git a/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp b/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp
index f9c03bf3b88c0..4d315fe735aef 100644
--- a/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp
@@ -351,6 +351,21 @@ LogicalResult AccessChainOp::verify() {
return verifyAccessChain(*this, getIndices());
}
+//===----------------------------------------------------------------------===//
+// spirv.InBoundsAccessChainOp
+//===----------------------------------------------------------------------===//
+
+void InBoundsAccessChainOp::build(OpBuilder &builder, OperationState &state,
+ Value basePtr, ValueRange indices) {
+ auto type = getElementPtrType(basePtr.getType(), indices, state.location);
+ assert(type && "Unable to deduce return type based on basePtr and indices");
+ build(builder, state, type, basePtr, indices);
+}
+
+LogicalResult InBoundsAccessChainOp::verify() {
+ return verifyAccessChain(*this, getIndices());
+}
+
//===----------------------------------------------------------------------===//
// spirv.LoadOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
index 2d5c4d7d3fd0e..12347d07abd63 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
@@ -121,6 +121,40 @@ void spirv::AccessChainOp::getCanonicalizationPatterns(
results.add<CombineChainedAccessChain>(context);
}
+namespace {
+
+/// Combines chained `spirv::InBoundsAccessChainOp` operations while retaining
+/// the in-bounds contract of both segments.
+struct CombineChainedInBoundsAccessChain final
+ : OpRewritePattern<spirv::InBoundsAccessChainOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(spirv::InBoundsAccessChainOp accessChainOp,
+ PatternRewriter &rewriter) const override {
+ auto parentAccessChainOp =
+ accessChainOp.getBasePtr()
+ .getDefiningOp<spirv::InBoundsAccessChainOp>();
+
+ if (!parentAccessChainOp)
+ return failure();
+
+ SmallVector<Value, 4> indices(parentAccessChainOp.getIndices());
+ llvm::append_range(indices, accessChainOp.getIndices());
+
+ rewriter.replaceOpWithNewOp<spirv::InBoundsAccessChainOp>(
+ accessChainOp, parentAccessChainOp.getBasePtr(), indices);
+
+ return success();
+ }
+};
+
+} // namespace
+
+void spirv::InBoundsAccessChainOp::getCanonicalizationPatterns(
+ RewritePatternSet &results, MLIRContext *context) {
+ results.add<CombineChainedInBoundsAccessChain>(context);
+}
+
//===----------------------------------------------------------------------===//
// spirv.IAddCarry / spirv.ISubBorrow
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index 0f823b9505c2a..6a27e9d16bb32 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -1296,6 +1296,69 @@ static void addNoWrapDecorations(Operation *op,
builder.getUnitAttr());
}
+static std::optional<uint64_t> getMaxLinearizedIndex(ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> strides,
+ int64_t offset) {
+ if (shape.size() != strides.size() || offset < 0)
+ return std::nullopt;
+
+ uint64_t maxLinearIndex = offset;
+ for (auto [dimension, stride] : llvm::zip(shape, strides)) {
+ if (dimension <= 0 || stride < 0)
+ return std::nullopt;
+ std::optional<uint64_t> nextMaxLinearIndex = llvm::checkedMulAddUnsigned(
+ static_cast<uint64_t>(dimension - 1), static_cast<uint64_t>(stride),
+ maxLinearIndex);
+ if (!nextMaxLinearIndex)
+ return std::nullopt;
+ maxLinearIndex = *nextMaxLinearIndex;
+ }
+ return maxLinearIndex;
+}
+
+static std::optional<uint64_t> getStorageBufferElementCount(Value basePtr) {
+ auto pointerType = dyn_cast<spirv::PointerType>(basePtr.getType());
+ if (!pointerType ||
+ pointerType.getStorageClass() != spirv::StorageClass::StorageBuffer)
+ return std::nullopt;
+
+ Type pointeeType = pointerType.getPointeeType();
+ if (auto structType = dyn_cast<spirv::StructType>(pointeeType)) {
+ if (structType.getNumElements() != 1)
+ return std::nullopt;
+ pointeeType = structType.getElementType(0);
+ }
+ auto arrayType = dyn_cast<spirv::ArrayType>(pointeeType);
+ if (!arrayType)
+ return std::nullopt;
+ return arrayType.getNumElements();
+}
+
+static bool shouldEmitInBoundsAccessChain(MemRefType baseType, Value basePtr,
+ ArrayRef<int64_t> strides,
+ int64_t offset,
+ uint64_t accessElementCount) {
+ // Sub-16-bit integer memrefs may be stored using a wider SPIR-V array element
+ // than the source element. Keep a plain access chain so later bitwidth
+ // emulation can adjust the final index in storage-element units.
+ if (auto integerType = dyn_cast<IntegerType>(baseType.getElementType()))
+ if (integerType.getWidth() < 16)
+ return false;
+
+ std::optional<uint64_t> maxLinearIndex =
+ getMaxLinearizedIndex(baseType.getShape(), strides, offset);
+ std::optional<uint64_t> objectElementCount =
+ getStorageBufferElementCount(basePtr);
+ if (!maxLinearIndex || !objectElementCount || !accessElementCount ||
+ accessElementCount > *objectElementCount)
+ return false;
+
+ // The source memory operation guarantees that its dynamic indices, including
+ // a vector access width, are in bounds. The static layout proof here ensures
+ // that this contract describes the same fixed-size SPIR-V buffer object.
+ return *maxLinearIndex < *objectElementCount;
+}
+
static spirv::LinearizedIndexNoWrapFlags
shouldEmitNoWrapDecorations(const SPIRVTypeConverter &typeConverter,
MemRefType baseType, ArrayRef<int64_t> strides,
@@ -1366,19 +1429,13 @@ mlir::spirv::getLinearizedIndexNoWrapFlags(ArrayRef<int64_t> shape,
int64_t offset, Type integerType) {
LinearizedIndexNoWrapFlags flags;
auto integer = dyn_cast<IntegerType>(integerType);
- if (!integer || shape.size() != strides.size() || offset < 0)
+ if (!integer)
return flags;
- std::optional<uint64_t> maxLinearIndex = static_cast<uint64_t>(offset);
- for (auto [dimension, stride] : llvm::zip(shape, strides)) {
- if (dimension <= 0 || stride < 0)
- return flags;
- maxLinearIndex = llvm::checkedMulAddUnsigned(
- static_cast<uint64_t>(dimension - 1), static_cast<uint64_t>(stride),
- *maxLinearIndex);
- if (!maxLinearIndex)
- return flags;
- }
+ std::optional<uint64_t> maxLinearIndex =
+ getMaxLinearizedIndex(shape, strides, offset);
+ if (!maxLinearIndex)
+ return flags;
flags.noSignedWrap =
*maxLinearIndex <=
@@ -1424,7 +1481,8 @@ Value mlir::spirv::linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
Value mlir::spirv::getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
MemRefType baseType, Value basePtr,
ValueRange indices, Location loc,
- OpBuilder &builder) {
+ OpBuilder &builder,
+ uint64_t accessElementCount) {
// Get base and offset of the MemRefType and verify they are static.
int64_t offset;
@@ -1454,9 +1512,21 @@ Value mlir::spirv::getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
// Interface memrefs are wrapped in a struct: index to its first elem.
if (isa<spirv::StructType>(pointeeType))
linearizedIndices.insert(linearizedIndices.begin(), zero);
+ if (shouldEmitInBoundsAccessChain(baseType, basePtr, strides, offset,
+ accessElementCount))
+ return spirv::InBoundsAccessChainOp::create(builder, loc, basePtr,
+ linearizedIndices);
return spirv::AccessChainOp::create(builder, loc, basePtr, linearizedIndices);
}
+Value mlir::spirv::getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
+ MemRefType baseType, Value basePtr,
+ ValueRange indices, Location loc,
+ OpBuilder &builder) {
+ return getVulkanElementPtr(typeConverter, baseType, basePtr, indices, loc,
+ builder, /*accessElementCount=*/1);
+}
+
Value mlir::spirv::getOpenCLElementPtr(const SPIRVTypeConverter &typeConverter,
MemRefType baseType, Value basePtr,
ValueRange indices, Location loc,
@@ -1497,7 +1567,8 @@ Value mlir::spirv::getOpenCLElementPtr(const SPIRVTypeConverter &typeConverter,
Value mlir::spirv::getElementPtr(const SPIRVTypeConverter &typeConverter,
MemRefType baseType, Value basePtr,
ValueRange indices, Location loc,
- OpBuilder &builder) {
+ OpBuilder &builder,
+ uint64_t accessElementCount) {
if (typeConverter.allows(spirv::Capability::Kernel)) {
return getOpenCLElementPtr(typeConverter, baseType, basePtr, indices, loc,
@@ -1505,7 +1576,15 @@ Value mlir::spirv::getElementPtr(const SPIRVTypeConverter &typeConverter,
}
return getVulkanElementPtr(typeConverter, baseType, basePtr, indices, loc,
- builder);
+ builder, accessElementCount);
+}
+
+Value mlir::spirv::getElementPtr(const SPIRVTypeConverter &typeConverter,
+ MemRefType baseType, Value basePtr,
+ ValueRange indices, Location loc,
+ OpBuilder &builder) {
+ return getElementPtr(typeConverter, baseType, basePtr, indices, loc, builder,
+ /*accessElementCount=*/1);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
index 5550c778b6f31..5588c7a47ac6a 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -24,12 +24,12 @@ func.func @load_store_zero_rank_float(%arg0: memref<f32, #spirv.storage_class<St
// CHECK-DAG: [[ARG0:%.*]] = builtin.unrealized_conversion_cast %[[OARG0]] : memref<f32, #spirv.storage_class<StorageBuffer>> to !spirv.ptr<!spirv.struct<(!spirv.array<1 x f32, stride=4> [0])>, StorageBuffer>
// CHECK-DAG: [[ARG1:%.*]] = builtin.unrealized_conversion_cast %[[OARG1]] : memref<f32, #spirv.storage_class<StorageBuffer>> to !spirv.ptr<!spirv.struct<(!spirv.array<1 x f32, stride=4> [0])>, StorageBuffer>
// CHECK: [[ZERO:%.*]] = spirv.Constant 0 : i32
- // CHECK: spirv.AccessChain [[ARG0]][
+ // CHECK: spirv.InBoundsAccessChain [[ARG0]][
// CHECK-SAME: [[ZERO]], [[ZERO]]
// CHECK-SAME: ] :
// CHECK: spirv.Load "StorageBuffer" %{{.*}} : f32
%0 = memref.load %arg0[] : memref<f32, #spirv.storage_class<StorageBuffer>>
- // CHECK: spirv.AccessChain [[ARG1]][
+ // CHECK: spirv.InBoundsAccessChain [[ARG1]][
// CHECK-SAME: [[ZERO]], [[ZERO]]
// CHECK-SAME: ] :
// CHECK: spirv.Store "StorageBuffer" %{{.*}} : f32
@@ -43,12 +43,12 @@ func.func @load_store_zero_rank_int(%arg0: memref<i32, #spirv.storage_class<Stor
// CHECK-DAG: [[ARG0:%.*]] = builtin.unrealized_conversion_cast %[[OARG0]] : memref<i32, #spirv.storage_class<StorageBuffer>> to !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer>
// CHECK-DAG: [[ARG1:%.*]] = builtin.unrealized_conversion_cast %[[OARG1]] : memref<i32, #spirv.storage_class<StorageBuffer>> to !spirv.ptr<!spirv.struct<(!spirv.array<1 x i32, stride=4> [0])>, StorageBuffer>
// CHECK: [[ZERO:%.*]] = spirv.Constant 0 : i32
- // CHECK: spirv.AccessChain [[ARG0]][
+ // CHECK: spirv.InBoundsAccessChain [[ARG0]][
// CHECK-SAME: [[ZERO]], [[ZERO]]
// CHECK-SAME: ] :
// CHECK: spirv.Load "StorageBuffer" %{{.*}} : i32
%0 = memref.load %arg0[] : memref<i32, #spirv.storage_class<StorageBuffer>>
- // CHECK: spirv.AccessChain [[ARG1]][
+ // CHECK: spirv.InBoundsAccessChain [[ARG1]][
// CHECK-SAME: [[ZERO]], [[ZERO]]
// CHECK-SAME: ] :
// CHECK: spirv.Store "StorageBuffer" %{{.*}} : i32
@@ -231,7 +231,7 @@ func.func @static_linearized_index(
// CHECK: %[[STRIDE:.+]] = spirv.Constant 4 : i32
// CHECK: %[[OFFSET:.+]] = spirv.IMul %{{.*}}, %[[STRIDE]] {no_signed_wrap, no_unsigned_wrap} : i32
// CHECK: %[[LINEAR:.+]] = spirv.IAdd %{{.*}}, %[[OFFSET]] {no_signed_wrap, no_unsigned_wrap} : i32
- // CHECK: spirv.AccessChain {{.*}}[%{{.*}}, %[[LINEAR]]]
+ // CHECK: spirv.InBoundsAccessChain {{.*}}[%{{.*}}, %[[LINEAR]]]
%0 = memref.load %arg0[%row, %column] : memref<2x4xf32, #spirv.storage_class<StorageBuffer>>
return %0 : f32
}
@@ -247,7 +247,7 @@ func.func @unsigned_only_linearized_index(
// CHECK-NOT: no_signed_wrap
// CHECK: %[[LINEAR:.+]] = spirv.IAdd %{{.*}}, %[[OFFSET]] {no_unsigned_wrap} : i32
// CHECK-NOT: no_signed_wrap
- // CHECK: spirv.AccessChain {{.*}}[%{{.*}}, %[[LINEAR]]]
+ // CHECK: spirv.InBoundsAccessChain {{.*}}[%{{.*}}, %[[LINEAR]]]
%0 = memref.load %arg0[%row, %column] : memref<2x1073741825xf32, #spirv.storage_class<StorageBuffer>>
return %0 : f32
}
diff --git a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
index f904dd9d35c37..d8e187e919bf1 100644
--- a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
+++ b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
@@ -1117,7 +1117,7 @@ module attributes {
// CHECK: %[[CST1:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST2:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST3:.+]] = spirv.Constant 1 : i32
-// CHECK: %[[S4:.+]] = spirv.AccessChain %[[S0]][%[[CST1]], %[[S1]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
+// CHECK: %[[S4:.+]] = spirv.InBoundsAccessChain %[[S0]][%[[CST1]], %[[S1]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
// CHECK: %[[S5:.+]] = spirv.Bitcast %[[S4]] : !spirv.ptr<f32, StorageBuffer> to !spirv.ptr<vector<4xf32>, StorageBuffer>
// CHECK: %[[R0:.+]] = spirv.Load "StorageBuffer" %[[S5]] : vector<4xf32>
// CHECK: return %[[R0]] : vector<4xf32>
@@ -1128,6 +1128,17 @@ func.func @vector_load(%arg0 : memref<4xf32, #spirv.storage_class<StorageBuffer>
return %0: vector<4xf32>
}
+// Dynamic StorageBuffer layouts do not identify a fixed base object, so keep
+// the pointer calculation conservative even though the vector access itself
+// has in-bounds source semantics.
+// CHECK-LABEL: @vector_load_dynamic
+// CHECK: spirv.AccessChain {{.*}} !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>
+func.func @vector_load_dynamic(%arg0 : memref<?xf32, #spirv.storage_class<StorageBuffer>>) -> vector<4xf32> {
+ %idx = arith.constant 0 : index
+ %0 = vector.load %arg0[%idx] : memref<?xf32, #spirv.storage_class<StorageBuffer>>, vector<4xf32>
+ return %0: vector<4xf32>
+}
+
// CHECK-LABEL: @vector_load_single_elem
// CHECK-SAME: (%[[ARG0:.*]]: memref<4xf32, #spirv.storage_class<StorageBuffer>>)
@@ -1137,7 +1148,7 @@ func.func @vector_load(%arg0 : memref<4xf32, #spirv.storage_class<StorageBuffer>
// CHECK: %[[CST1:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST2:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST3:.+]] = spirv.Constant 1 : i32
-// CHECK: %[[S4:.+]] = spirv.AccessChain %[[S0]][%[[CST1]], %[[S1]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
+// CHECK: %[[S4:.+]] = spirv.InBoundsAccessChain %[[S0]][%[[CST1]], %[[S1]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
// CHECK: %[[S5:.+]] = spirv.Load "StorageBuffer" %[[S4]] : f32
// CHECK: %[[R0:.+]] = builtin.unrealized_conversion_cast %[[S5]] : f32 to vector<1xf32>
// CHECK: return %[[R0]] : vector<1xf32>
@@ -1170,7 +1181,7 @@ func.func @vector_load_aligned(%arg0 : memref<4xf32, #spirv.storage_class<Storag
// CHECK: %[[S3:.+]] = spirv.IMul %[[S1]], %[[CST4]] : i32
// CHECK: %[[CST1:.+]] = spirv.Constant 1 : i32
// CHECK: %[[S6:.+]] = spirv.IAdd %[[S2]], %[[S3]] : i32
-// CHECK: %[[S7:.+]] = spirv.AccessChain %[[S0]][%[[CST0_1]], %[[S6]]] : !spirv.ptr<!spirv.struct<(!spirv.array<16 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
+// CHECK: %[[S7:.+]] = spirv.InBoundsAccessChain %[[S0]][%[[CST0_1]], %[[S6]]] : !spirv.ptr<!spirv.struct<(!spirv.array<16 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
// CHECK: %[[S8:.+]] = spirv.Bitcast %[[S7]] : !spirv.ptr<f32, StorageBuffer> to !spirv.ptr<vector<4xf32>, StorageBuffer>
// CHECK: %[[R0:.+]] = spirv.Load "StorageBuffer" %[[S8]] : vector<4xf32>
// CHECK: return %[[R0]] : vector<4xf32>
@@ -1190,7 +1201,7 @@ func.func @vector_load_2d(%arg0 : memref<4x4xf32, #spirv.storage_class<StorageBu
// CHECK: %[[CST1:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST2:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST3:.+]] = spirv.Constant 1 : i32
-// CHECK: %[[S4:.+]] = spirv.AccessChain %[[S0]][%[[CST1]], %[[S1]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
+// CHECK: %[[S4:.+]] = spirv.InBoundsAccessChain %[[S0]][%[[CST1]], %[[S1]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
// CHECK: %[[S5:.+]] = spirv.Bitcast %[[S4]] : !spirv.ptr<f32, StorageBuffer> to !spirv.ptr<vector<4xf32>, StorageBuffer>
// CHECK: spirv.Store "StorageBuffer" %[[S5]], %[[ARG1]] : vector<4xf32>
func.func @vector_store(%arg0 : memref<4xf32, #spirv.storage_class<StorageBuffer>>, %arg1 : vector<4xf32>) {
@@ -1218,7 +1229,7 @@ func.func @vector_store_aligned(%arg0 : memref<4xf32, #spirv.storage_class<Stora
// CHECK: %[[CST1:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST2:.+]] = spirv.Constant 0 : i32
// CHECK: %[[CST3:.+]] = spirv.Constant 1 : i32
-// CHECK: %[[S4:.+]] = spirv.AccessChain %[[S0]][%[[CST1]], %[[S2]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<f32, StorageBuffer>
+// CHECK: %[[S4:.+]] = spirv.InBoundsAccessChain %[[S0]][%[[CST1]], %[[S2]]] : !spirv.ptr<!spirv.struct<(!spirv.array<4 x f32, stride=4> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<f32, StorageBuffer>
// CHECK: spirv.Store "StorageBuffer" %[[S4]], %[[S1]] : f32
func.func @vector_store_single_elem(%arg0 : memref<4xf32, #spirv.storage_class<StorageBuffer>>, %arg1 : vector<1xf32>) {
%idx = arith.constant 0 : index
@@ -1240,7 +1251,7 @@ func.func @vector_store_single_elem(%arg0 : memref<4xf32, #spirv.storage_class<S
// CHECK: %[[S3:.+]] = spirv.IMul %[[S1]], %[[CST4]] : i32
// CHECK: %[[CST1:.+]] = spirv.Constant 1 : i32
// CHECK: %[[S6:.+]] = spirv.IAdd %[[S2]], %[[S3]] : i32
-// CHECK: %[[S7:.+]] = spirv.AccessChain %[[S0]][%[[CST0_1]], %[[S6]]] : !spirv.ptr<!spirv.struct<(!spirv.array<16 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
+// CHECK: %[[S7:.+]] = spirv.InBoundsAccessChain %[[S0]][%[[CST0_1]], %[[S6]]] : !spirv.ptr<!spirv.struct<(!spirv.array<16 x f32, stride=4> [0])>, StorageBuffer>, i32, i32
// CHECK: %[[S8:.+]] = spirv.Bitcast %[[S7]] : !spirv.ptr<f32, StorageBuffer> to !spirv.ptr<vector<4xf32>, StorageBuffer>
// CHECK: spirv.Store "StorageBuffer" %[[S8]], %[[ARG1]] : vector<4xf32>
func.func @vector_store_2d(%arg0 : memref<4x4xf32, #spirv.storage_class<StorageBuffer>>, %arg1 : vector<4xf32>) {
diff --git a/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir b/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
index a3b96c698a344..c58f22cea5a4b 100644
--- a/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
@@ -35,6 +35,17 @@ func.func @access_chain_2D_array_2(%arg0 : i32) -> () {
return
}
+//===----------------------------------------------------------------------===//
+// spirv.InBoundsAccessChain
+//===----------------------------------------------------------------------===//
+
+func.func @inbounds_access_chain(%arg0 : i32) -> () {
+ %0 = spirv.Variable : !spirv.ptr<!spirv.array<4xf32>, Function>
+ // CHECK: spirv.InBoundsAccessChain {{.*}}[{{.*}}] : !spirv.ptr<!spirv.array<4 x f32>, Function>
+ %1 = spirv.InBoundsAccessChain %0[%arg0] : !spirv.ptr<!spirv.array<4xf32>, Function>, i32 -> !spirv.ptr<f32, Function>
+ return
+}
+
func.func @access_chain_rtarray(%arg0 : i32) -> () {
%0 = spirv.Variable : !spirv.ptr<!spirv.rtarray<f32>, Function>
// CHECK: spirv.AccessChain {{.*}}[{{.*}}] : !spirv.ptr<!spirv.rtarray<f32>, Function>
diff --git a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
index 94d9c53db0bbc..e80145419f8a8 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
@@ -19,6 +19,25 @@ func.func @combine_full_access_chain() -> f32 {
// -----
+//===----------------------------------------------------------------------===//
+// spirv.InBoundsAccessChain
+//===----------------------------------------------------------------------===//
+
+func.func @combine_full_inbounds_access_chain() -> f32 {
+ // CHECK: %[[INDEX:.*]] = spirv.Constant 0
+ // CHECK-NEXT: %[[VAR:.*]] = spirv.Variable
+ // CHECK-NEXT: %[[PTR:.*]] = spirv.InBoundsAccessChain %[[VAR]][%[[INDEX]], %[[INDEX]], %[[INDEX]]]
+ // CHECK-NEXT: spirv.Load "Function" %[[PTR]]
+ %c0 = spirv.Constant 0: i32
+ %0 = spirv.Variable : !spirv.ptr<!spirv.struct<(!spirv.array<4x!spirv.array<4xf32>>, !spirv.array<4xi32>)>, Function>
+ %1 = spirv.InBoundsAccessChain %0[%c0] : !spirv.ptr<!spirv.struct<(!spirv.array<4x!spirv.array<4xf32>>, !spirv.array<4xi32>)>, Function>, i32 -> !spirv.ptr<!spirv.array<4x!spirv.array<4xf32>>, Function>
+ %2 = spirv.InBoundsAccessChain %1[%c0, %c0] : !spirv.ptr<!spirv.array<4x!spirv.array<4xf32>>, Function>, i32, i32 -> !spirv.ptr<f32, Function>
+ %3 = spirv.Load "Function" %2 : f32
+ spirv.ReturnValue %3 : f32
+}
+
+// -----
+
func.func @combine_access_chain_multi_use() -> !spirv.array<4xf32> {
// CHECK: %[[INDEX:.*]] = spirv.Constant 0
// CHECK-NEXT: %[[VAR:.*]] = spirv.Variable
diff --git a/mlir/test/Target/SPIRV/memory-ops.mlir b/mlir/test/Target/SPIRV/memory-ops.mlir
index 2d18394818611..fa786f18aac18 100644
--- a/mlir/test/Target/SPIRV/memory-ops.mlir
+++ b/mlir/test/Target/SPIRV/memory-ops.mlir
@@ -40,6 +40,17 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
// -----
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+ spirv.func @inbounds_access_chain(%arg0 : !spirv.ptr<!spirv.array<4xf32>, Function>, %arg1 : i32) "None" {
+ // CHECK: {{%.*}} = spirv.InBoundsAccessChain {{%.*}}[{{%.*}}] : !spirv.ptr<!spirv.array<4 x f32>, Function>
+ %0 = spirv.InBoundsAccessChain %arg0[%arg1] : !spirv.ptr<!spirv.array<4xf32>, Function>, i32 -> !spirv.ptr<f32, Function>
+ %1 = spirv.Load "Function" %0 : f32
+ spirv.Return
+ }
+}
+
+// -----
+
spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], [SPV_KHR_storage_buffer_storage_class]> {
spirv.func @load_store_zero_rank_float(%arg0: !spirv.ptr<!spirv.struct<(!spirv.array<1 x f32, stride=4> [0]), Block>, StorageBuffer>, %arg1: !spirv.ptr<!spirv.struct<(!spirv.array<1 x f32, stride=4> [0]), Block>, StorageBuffer>) "None" {
// CHECK: [[LOAD_PTR:%.*]] = spirv.AccessChain {{%.*}}[{{%.*}}, {{%.*}}] : !spirv.ptr<!spirv.struct<(!spirv.array<1 x f32, stride=4> [0]), Block>, StorageBuffer>
@@ -122,4 +133,3 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.4, [Shader, Linkage], []> {
spirv.Return
}
}
-
More information about the Mlir-commits
mailing list