[Mlir-commits] [mlir] [mlir][SPIR-V] Derive Aligned memory operand for PhysicalStorageBuffer in vector.load/store (PR #214200)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 5 05:06:01 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-spirv
Author: Arseniy Obolenskiy (aobolensk)
<details>
<summary>Changes</summary>
Mirror calculateMemoryRequirements from MemRefToSPIRV.cpp, since PhysicalStorageBuffer pointers carry no implicit alignment guarantee
---
Full diff: https://github.com/llvm/llvm-project/pull/214200.diff
2 Files Affected:
- (modified) mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp (+62-37)
- (modified) mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir (+30)
``````````diff
diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 78693e924c4d9..22b3619b27e8c 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
@@ -53,6 +53,50 @@ static int getNumBits(Type type) {
return type.getIntOrFloatBitWidth();
}
+namespace {
+struct MemoryRequirements {
+ spirv::MemoryAccessAttr memoryAccess;
+ IntegerAttr alignment;
+};
+} // namespace
+
+/// PhysicalStorageBuffer pointers always need Aligned, derived from pointee
+/// size unless preferredAlignment is set; other storage classes may omit it.
+static FailureOr<MemoryRequirements>
+calculateMemoryRequirements(Value accessedPtr, uint64_t preferredAlignment) {
+ if (preferredAlignment >= std::numeric_limits<uint32_t>::max())
+ return failure();
+
+ MLIRContext *ctx = accessedPtr.getContext();
+ auto ptrType = cast<spirv::PointerType>(accessedPtr.getType());
+ bool mayOmitAlignment =
+ !preferredAlignment &&
+ ptrType.getStorageClass() != spirv::StorageClass::PhysicalStorageBuffer;
+ if (mayOmitAlignment)
+ return MemoryRequirements{spirv::MemoryAccessAttr{}, IntegerAttr{}};
+
+ // PhysicalStorageBuffer pointers require Aligned.
+ std::optional<int64_t> sizeInBytes;
+ Type pointeeType = ptrType.getPointeeType();
+ if (auto scalarType = dyn_cast<spirv::ScalarType>(pointeeType)) {
+ sizeInBytes = scalarType.getSizeInBytes();
+ } else if (auto vecType = dyn_cast<VectorType>(pointeeType)) {
+ if (auto scalarElem = dyn_cast<spirv::ScalarType>(vecType.getElementType()))
+ if (auto elemSize = scalarElem.getSizeInBytes())
+ sizeInBytes = *elemSize * vecType.getNumElements();
+ }
+
+ if (!sizeInBytes)
+ return failure();
+
+ auto memoryAccess =
+ spirv::MemoryAccessAttr::get(ctx, spirv::MemoryAccess::Aligned);
+ uint64_t alignmentValue =
+ preferredAlignment ? preferredAlignment : *sizeInBytes;
+ auto alignment = IntegerAttr::get(IntegerType::get(ctx, 32), alignmentValue);
+ return MemoryRequirements{memoryAccess, alignment};
+}
+
namespace {
struct VectorShapeCast final : public OpConversionPattern<vector::ShapeCastOp> {
@@ -760,22 +804,6 @@ struct VectorLoadOpConverter final
auto vectorPtrType = spirv::PointerType::get(spirvVectorType, storageClass);
- std::optional<uint64_t> alignment = loadOp.getAlignment();
- if (alignment > std::numeric_limits<uint32_t>::max()) {
- return rewriter.notifyMatchFailure(loadOp,
- "invalid alignment requirement");
- }
-
- auto memoryAccess = spirv::MemoryAccess::None;
- spirv::MemoryAccessAttr memoryAccessAttr;
- IntegerAttr alignmentAttr;
- if (alignment.has_value()) {
- memoryAccess |= spirv::MemoryAccess::Aligned;
- memoryAccessAttr =
- spirv::MemoryAccessAttr::get(rewriter.getContext(), memoryAccess);
- alignmentAttr = rewriter.getI32IntegerAttr(alignment.value());
- }
-
// For single element vectors, we don't need to bitcast the access chain to
// the original vector type. Both is going to be the same, a pointer
// to a scalar.
@@ -785,9 +813,15 @@ struct VectorLoadOpConverter final
: spirv::BitcastOp::create(rewriter, loc, vectorPtrType,
accessChain);
- rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp, spirvVectorType,
- castedAccessChain,
- memoryAccessAttr, alignmentAttr);
+ auto memoryRequirements = calculateMemoryRequirements(
+ castedAccessChain, loadOp.getAlignment().value_or(0));
+ if (failed(memoryRequirements))
+ return rewriter.notifyMatchFailure(
+ loadOp, "failed to determine memory requirements");
+
+ auto [memoryAccess, alignment] = *memoryRequirements;
+ rewriter.replaceOpWithNewOp<spirv::LoadOp>(
+ loadOp, spirvVectorType, castedAccessChain, memoryAccess, alignment);
return success();
}
@@ -816,12 +850,6 @@ struct VectorStoreOpConverter final
return rewriter.notifyMatchFailure(
storeOp, "failed to get memref element pointer");
- std::optional<uint64_t> alignment = storeOp.getAlignment();
- if (alignment > std::numeric_limits<uint32_t>::max()) {
- return rewriter.notifyMatchFailure(storeOp,
- "invalid alignment requirement");
- }
-
spirv::StorageClass storageClass = attr.getValue();
auto vectorType = storeOp.getVectorType();
// Use the converted vector type instead of original (single element vector
@@ -841,19 +869,16 @@ struct VectorStoreOpConverter final
: spirv::BitcastOp::create(rewriter, loc, vectorPtrType,
accessChain);
- auto memoryAccess = spirv::MemoryAccess::None;
- spirv::MemoryAccessAttr memoryAccessAttr;
- IntegerAttr alignmentAttr;
- if (alignment.has_value()) {
- memoryAccess |= spirv::MemoryAccess::Aligned;
- memoryAccessAttr =
- spirv::MemoryAccessAttr::get(rewriter.getContext(), memoryAccess);
- alignmentAttr = rewriter.getI32IntegerAttr(alignment.value());
- }
+ auto memoryRequirements = calculateMemoryRequirements(
+ castedAccessChain, storeOp.getAlignment().value_or(0));
+ if (failed(memoryRequirements))
+ return rewriter.notifyMatchFailure(
+ storeOp, "failed to determine memory requirements");
- rewriter.replaceOpWithNewOp<spirv::StoreOp>(
- storeOp, castedAccessChain, adaptor.getValueToStore(), memoryAccessAttr,
- alignmentAttr);
+ auto [memoryAccess, alignment] = *memoryRequirements;
+ rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, castedAccessChain,
+ adaptor.getValueToStore(),
+ memoryAccess, alignment);
return success();
}
diff --git a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
index f904dd9d35c37..2beb59d16b631 100644
--- a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
+++ b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
@@ -1276,3 +1276,33 @@ func.func @vector_load(%arg0 : memref<4xf32, #spirv.storage_class<StorageBuffer>
%0 = vector.load %arg0[%idx] : memref<4xf32, #spirv.storage_class<StorageBuffer>>, vector<4xf32>
return %0: vector<4xf32>
}
+
+// -----
+
+// PhysicalStorageBuffer pointers get an Aligned memory operand, unlike StorageBuffer above.
+
+module attributes {
+ spirv.target_env = #spirv.target_env<
+ #spirv.vce<v1.6, [Shader, PhysicalStorageBufferAddresses],
+ [SPV_KHR_physical_storage_buffer]>, #spirv.resource_limits<>>
+ } {
+
+// CHECK-LABEL: @vector_load_physical
+// CHECK: spirv.Load "PhysicalStorageBuffer" {{%.*}} ["Aligned", 16] : vector<4xf32>
+func.func @vector_load_physical(%arg0 : memref<4xf32, #spirv.storage_class<PhysicalStorageBuffer>>
+ {spirv.decoration = #spirv.decoration<Aliased>}) -> vector<4xf32> {
+ %idx = arith.constant 0 : index
+ %0 = vector.load %arg0[%idx] : memref<4xf32, #spirv.storage_class<PhysicalStorageBuffer>>, vector<4xf32>
+ return %0: vector<4xf32>
+}
+
+// CHECK-LABEL: @vector_store_physical
+// CHECK: spirv.Store "PhysicalStorageBuffer" {{%.*}}, {{%.*}} ["Aligned", 16] : vector<4xf32>
+func.func @vector_store_physical(%arg0 : memref<4xf32, #spirv.storage_class<PhysicalStorageBuffer>>
+ {spirv.decoration = #spirv.decoration<Aliased>}, %arg1 : vector<4xf32>) {
+ %idx = arith.constant 0 : index
+ vector.store %arg1, %arg0[%idx] : memref<4xf32, #spirv.storage_class<PhysicalStorageBuffer>>, vector<4xf32>
+ return
+}
+
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/214200
More information about the Mlir-commits
mailing list