[mlir][spirv] Mark in-bounds linearized indices no-wrap - #215834
Open
Hsiangkai wants to merge 1 commit into
Open
[mlir][spirv] Mark in-bounds linearized indices no-wrap#215834Hsiangkai wants to merge 1 commit into
Hsiangkai wants to merge 1 commit into
Conversation
Preserve signed and unsigned no-wrap guarantees for statically bounded SPIR-V index linearization, allowing buffer address arithmetic to be reassociated safely.
Hsiangkai
requested review from
inbelic
and removed request for
IgWod,
antiagainst and
kuhar
August 12, 2026 16:02
|
@llvm/pr-subscribers-mlir-spirv @llvm/pr-subscribers-mlir Author: Hsiangkai Wang (Hsiangkai) ChangesPreserve signed and unsigned no-wrap guarantees for statically bounded SPIR-V index linearization, allowing buffer address arithmetic to be reassociated safely. Full diff: https://github.com/llvm/llvm-project/pull/215834.diff 4 Files Affected:
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.
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Preserve signed and unsigned no-wrap guarantees for statically bounded SPIR-V index linearization, allowing buffer address arithmetic to be reassociated safely.