From b3b7e385a1d0513b249a9db6ff7e93619b23f64f Mon Sep 17 00:00:00 2001 From: Hsiangkai Wang Date: Mon, 15 Sep 2025 14:42:32 +0100 Subject: [PATCH] [mlir][gpu][spirv] Add conversion for gpu.subgroup_mma_elementwise mulf gpu.subgroup_mma_elementwise supports mulf op type. Add conversion for it. --- mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp | 3 +++ .../GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir | 7 +++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp index 51dc50048024f..c1b09f18af560 100644 --- a/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp @@ -55,6 +55,9 @@ static bool createElementwiseOp(ConversionPatternRewriter &builder, case gpu::MMAElementwiseOp::SUBI: builder.replaceOpWithNewOp(op, coopType, operands); return true; + case gpu::MMAElementwiseOp::MULF: + builder.replaceOpWithNewOp(op, coopType, operands); + return true; case gpu::MMAElementwiseOp::DIVF: builder.replaceOpWithNewOp(op, coopType, operands); return true; 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 7ef3711ebe28b..6dba9c3486c7b 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 @@ -136,14 +136,17 @@ module attributes { // CHECK: {{%.*}} = spirv.FDiv {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc> %E = gpu.subgroup_mma_elementwise divf %D, %A : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp"> + // CHECK: {{%.*}} = spirv.FMul {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc> + %F = gpu.subgroup_mma_elementwise mulf %E, %A : + (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp"> // CHECK: {{%.*}} = spirv.FConvert {{%.*}} : // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup, MatrixAcc> to !spirv.coopmatrix<16x16xf32, Subgroup, MatrixAcc> - %F = gpu.subgroup_mma_elementwise extf %E : + %G = gpu.subgroup_mma_elementwise extf %F : (!gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp"> %i = arith.constant 0 : index // CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %{{.+}}, %{{.+}}, - gpu.subgroup_mma_store_matrix %F, %ptr[%i,%i] {leadDimension = 32 : index} : + gpu.subgroup_mma_store_matrix %G, %ptr[%i,%i] {leadDimension = 32 : index} : !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32, #spirv.storage_class> // CHECK: spirv.Return gpu.return