[mlir][spirv] Add GPU subgroup MMA to spirv.MMAMatrixTimesScalar

[antiagainst]

Summary

Along the way, make the default pattern fail instead of crashing
when an elementwise op is not supported yet.

[kuhar]

Just some nits

mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp

31

nit: s/element wise/elementwise/ for consistency

32

nit: createElementwiseOp to much the capitalization of SubgroupMmaElementwiseOp and other names in mlir?

207–208

Maybe move this above the previous check? I'd assume that it's more useful to know that the number of operands is off.

214

nit: can you split this into two separate definitions?

231

nit: flip the if condition and use early return? IE auto cc = splat.get...; if (!cc) return failure()

[antiagainst]

Address comments

mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp

 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/TypeUtilities.h"
 
 using namespace mlir;
 
-// See SPV_NV_cooperative_matrix for supported element wise ops.
-static void createElementWiseOp(ConversionPatternRewriter &builder,
+/// 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.
+///
+/// See SPV_NV_cooperative_matrix for supported element wise ops.
+static bool createElementWiseOp(ConversionPatternRewriter &builder,
                                 gpu::SubgroupMmaElementwiseOp op,
                                 spirv::CooperativeMatrixNVType coopType,
                                 ValueRange operands) {
   switch (op.getOpType()) {
   case gpu::MMAElementwiseOp::ADDF:
     builder.replaceOpWithNewOp<spirv::FAddOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::ADDI:
     builder.replaceOpWithNewOp<spirv::IAddOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::SUBF:
     builder.replaceOpWithNewOp<spirv::FSubOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::SUBI:
     builder.replaceOpWithNewOp<spirv::ISubOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::DIVF:
     builder.replaceOpWithNewOp<spirv::FDivOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::DIVS:
     builder.replaceOpWithNewOp<spirv::SDivOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::DIVU:
     builder.replaceOpWithNewOp<spirv::UDivOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::NEGATEF:
     builder.replaceOpWithNewOp<spirv::FNegateOp>(op, coopType, operands);
-    return;
+    return true;
   case gpu::MMAElementwiseOp::NEGATES:
     builder.replaceOpWithNewOp<spirv::SNegateOp>(op, coopType, operands);
-    return;
+    return true;
   default:
-    llvm_unreachable("unknown op");
+    break;
   }
+  return false;
 }
 
 namespace {
 
 /// This class implements the conversion of GPU MMA loadOp to
 /// CooperativeMatrixLoad op in the SPIRV dialect.
 struct WmmaLoadOpToSPIRVLowering
     : public OpConversionPattern<gpu::SubgroupMmaLoadMatrixOp> {
matchAndRewrite(gpu::SubgroupMmaConstantMatrixOp subgroupMmaConstantMatrixOp,
     auto coopType = convertMMAToSPIRVType(
         subgroupMmaConstantMatrixOp.getType().cast<gpu::MMAMatrixType>());
     rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
         subgroupMmaConstantMatrixOp, coopType, cst);
     return success();
   }
 };
 
-/// Converts elementwise ops to SPIR-V cooperative matrix elementwise ops.
-struct WmmaElementwiseOpToSPIRVLowering
+/// Converts elementwise ops to SPIR-V cooperative matrix elementwise ops for
+/// the default case.
+struct WmmaElementwiseOpToSPIRVDefaultLowering
     : public OpConversionPattern<gpu::SubgroupMmaElementwiseOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(gpu::SubgroupMmaElementwiseOp subgroupMmaElementwiseOp,
+  matchAndRewrite(gpu::SubgroupMmaElementwiseOp elementwiseOp,
                   OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // All operands should be of cooperative matrix types.
     for (Value operand : adaptor.getOperands()) {
       if (!operand.getType().isa<spirv::CooperativeMatrixNVType>())
         return failure();
     }
     auto coopType = convertMMAToSPIRVType(
-        subgroupMmaElementwiseOp.getType().cast<gpu::MMAMatrixType>());
-    createElementWiseOp(rewriter, subgroupMmaElementwiseOp, coopType,
-                        adaptor.getOperands());
+        elementwiseOp.getType().cast<gpu::MMAMatrixType>());
+    return success(createElementWiseOp(rewriter, elementwiseOp, coopType,
+                                       adaptor.getOperands()));
+  }
+};
+
+/// Converts elementwise ops to SPIR-V cooperative matrix elementwise ops for
+/// matrix times scalar case.
+struct WmmaElementwiseOpToSPIRVScalarMulLowering
+    : public OpConversionPattern<gpu::SubgroupMmaElementwiseOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(gpu::SubgroupMmaElementwiseOp elementwiseOp,
+                  OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // All operands should be of cooperative matrix types.
+    for (Value operand : adaptor.getOperands()) {
+      if (!operand.getType().isa<spirv::CooperativeMatrixNVType>())
+        return failure();
+    }
+    if (adaptor.getOperands().size() != 2)
+      return failure();
+
+    // Use the original operands to check whether one of the operands is a splat
+    // scalar value.
+    Value lhs = elementwiseOp.getOperands().front();
+    Value rhs = elementwiseOp.getOperands().back();
+    Value splat = nullptr, matrix = nullptr;
+    if (lhs.getDefiningOp<gpu::SubgroupMmaConstantMatrixOp>()) {
+      splat = adaptor.getOperands().front();
+      matrix = adaptor.getOperands().back();
+    } else if (rhs.getDefiningOp<gpu::SubgroupMmaConstantMatrixOp>()) {
+      matrix = adaptor.getOperands().front();
+      splat = adaptor.getOperands().back();
+    }
+    if (!splat || !matrix)
+      return failure();
+
+    // Constant MMA matrix ops are converted to spirv.CompositeConstruct ops.
+    Value scalar = nullptr;
+    if (auto cc = splat.getDefiningOp<spirv::CompositeConstructOp>()) {
+      assert(cc.getConstituents().size() == 1);
+      scalar = cc.getConstituents().front();
+    } else {
+      return failure();
+    }
+
+    auto coopType = convertMMAToSPIRVType(
+        elementwiseOp.getType().cast<gpu::MMAMatrixType>());
+    rewriter.replaceOpWithNewOp<spirv::MatrixTimesScalarOp>(
+        elementwiseOp, coopType, ValueRange{matrix, scalar});
     return success();
   }
 };
 
 } // namespace
 
 /// Return the LLVMStructureType corresponding to the MMAMatrixType `type`.
 mlir::spirv::CooperativeMatrixNVType
 mlir::convertMMAToSPIRVType(gpu::MMAMatrixType type) {
   ArrayRef<int64_t> retTypeShape = type.getShape();
   Type elementType = type.getElementType();
   return spirv::CooperativeMatrixNVType::get(
       elementType, spirv::Scope::Subgroup, retTypeShape[0], retTypeShape[1]);
 }
 
 void mlir::populateGpuWMMAToSPIRVConversionPatterns(
     SPIRVTypeConverter &converter, RewritePatternSet &patterns) {
+  MLIRContext *context = patterns.getContext();
   patterns.add<WmmaLoadOpToSPIRVLowering, WmmaMmaOpToSPIRVLowering,
                WmmaStoreOpToSPIRVLowering, WmmaConstantOpToSPIRVLowering,
-               WmmaElementwiseOpToSPIRVLowering>(converter,
-                                                 patterns.getContext());
+               WmmaElementwiseOpToSPIRVDefaultLowering>(converter, context);
+  // Give the following patterns higher benefit to prevail over the default one.
+  patterns.add<WmmaElementwiseOpToSPIRVScalarMulLowering>(converter, context,
+                                                          /*benefit=*/2);
 }

mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -convert-gpu-to-spirv -split-input-file -verify-diagnostics %s | FileCheck %s
 
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, CooperativeMatrixNV, Float16], [SPV_KHR_storage_buffer_storage_class, SPV_NV_cooperative_matrix]>, #spirv.resource_limits<>>} {
   gpu.module @kernels {
-    // CHECK:       spirv.module @{{.*}} Logical GLSL450 {
     // CHECK-LABEL: spirv.func @gpu_wmma_load_op
-    // CHECK-SAME: {{%.*}}: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
+    // CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer>
     gpu.func @gpu_wmma_load_op(%arg0 : memref<32x32xf16, #spirv.storage_class<StorageBuffer>>) kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       %i = arith.constant 16 : index
       %j = arith.constant 16 : index
       // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixLoad {{%.*}}, {{%.*}}, {{%.*}} :  !spirv.ptr<f32, StorageBuffer> as !spirv.coopmatrix<16x16xf16, Subgroup>
       %0 = gpu.subgroup_mma_load_matrix %arg0[%i, %j] {leadDimension = 32 : index} : memref<32x32xf16, #spirv.storage_class<StorageBuffer>> -> !gpu.mma_matrix<16x16xf16, "COp">
       // CHECK: spirv.Return
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, CooperativeMatrixNV, Float16], [SPV_KHR_storage_buffer_storage_class, SPV_NV_cooperative_matrix]>, #spirv.resource_limits<>>} {
   gpu.module @kernels {
-    // CHECK:       spirv.module @{{.*}} Logical GLSL450 {
     // CHECK-LABEL: spirv.func @gpu_wmma_store_op
-    // CHECK-SAME: {{%.*}}: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
-    // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>})
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
+    // CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<512 x f32, stride=4> [0])>, StorageBuffer>
+    // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup>
     gpu.func @gpu_wmma_store_op(%arg0 : memref<32x32xf16, #spirv.storage_class<StorageBuffer>>, %arg1 : !gpu.mma_matrix<16x16xf16, "COp">) kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       %i = arith.constant 16 : index
       %j = arith.constant 16 : index
       //  CHECK: spirv.NV.CooperativeMatrixStore {{%.*}}, {{%.*}}, {{%.*}}, {{%.*}} : !spirv.ptr<f32, StorageBuffer>, !spirv.coopmatrix<16x16xf16, Subgroup>
       gpu.subgroup_mma_store_matrix %arg1, %arg0[%i,%j] {leadDimension= 32 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<32x32xf16,  #spirv.storage_class<StorageBuffer>>
       // CHECK: spirv.Return
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, CooperativeMatrixNV, Float16], [SPV_KHR_storage_buffer_storage_class, SPV_NV_cooperative_matrix]>, #spirv.resource_limits<>>} {
   gpu.module @kernels {
-    // CHECK:       spirv.module @{{.*}} Logical GLSL450 {
     // CHECK-LABEL: spirv.func @gpu_wmma_mma_op
-    // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
-    // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>}
-    // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 2)>})
-    // CHECK-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
+    // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup>
+    // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup>
+    // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup>
     gpu.func @gpu_wmma_mma_op(%A : !gpu.mma_matrix<16x16xf16, "AOp">, %B : !gpu.mma_matrix<16x16xf16, "BOp">, %C : !gpu.mma_matrix<16x16xf16, "COp">) kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup>, !spirv.coopmatrix<16x16xf16, Subgroup> -> !spirv.coopmatrix<16x16xf16, Subgroup>
       %D = gpu.subgroup_mma_compute %A, %B, %C : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
       // CHECK: spirv.Return
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, CooperativeMatrixNV, Float16], [SPV_KHR_storage_buffer_storage_class, SPV_NV_cooperative_matrix]>, #spirv.resource_limits<>>} {
   gpu.module @kernels {
-    // CHECK:       spirv.module @{{.*}} Logical GLSL450 {
     // CHECK-LABEL: spirv.func @gpu_wmma_constant_op
     gpu.func @gpu_wmma_constant_op() kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       // CHECK: {{%.*}} = spirv.Constant
       %cst = arith.constant 1.0 : f16
       // CHECK: {{%.*}} = spirv.CompositeConstruct {{%.*}} : (f16) -> !spirv.coopmatrix<16x16xf16, Subgroup>
       %C = gpu.subgroup_mma_constant_matrix %cst : !gpu.mma_matrix<16x16xf16, "COp">
       // CHECK: spirv.Return
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {
   gpu.container_module,
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, CooperativeMatrixNV, Float16], [SPV_KHR_storage_buffer_storage_class, SPV_NV_cooperative_matrix]>, #spirv.resource_limits<>>} {
   gpu.module @kernels {
-    // CHECK:       spirv.module @{{.*}} Logical GLSL450 {
-    // CHECK-LABEL: spirv.func @gpu_wmma_elementwise_op
-    // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>}
-    // CHECK-SAME: {{%.*}}: !spirv.coopmatrix<16x16xf16, Subgroup> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 1)>})
-    gpu.func @gpu_wmma_elementwise_op(%A : !gpu.mma_matrix<16x16xf16, "COp">, %B : !gpu.mma_matrix<16x16xf16, "COp">) kernel
+    // CHECK-LABEL: spirv.func @gpu_wmma_elementwise_op_default
+    // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup>
+    // CHECK-SAME: !spirv.coopmatrix<16x16xf16, Subgroup>
+    gpu.func @gpu_wmma_elementwise_op_default(%A : !gpu.mma_matrix<16x16xf16, "COp">, %B : !gpu.mma_matrix<16x16xf16, "COp">) kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       // CHECK:  {{%.*}} = spirv.FAdd {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup>
       %C = gpu.subgroup_mma_elementwise addf %A, %B : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
       // CHECK:  {{%.*}} = spirv.FNegate {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup>
       %D = gpu.subgroup_mma_elementwise negatef %C : (!gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
       // CHECK:  {{%.*}} = spirv.FDiv {{%.*}}, {{%.*}} : !spirv.coopmatrix<16x16xf16, Subgroup>
       %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.Return
       gpu.return
     }
   }
-}
+}
+
+// -----
+
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, CooperativeMatrixNV, Float16], [SPV_KHR_storage_buffer_storage_class, SPV_NV_cooperative_matrix]>, #spirv.resource_limits<>>} {
+  gpu.module @kernels {
+    // CHECK-LABEL: spirv.func @gpu_wmma_elementwise_op_matrix_times_scalar
+    // CHECK-SAME:    %[[A:.+]]: !spirv.coopmatrix<16x16xf16, Subgroup>
+    // CHECK-SAME:    %[[S:.+]]: f16
+    gpu.func @gpu_wmma_elementwise_op_matrix_times_scalar(%A : !gpu.mma_matrix<16x16xf16, "COp">, %scalar : f16) kernel
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
+      %B = gpu.subgroup_mma_constant_matrix %scalar : !gpu.mma_matrix<16x16xf16, "COp">
+      // CHECK: %{{.+}} = spirv.MatrixTimesScalar %[[A]], %[[S]] : !spirv.coopmatrix<16x16xf16, Subgroup>, f16
+      %C = gpu.subgroup_mma_elementwise mulf %A, %B : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
+      // CHECK: %{{.+}} = spirv.MatrixTimesScalar %[[A]], %[[S]] : !spirv.coopmatrix<16x16xf16, Subgroup>, f16
+      %D = gpu.subgroup_mma_elementwise mulf %B, %A : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
+      gpu.return
+    }
+  }
+}