[Mlir-commits] [mlir] [mlir][SPIRV] Handle vector stores when unifying aliased resources (PR #232185)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 14 23:03:00 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-spirv
Author: Anutosh Bhat (anutosh491)
<details>
<summary>Changes</summary>
I hit this while compiling JAX training code to WebGPU through IREE.
The alias-unification pass can select a smaller vector resource, but its store conversion currently only handles scalars. Handle vector stores by splitting them into stores of the canonical vector type, preserving memory-access flags and alignment.
Assisted by codex (GPT 6.1 Sol)
---
Full diff: https://github.com/llvm/llvm-project/pull/232185.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp (+61-6)
- (modified) mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir (+120)
``````````diff
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
index 99eaa3f0b16654..2cc1f0a4a821a8 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
@@ -24,6 +24,7 @@
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Alignment.h"
#include <iterator>
namespace mlir {
@@ -550,15 +551,69 @@ struct ConvertStore : public ConvertAliasResource<spirv::StoreOp> {
cast<spirv::PointerType>(storeOp.getPtr().getType()).getPointeeType();
auto dstElemType =
cast<spirv::PointerType>(adaptor.getPtr().getType()).getPointeeType();
- if (!srcElemType.isIntOrFloat() || !dstElemType.isIntOrFloat())
- return rewriter.notifyMatchFailure(storeOp, "not scalar type");
- if (!areSameBitwidthScalarType(srcElemType, dstElemType))
- return rewriter.notifyMatchFailure(storeOp, "different bitwidth");
-
Location loc = storeOp.getLoc();
Value value = adaptor.getValue();
- if (srcElemType != dstElemType)
+ auto srcVector = dyn_cast<VectorType>(srcElemType);
+ auto dstVector = dyn_cast<VectorType>(dstElemType);
+ if (srcVector && dstVector && srcVector != dstVector) {
+ if (srcVector.getElementTypeBitWidth() !=
+ dstVector.getElementTypeBitWidth() ||
+ srcVector.getNumElements() % dstVector.getNumElements() != 0)
+ return rewriter.notifyMatchFailure(storeOp,
+ "incompatible vector types");
+ auto acOp = adaptor.getPtr().getDefiningOp<spirv::AccessChainOp>();
+ if (!acOp)
+ return rewriter.notifyMatchFailure(storeOp,
+ "ptr not spirv.AccessChain");
+ if (srcVector.getElementType() != dstVector.getElementType())
+ value = spirv::BitcastOp::create(
+ rewriter, loc,
+ VectorType::get(srcVector.getShape(), dstVector.getElementType()),
+ value);
+ unsigned count = dstVector.getNumElements();
+ unsigned ratio = srcVector.getNumElements() / count;
+ unsigned bytes = count * dstVector.getElementTypeBitWidth() / 8;
+ for (unsigned i = 0; i < ratio; ++i) {
+ Value part = value;
+ if (ratio > 1) {
+ SmallVector<int32_t> components;
+ for (unsigned j = 0; j < count; ++j)
+ components.push_back(i * count + j);
+ part = spirv::VectorShuffleOp::create(
+ rewriter, loc, dstVector, value, value,
+ rewriter.getI32ArrayAttr(components));
+ }
+ Value ptr = adaptor.getPtr();
+ auto alignment = storeOp.getAlignmentAttr();
+ if (i) {
+ auto indices = llvm::to_vector(acOp.getIndices());
+ Type indexType = indices.back().getType();
+ auto offset = spirv::ConstantOp::create(
+ rewriter, loc, indexType, rewriter.getIntegerAttr(indexType, i));
+ indices.back() = spirv::IAddOp::create(rewriter, loc, indexType,
+ indices.back(), offset);
+ ptr = spirv::AccessChainOp::create(rewriter, loc, acOp.getBasePtr(),
+ indices);
+ if (alignment)
+ alignment = rewriter.getI32IntegerAttr(
+ llvm::commonAlignment(llvm::Align(alignment.getInt()),
+ i * bytes)
+ .value());
+ }
+ auto store = spirv::StoreOp::create(
+ rewriter, loc, ptr, part, storeOp.getMemoryAccessAttr(), alignment);
+ store->setDiscardableAttrs(storeOp->getDiscardableAttrDictionary());
+ }
+ rewriter.eraseOp(storeOp);
+ return success();
+ }
+ if (srcElemType != dstElemType) {
+ if (!srcElemType.isIntOrFloat() || !dstElemType.isIntOrFloat())
+ return rewriter.notifyMatchFailure(storeOp, "not scalar type");
+ if (!areSameBitwidthScalarType(srcElemType, dstElemType))
+ return rewriter.notifyMatchFailure(storeOp, "different bitwidth");
value = spirv::BitcastOp::create(rewriter, loc, dstElemType, value);
+ }
auto newStore = rewriter.replaceOpWithNewOp<spirv::StoreOp>(
storeOp, adaptor.getPtr(), value, storeOp.getMemoryAccessAttr(),
storeOp.getAlignmentAttr());
diff --git a/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir b/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir
index d7cc1d0cc1f849..e0567272e57de1 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir
@@ -581,3 +581,123 @@ spirv.module Logical GLSL450 {
// CHECK-LABEL: spirv.func @main
// CHECK-SAME: (%{{.+}}: !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>) "None"
// CHECK: spirv.Return
+
+// -----
+
+spirv.module Logical GLSL450 {
+ spirv.GlobalVariable @canonical bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf32>, stride=8> [0])>, StorageBuffer>
+ spirv.GlobalVariable @wide bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ spirv.func @store_wide_vector(%value: vector<4xf32>) "None" {
+ %c0 = spirv.Constant 0 : i32
+ %c3 = spirv.Constant 3 : i32
+ %base = spirv.mlir.addressof @wide : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ %ptr = spirv.AccessChain %base[%c0, %c3] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<vector<4xf32>, StorageBuffer>
+ spirv.Store "StorageBuffer" %ptr, %value : vector<4xf32>
+ spirv.Return
+ }
+}
+// CHECK-LABEL: spirv.module
+// CHECK: spirv.GlobalVariable @canonical bind(0, 0)
+// CHECK-NOT: spirv.GlobalVariable @wide
+// CHECK: spirv.func @store_wide_vector(%[[VALUE:.+]]: vector<4xf32>)
+// CHECK-DAG: %[[ZERO:.+]] = spirv.Constant 0 : i32
+// CHECK-DAG: %[[THREE:.+]] = spirv.Constant 3 : i32
+// CHECK-DAG: %[[BASE:.+]] = spirv.mlir.addressof @canonical
+// CHECK-DAG: %[[TWO:.+]] = spirv.Constant 2 : i32
+// CHECK: %[[INDEX:.+]] = spirv.IMul %[[THREE]], %[[TWO]] : i32
+// CHECK: %[[PTR0:.+]] = spirv.AccessChain %[[BASE]][%[[ZERO]], %[[INDEX]]]
+// CHECK: %[[LOW:.+]] = spirv.VectorShuffle [0 : i32, 1 : i32] %[[VALUE]], %[[VALUE]]
+// CHECK: spirv.Store "StorageBuffer" %[[PTR0]], %[[LOW]] : vector<2xf32>
+// CHECK: %[[HIGH:.+]] = spirv.VectorShuffle [2 : i32, 3 : i32] %[[VALUE]], %[[VALUE]]
+// CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
+// CHECK: %[[NEXT:.+]] = spirv.IAdd %[[INDEX]], %[[ONE]] : i32
+// CHECK: %[[PTR1:.+]] = spirv.AccessChain %[[BASE]][%[[ZERO]], %[[NEXT]]]
+// CHECK: spirv.Store "StorageBuffer" %[[PTR1]], %[[HIGH]] : vector<2xf32>
+
+// -----
+
+spirv.module Logical GLSL450 {
+ spirv.GlobalVariable @canonical bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf32>, stride=8> [0])>, StorageBuffer>
+ spirv.GlobalVariable @wide bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xi32>, stride=16> [0])>, StorageBuffer>
+ spirv.func @store_wide_bitcast(%value: vector<4xi32>) "None" {
+ %c0 = spirv.Constant 0 : i32
+ %c3 = spirv.Constant 3 : i32
+ %base = spirv.mlir.addressof @wide : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xi32>, stride=16> [0])>, StorageBuffer>
+ %ptr = spirv.AccessChain %base[%c0, %c3] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xi32>, stride=16> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<vector<4xi32>, StorageBuffer>
+ spirv.Store "StorageBuffer" %ptr, %value : vector<4xi32>
+ spirv.Return
+ }
+}
+// CHECK-LABEL: spirv.module
+// CHECK: spirv.GlobalVariable @canonical bind(0, 0)
+// CHECK-NOT: spirv.GlobalVariable @wide
+// CHECK: spirv.func @store_wide_bitcast
+// CHECK: spirv.Bitcast
+// CHECK: spirv.VectorShuffle
+// CHECK: spirv.Store "StorageBuffer" {{.+}} : vector<2xf32>
+// CHECK: spirv.VectorShuffle
+// CHECK: spirv.Store "StorageBuffer" {{.+}} : vector<2xf32>
+
+// -----
+
+spirv.module Logical GLSL450 {
+ spirv.GlobalVariable @canonical bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf32>, stride=8> [0])>, StorageBuffer>
+ spirv.GlobalVariable @wide bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ spirv.func @store_wide_aligned(%value: vector<4xf32>) "None" {
+ %c0 = spirv.Constant 0 : i32
+ %c3 = spirv.Constant 3 : i32
+ %base = spirv.mlir.addressof @wide : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ %ptr = spirv.AccessChain %base[%c0, %c3] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<vector<4xf32>, StorageBuffer>
+ spirv.Store "StorageBuffer" %ptr, %value ["Volatile|Aligned", 16] : vector<4xf32>
+ spirv.Return
+ }
+}
+// CHECK-LABEL: spirv.module
+// CHECK: spirv.GlobalVariable @canonical bind(0, 0)
+// CHECK-NOT: spirv.GlobalVariable @wide
+// CHECK: spirv.func @store_wide_aligned
+// CHECK: spirv.VectorShuffle
+// CHECK: spirv.Store "StorageBuffer" {{.+}}["Volatile|Aligned", 16] : vector<2xf32>
+// CHECK: spirv.VectorShuffle
+// CHECK: spirv.Store "StorageBuffer" {{.+}}["Volatile|Aligned", 8] : vector<2xf32>
+
+// -----
+
+spirv.module Logical GLSL450 {
+ spirv.GlobalVariable @canonical bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ spirv.GlobalVariable @wide bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xi32>, stride=16> [0])>, StorageBuffer>
+ spirv.func @store_equal_vector_bitcast(%value: vector<4xi32>) "None" {
+ %c0 = spirv.Constant 0 : i32
+ %c3 = spirv.Constant 3 : i32
+ %base = spirv.mlir.addressof @wide : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xi32>, stride=16> [0])>, StorageBuffer>
+ %ptr = spirv.AccessChain %base[%c0, %c3] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xi32>, stride=16> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<vector<4xi32>, StorageBuffer>
+ spirv.Store "StorageBuffer" %ptr, %value : vector<4xi32>
+ spirv.Return
+ }
+}
+// CHECK-LABEL: spirv.module
+// CHECK: spirv.GlobalVariable @canonical bind(0, 0)
+// CHECK-NOT: spirv.GlobalVariable @wide
+// CHECK: spirv.func @store_equal_vector_bitcast
+// CHECK: spirv.Bitcast
+// CHECK: spirv.Store "StorageBuffer" {{.+}} : vector<4xf32>
+
+// -----
+
+spirv.module Logical GLSL450 {
+ spirv.GlobalVariable @canonical bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ spirv.GlobalVariable @alias bind(0, 0) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ spirv.func @store_equal_vector(%value: vector<4xf32>) "None" {
+ %c0 = spirv.Constant 0 : i32
+ %base = spirv.mlir.addressof @alias : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+ %ptr = spirv.AccessChain %base[%c0, %c0] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<vector<4xf32>, StorageBuffer>
+ spirv.Store "StorageBuffer" %ptr, %value : vector<4xf32>
+ spirv.Return
+ }
+}
+// CHECK-LABEL: spirv.module
+// CHECK: spirv.GlobalVariable @canonical bind(0, 0)
+// CHECK-NOT: spirv.GlobalVariable @alias
+// CHECK: spirv.func @store_equal_vector(%[[VALUE:.+]]: vector<4xf32>)
+// CHECK: %[[PTR:.+]] = spirv.AccessChain
+// CHECK: spirv.Store "StorageBuffer" %[[PTR]], %[[VALUE]] : vector<4xf32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/232185
More information about the Mlir-commits
mailing list