[Mlir-commits] [mlir] [mlir][spirv] Mark in-bounds linearized indices no-wrap (PR #215834)
Hsiangkai Wang
llvmlistbot at llvm.org
Mon Aug 17 03:57:56 PDT 2026
https://github.com/Hsiangkai updated https://github.com/llvm/llvm-project/pull/215834
>From 4670ba51d9dfc0073b3daf79b3dc6d8e21f4693e 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 | 88 +++++++++++++++++--
.../MemRefToSPIRV/memref-to-spirv.mlir | 50 +++++++++++
4 files changed, 156 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..1e2bfc5557796 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,48 @@ 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 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 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 +1380,32 @@ 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)
+ return flags;
+
+ std::optional<uint64_t> maxLinearIndex =
+ getMaxLinearizedIndex(shape, strides, offset);
+ 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 +1422,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 +1450,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 +1459,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 +1486,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 ff29fec9cd20931c81c0e44ddd8b023ab296865a Mon Sep 17 00:00:00 2001
From: Hsiangkai Wang <hsiangkai.wang at arm.com>
Date: Mon, 17 Aug 2026 11:53:40 +0100
Subject: [PATCH 2/2] [mlir][spirv] Test no-wrap decoration availability
Cover both SPIR-V 1.4, where SPV_KHR_no_integer_wrap_decoration is implied,
and pre-1.4 targets without the extension.
---
.../MemRefToSPIRV/memref-to-spirv.mlir | 44 +++++++++++++++++++
1 file changed, 44 insertions(+)
diff --git a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
index 5550c778b6f31..ebddeadf3a31a 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -270,6 +270,50 @@ func.func @dynamic_linearized_index(
// -----
+module attributes {
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @static_linearized_index_with_spirv_1_4
+func.func @static_linearized_index_with_spirv_1_4(
+ %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
+}
+
+}
+
+// -----
+
+module attributes {
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.1, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @static_linearized_index_without_no_wrap_extension
+func.func @static_linearized_index_without_no_wrap_extension(
+ %arg0: memref<2x4xf32, #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: %[[LINEAR:.+]] = spirv.IAdd %{{.*}}, %[[OFFSET]] : i32
+ // CHECK-NOT: no_signed_wrap
+ // CHECK-NOT: no_unsigned_wrap
+ // CHECK: spirv.{{.*}}AccessChain {{.*}}[%{{.*}}, %[[LINEAR]]]
+ %0 = memref.load %arg0[%row, %column] : memref<2x4xf32, #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.
More information about the Mlir-commits
mailing list