[Mlir-commits] [mlir] [mlir][gpu][spirv] Convert memref<mma_matrix> to spv.array<coopmatrix> (PR #212806)

[llvmlistbot at llvm.org]

fabrizio-indirli wrote:

The reason for this patch is that I noticed that in KHR SPIRV and in GLSL it's possible to instantiate arrays of CoopMatrices , for example:
GLSL: `coopmat<float, gl_ScopeSubgroup, H, W, gl_MatrixUseAccumulator> out[8]`
SPIRV : `<8x!spirv.coopmatrix<HxWxf32, Subgroup, MatrixAcc>>`

The MLIR GPU dialect represents CoopMatrices as `gpu.mma_matrix` and exposes operations on them, but I couldn't find a way to build arrays of such matrices.

Since SPIRV allows creating Variables containg CoopMatrices only with Function-local (or Private) storage, and since memrefs are lowered naturally to SPIRV arrays, I thought that allowing `memref.alloca() : memref<Nx!gpu.mma_matrix<HxW, ...>` and lowering it to `spirv.Variable: spirv.ptr<spirv.array<N x coopmatrix>` could work here. I’m not completely certain that a memref is the ideal abstraction, but it seems preferable to introducing a dedicated `gpu.array_of_mma_matrices` type, since I assume we wouldn't want composite types for every possible element type when they would all lower to `spirv.array` (or similar constructs in other backends) anyway. Happy to know your opinions!

https://github.com/llvm/llvm-project/pull/212806