[Mlir-commits] [mlir] ddda690 - [mlir][spirv] Mark in-bounds linearized indices no-wrap (#215834)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Aug 18 01:01:37 PDT 2026
Author: Hsiangkai Wang
Date: 2026-08-18T09:01:32+01:00
New Revision: ddda6903e9d10cd8b835edd5be5e5c7d3e9c6c83
URL: https://github.com/llvm/llvm-project/commit/ddda6903e9d10cd8b835edd5be5e5c7d3e9c6c83
DIFF: https://github.com/llvm/llvm-project/commit/ddda6903e9d10cd8b835edd5be5e5c7d3e9c6c83.diff
LOG: [mlir][spirv] Mark in-bounds linearized indices no-wrap (#215834)
Preserve signed and unsigned no-wrap guarantees for statically bounded
SPIR-V index linearization, allowing buffer address arithmetic to be
reassociated safely.
Added:
Modified:
mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
index 03ae54a8ae30a..577fa2c44496b 100644
--- a/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
+++ b/mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
@@ -168,11 +168,25 @@ 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`, if supported by `targetEnv`.
+LinearizedIndexNoWrapFlags getLinearizedIndexNoWrapFlags(
+ const TargetEnv &targetEnv, 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..236d0ad024e42 100644
--- a/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
+++ b/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
@@ -81,8 +81,13 @@ class TensorExtractPattern final
auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
auto indexType = typeConverter.getIndexType();
+ spirv::LinearizedIndexNoWrapFlags noWrapFlags =
+ spirv::getLinearizedIndexNoWrapFlags(typeConverter.getTargetEnv(),
+ 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..1c8a22fb35639 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,37 @@ 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;
+}
+
} // namespace
//===----------------------------------------------------------------------===//
@@ -1336,9 +1369,36 @@ Value spirv::getPushConstantValue(Operation *op, unsigned elementCount,
// Public functions for index calculation
//===----------------------------------------------------------------------===//
+mlir::spirv::LinearizedIndexNoWrapFlags
+mlir::spirv::getLinearizedIndexNoWrapFlags(const TargetEnv &targetEnv,
+ ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> strides,
+ int64_t offset, Type integerType) {
+ LinearizedIndexNoWrapFlags flags;
+ if (!targetEnv.allows(Extension::SPV_KHR_no_integer_wrap_decoration))
+ return 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 +1415,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 +1443,9 @@ Value mlir::spirv::getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
}
auto indexType = typeConverter.getIndexType();
+ LinearizedIndexNoWrapFlags noWrapFlags = getLinearizedIndexNoWrapFlags(
+ typeConverter.getTargetEnv(), baseType.getShape(), strides, offset,
+ indexType);
SmallVector<Value, 2> linearizedIndices;
auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
@@ -1383,8 +1453,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 +1480,17 @@ Value mlir::spirv::getOpenCLElementPtr(const SPIRVTypeConverter &typeConverter,
}
auto indexType = typeConverter.getIndexType();
+ LinearizedIndexNoWrapFlags noWrapFlags = getLinearizedIndexNoWrapFlags(
+ typeConverter.getTargetEnv(), baseType.getShape(), 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..ebddeadf3a31a 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -220,6 +220,100 @@ 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
+}
+
+}
+
+// -----
+
+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