[Mlir-commits] [mlir] [mlir][gpu][spirv] Convert memref<mma_matrix> to spv.array<coopmatrix> (PR #212806)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 29 08:52:32 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-spirv
Author: fabrizio-indirli
<details>
<summary>Changes</summary>
In a gpu.func, allow allocating local arrays of gpu.mma_matrix, e.g. `memref.alloca() : memref<Nx!gpu.mma_matrix<HxW, ...>`. Extend the SPIRV conversion to convert such memrefs to `%ARR = spirv.Variable: spv.ptr<spv.array<N x coopmatrix>, Function>`, which can be conveniently indexed with `spirv.AccessChain %ARR[i]`. This enables arrays of CoopMMA matrices both at the `gpu` and `spirv` levels.
---
Full diff: https://github.com/llvm/llvm-project/pull/212806.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h (+5-1)
- (modified) mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp (+38-13)
- (modified) mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp (+28-8)
- (modified) mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp (+6-3)
- (modified) mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir (+53)
- (modified) mlir/test/Conversion/MemRefToSPIRV/alloca.mlir (+19)
``````````diff
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
index 0e135781d9d92..5880c78136566 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
@@ -134,9 +134,13 @@ struct MMAMatrixStorageType : public TypeStorage {
///
/// gpu.subgroup_mma_store_matrix %3, %arg22[%c0, %c0] {leadDimension = 16
/// : index}: !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32>
+///
+/// MMA matrices may be memref elements on targets that support storing them,
+/// for example SPIRV can represent arrays of MMA matrices.
// TODO: consider moving this to ODS.
class MMAMatrixType
- : public Type::TypeBase<MMAMatrixType, Type, MMAMatrixStorageType> {
+ : public Type::TypeBase<MMAMatrixType, Type, MMAMatrixStorageType,
+ MemRefElementTypeInterface::Trait> {
public:
using Base::Base;
diff --git a/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
index 956134dfee55d..f3d57cf317c2a 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
@@ -26,12 +26,48 @@
#include "llvm/ADT/StringSwitch.h"
#include <cassert>
+#include <limits>
namespace mlir {
//===----------------------------------------------------------------------===//
// Patterns and helpers.
//===----------------------------------------------------------------------===//
+static spirv::CooperativeMatrixType
+convertMMAMatrixType(gpu::MMAMatrixType type) {
+ ArrayRef<int64_t> shape = type.getShape();
+ auto use =
+ llvm::StringSwitch<spirv::CooperativeMatrixUseKHR>(type.getOperand())
+ .Case("AOp", spirv::CooperativeMatrixUseKHR::MatrixA)
+ .Case("BOp", spirv::CooperativeMatrixUseKHR::MatrixB)
+ .Default(spirv::CooperativeMatrixUseKHR::MatrixAcc);
+ return spirv::CooperativeMatrixType::get(
+ type.getElementType(), shape[0], shape[1], spirv::Scope::Subgroup, use);
+}
+
+// Convert a memref of `gpu.mma_matrix` into a SPIR-V pointer to an array
+// of spirv::CooperativeMatrix.
+static std::optional<Type> convertMemrefOfMMAMatrixType(MemRefType type) {
+ auto matrixType = dyn_cast<gpu::MMAMatrixType>(type.getElementType());
+ if (!matrixType)
+ return std::nullopt;
+ // SPV Cooperative matrix types, and composites containing them, may only be
+ // allocated in Function or Private storage classes. For now, support only
+ // function-local arrays.
+ auto storageClass =
+ dyn_cast_or_null<spirv::StorageClassAttr>(type.getMemorySpace());
+ if (!storageClass ||
+ storageClass.getValue() != spirv::StorageClass::Function ||
+ !type.hasStaticShape() || !type.getLayout().isIdentity() ||
+ type.getNumElements() <= 0 ||
+ type.getNumElements() > std::numeric_limits<unsigned>::max())
+ return Type();
+ auto arrayType =
+ spirv::ArrayType::get(convertMMAMatrixType(matrixType),
+ static_cast<unsigned>(type.getNumElements()));
+ return spirv::PointerType::get(arrayType, spirv::StorageClass::Function);
+}
+
/// Creates a SPIR-V op to replace the given GPU subgroup mma elementwise op
/// when the elementwise op directly supports with cooperative matrix type.
/// Returns false if cannot.
@@ -413,17 +449,6 @@ void mlir::populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(
void mlir::populateMMAToSPIRVCoopMatrixTypeConversion(
mlir::SPIRVTypeConverter &typeConverter) {
- typeConverter.addConversion([](gpu::MMAMatrixType type) {
- ArrayRef<int64_t> retTypeShape = type.getShape();
- Type elementType = type.getElementType();
- auto use =
- llvm::StringSwitch<spirv::CooperativeMatrixUseKHR>(type.getOperand())
- .Case("AOp", spirv::CooperativeMatrixUseKHR::MatrixA)
- .Case("BOp", spirv::CooperativeMatrixUseKHR::MatrixB)
- .Default(spirv::CooperativeMatrixUseKHR::MatrixAcc);
-
- return spirv::CooperativeMatrixType::get(elementType, retTypeShape[0],
- retTypeShape[1],
- spirv::Scope::Subgroup, use);
- });
+ typeConverter.addConversion(convertMMAMatrixType);
+ typeConverter.addConversion(convertMemrefOfMMAMatrixType);
}
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index 7eca3d65052a3..382e32f90ff9f 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
+++ b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
@@ -21,6 +21,7 @@
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Visitors.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
#include <cassert>
#include <limits>
#include <optional>
@@ -122,6 +123,10 @@ static Value shiftValue(Location loc, Value value, Value offset, Value mask,
/// Returns true if the allocations of memref `type` generated from `allocOp`
/// can be lowered to SPIR-V.
static bool isAllocationSupported(Operation *allocOp, MemRefType type) {
+ // Currently only support static shape
+ if (!type.hasStaticShape())
+ return false;
+
if (isa<memref::AllocOp, memref::DeallocOp>(allocOp)) {
auto sc = dyn_cast_or_null<spirv::StorageClassAttr>(type.getMemorySpace());
if (!sc || sc.getValue() != spirv::StorageClass::Workgroup)
@@ -130,15 +135,15 @@ static bool isAllocationSupported(Operation *allocOp, MemRefType type) {
auto sc = dyn_cast_or_null<spirv::StorageClassAttr>(type.getMemorySpace());
if (!sc || sc.getValue() != spirv::StorageClass::Function)
return false;
+ // Function allocations of memref-compatible element types may be lowered
+ // to SPIRV pointers/arrays of the corresponding SPIRV element type.
+ if (isa<MemRefElementTypeInterface>(type.getElementType()))
+ return true;
} else {
return false;
}
- // Currently only support static shape and int or float, complex of int or
- // float, or vector of int or float element type.
- if (!type.hasStaticShape())
- return false;
-
+ // Support memref of int or float types, or their vector/complex types.
Type elementType = type.getElementType();
if (auto vecType = dyn_cast<VectorType>(elementType))
elementType = vecType.getElementType();
@@ -410,9 +415,24 @@ AllocaOpPattern::matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
if (!spirvType)
return rewriter.notifyMatchFailure(allocaOp, "type conversion failed");
- rewriter.replaceOpWithNewOp<spirv::VariableOp>(allocaOp, spirvType,
- spirv::StorageClass::Function,
- /*initializer=*/nullptr);
+ auto function = allocaOp->getParentOfType<FunctionOpInterface>();
+ if (!function)
+ return rewriter.notifyMatchFailure(allocaOp,
+ "requires a containing function");
+
+ // SPIR-V requires Function variables to be declared in the first block.
+ OpBuilder::InsertionGuard guard(rewriter);
+ Block &entryBlock = function->getRegion(0).front();
+ auto insertionPoint = entryBlock.begin();
+ // Insert the variable after any existing ones to preserve ordering.
+ while (insertionPoint != entryBlock.end() &&
+ isa<spirv::VariableOp>(*insertionPoint))
+ ++insertionPoint;
+ rewriter.setInsertionPoint(&entryBlock, insertionPoint);
+ Value variable = spirv::VariableOp::create(
+ rewriter, allocaOp.getLoc(), spirvType, spirv::StorageClass::Function,
+ /*initializer=*/nullptr);
+ rewriter.replaceOp(allocaOp, variable);
return success();
}
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index 858b89728d432..ef4d79c827bb2 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -1380,15 +1380,18 @@ Value mlir::spirv::getVulkanElementPtr(const SPIRVTypeConverter &typeConverter,
SmallVector<Value, 2> linearizedIndices;
auto zero = spirv::ConstantOp::getZero(indexType, loc, builder);
- // Add a '0' at the start to index into the struct.
- linearizedIndices.push_back(zero);
-
if (baseType.getRank() == 0) {
linearizedIndices.push_back(zero);
} else {
linearizedIndices.push_back(
linearizeIndex(indices, strides, offset, indexType, loc, builder));
}
+
+ const Type pointeeType =
+ cast<spirv::PointerType>(basePtr.getType()).getPointeeType();
+ // Interface memrefs are wrapped in a struct: index to its first elem.
+ if (isa<spirv::StructType>(pointeeType))
+ linearizedIndices.insert(linearizedIndices.begin(), zero);
return spirv::AccessChainOp::create(builder, loc, basePtr, linearizedIndices);
}
diff --git a/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir b/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir
index 30ea121a54ead..c0115c4337f7f 100644
--- a/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir
@@ -228,5 +228,58 @@ module attributes {
// CHECK: spirv.Return
gpu.return
}
+
+ // CHECK-LABEL: spirv.func @gpu_wmma_function_array
+ gpu.func @gpu_wmma_function_array(%index: index, %scalar: f32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // CHECK: %[[ARRAY:.+]] = spirv.Variable : !spirv.ptr<!spirv.array<16 x !spirv.coopmatrix<4x4xf32, Subgroup, MatrixAcc>>, Function>
+ // CHECK: %[[STORE_PTR:.+]] = spirv.AccessChain %[[ARRAY]][%{{.+}}] : !spirv.ptr<!spirv.array<16 x !spirv.coopmatrix<4x4xf32, Subgroup, MatrixAcc>>, Function>
+ // CHECK: spirv.Store "Function" %[[STORE_PTR]], %{{.+}} : !spirv.coopmatrix<4x4xf32, Subgroup, MatrixAcc>
+ // CHECK: spirv.Load "Function" %[[STORE_PTR]] : !spirv.coopmatrix<4x4xf32, Subgroup, MatrixAcc>
+ %array = memref.alloca() : memref<16x!gpu.mma_matrix<4x4xf32, "COp">, #spirv.storage_class<Function>>
+ %value = gpu.subgroup_mma_constant_matrix %scalar : !gpu.mma_matrix<4x4xf32, "COp">
+ memref.store %value, %array[%index] : memref<16x!gpu.mma_matrix<4x4xf32, "COp">, #spirv.storage_class<Function>>
+ %loaded = memref.load %array[%index] : memref<16x!gpu.mma_matrix<4x4xf32, "COp">, #spirv.storage_class<Function>>
+ gpu.return
+ }
+
+ }
+}
+
+// -----
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.6,
+ [Shader, CooperativeMatrixKHR],
+ [SPV_KHR_storage_buffer_storage_class, SPV_KHR_cooperative_matrix]>,
+ #spirv.resource_limits<>>} {
+
+ gpu.module @kernels {
+ gpu.func @negative_gpu_wmma_workgroup_array() kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // expected-error @+1 {{failed to legalize operation 'memref.alloc'}}
+ %array = memref.alloc() : memref<16x!gpu.mma_matrix<4x4xf32, "COp">, #spirv.storage_class<Workgroup>>
+ gpu.return
+ }
+ }
+}
+
+// -----
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.6,
+ [Shader, CooperativeMatrixKHR],
+ [SPV_KHR_storage_buffer_storage_class, SPV_KHR_cooperative_matrix]>,
+ #spirv.resource_limits<>>} {
+
+ gpu.module @kernels {
+ gpu.func @negative_gpu_wmma_strided_function_array() kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // expected-error @+1 {{failed to legalize operation 'memref.alloca'}}
+ %array = memref.alloca() : memref<16x!gpu.mma_matrix<4x4xf32, "COp">, strided<[2]>, #spirv.storage_class<Function>>
+ gpu.return
+ }
}
}
diff --git a/mlir/test/Conversion/MemRefToSPIRV/alloca.mlir b/mlir/test/Conversion/MemRefToSPIRV/alloca.mlir
index 58847d114df00..dd65c09269537 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/alloca.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/alloca.mlir
@@ -45,6 +45,25 @@ module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader
// CHECK-DAG: spirv.Variable : !spirv.ptr<!spirv.struct<(!spirv.array<4 x vector<4xf32>>)>, Function>
+// -----
+
+module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader], []>, #spirv.resource_limits<>>} {
+ func.func @nested_alloca(%condition: i1, %index: index) {
+ scf.if %condition {
+ %array = memref.alloca() : memref<4xf32, #spirv.storage_class<Function>>
+ %value = memref.load %array[%index] : memref<4xf32, #spirv.storage_class<Function>>
+ memref.store %value, %array[%index] : memref<4xf32, #spirv.storage_class<Function>>
+ }
+ return
+ }
+}
+
+// CHECK-LABEL: func @nested_alloca
+// CHECK: %[[ARRAY:.+]] = spirv.Variable
+// CHECK: scf.if
+// CHECK: spirv.Load "Function" {{.*}} : f32
+// CHECK: spirv.Store "Function" {{.*}} : f32
+
// -----
module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader], []>, #spirv.resource_limits<>>} {
``````````
</details>
https://github.com/llvm/llvm-project/pull/212806
More information about the Mlir-commits
mailing list