Revert "[MLIR][GPU] subgroup_mma fp64 extension" (#169049)

Reverts #165873

The revert is triggered by a failing integration test on a couple of
buildbots.

Author: fabianmcg

mlir/include/mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h

@@ -27,7 +27,7 @@ class MMAMatrixType;
 #define GEN_PASS_DECL_CONVERTGPUOPSTONVVMOPS
 #include "mlir/Conversion/Passes.h.inc"
 
-Type convertMMAToLLVMType(gpu::MMAMatrixType type);
+LLVM::LLVMStructType convertMMAToLLVMType(gpu::MMAMatrixType type);
 
 /// Configure target to convert from the GPU dialect to NVVM.
 void configureGpuToNVVMConversionLegality(ConversionTarget &target);

mlir/include/mlir/Dialect/GPU/IR/GPUBase.td

@@ -114,7 +114,7 @@ def GPU_MMAMatrix : DialectType<
   GPU_Dialect, IsMMAMatrixTypePred, "MMAMatrix type">;
 
 // Memref type acceptable to gpu.subgroup_mma_{load|store}_matrix ops.
-def GPU_MMAMemRef : MemRefOf<[I8, I32, F16, F32, F64, VectorOfRankAndType<[1], [I8, I32, F16, F32, F64]>]>;
+def GPU_MMAMemRef : MemRefOf<[I8, I32, F16, F32, VectorOfRankAndType<[1], [I8, I32, F16, F32]>]>;
 
 class MMAMatrixOf<list<Type> allowedTypes> :
   ContainerType<AnyTypeOf<allowedTypes>, IsMMAMatrixTypePred,

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -1872,7 +1872,7 @@ def GPU_SubgroupMmaStoreMatrixOp : GPU_Op<"subgroup_mma_store_matrix",
     ```
   }];
 
-  let arguments = (ins Arg<MMAMatrixOf<[SI8, UI8, I32, F16, F32, F64]>>:$src,
+  let arguments = (ins Arg<MMAMatrixOf<[SI8, UI8, I32, F16, F32]>>:$src,
                   Arg<GPU_MMAMemRef, "",[MemWriteAt<0, FullEffect>]>:$dstMemref,
                   Variadic<Index>:$indices,
                   IndexAttr:$leadDimension,
@@ -1919,9 +1919,9 @@ def GPU_SubgroupMmaComputeOp
     ```
   }];
 
-  let arguments = (ins Arg<MMAMatrixOf<[SI8, UI8, F16, F32, F64]>>:$opA,
-                  Arg<MMAMatrixOf<[SI8, UI8, F16, F32, F64]>>:$opB,
-                  Arg<MMAMatrixOf<[I32, F16, F32, F64]>>:$opC,
+  let arguments = (ins Arg<MMAMatrixOf<[SI8, UI8, F16, F32]>>:$opA,
+                  Arg<MMAMatrixOf<[SI8, UI8, F16, F32]>>:$opB,
+                  Arg<MMAMatrixOf<[I32, F16, F32]>>:$opC,
                   OptionalAttr<UnitAttr>:$a_transpose,
                   OptionalAttr<UnitAttr>:$b_transpose);
 

mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp

@@ -17,7 +17,6 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/IR/TypeUtilities.h"
-#include "mlir/IR/Types.h"
 
 using namespace mlir;
 
@@ -58,8 +57,7 @@ static NVVM::MMATypes getElementType(gpu::MMAMatrixType type) {
   if (type.getElementType().isF32())
     return type.getOperand() == "COp" ? NVVM::MMATypes::f32
                                       : NVVM::MMATypes::tf32;
-  if (type.getElementType().isF64())
-    return NVVM::MMATypes::f64;
+
   if (type.getElementType().isSignedInteger(8))
     return NVVM::MMATypes::s8;
   if (type.getElementType().isUnsignedInteger(8))
@@ -214,13 +212,8 @@ struct WmmaMmaOpToNVVMLowering
     // then passed on to the intrinsic call. Emit llvm ops to extract individual
     // values form lowered memrefs.
     SmallVector<Value> unpackedOps;
+
     auto unpackOp = [&](Value operand) {
-      // f64 a and b fragments are not structs but scalars.
-      if (!isa<LLVM::LLVMStructType>(operand.getType())) {
-        unpackedOps.push_back(operand);
-        return;
-      }
-      // every other type is lowered to an LLVM struct, extract the values.
       auto structType = cast<LLVM::LLVMStructType>(operand.getType());
       for (size_t i = 0, e = structType.getBody().size(); i < e; ++i) {
         Value toUse = LLVM::ExtractValueOp::create(rewriter, loc, operand, i);
@@ -283,16 +276,10 @@ struct WmmaConstantOpToNVVMLowering
       return failure();
     Location loc = subgroupMmaConstantOp.getLoc();
     Value cst = adaptor.getOperands()[0];
-    Type type = convertMMAToLLVMType(
+    LLVM::LLVMStructType type = convertMMAToLLVMType(
         cast<gpu::MMAMatrixType>(subgroupMmaConstantOp.getType()));
-    // If the element is not a struct, it means it's a scalar f64.
-    auto structType = dyn_cast<LLVM::LLVMStructType>(type);
-    if (!structType) {
-      rewriter.replaceOp(subgroupMmaConstantOp, cst);
-      return success();
-    }
     // If the element type is a vector create a vector from the operand.
-    if (auto vecType = dyn_cast<VectorType>(structType.getBody()[0])) {
+    if (auto vecType = dyn_cast<VectorType>(type.getBody()[0])) {
       Value vecCst = LLVM::PoisonOp::create(rewriter, loc, vecType);
       for (int64_t vecEl = 0; vecEl < vecType.getNumElements(); vecEl++) {
         Value idx = LLVM::ConstantOp::create(rewriter, loc,
@@ -302,8 +289,8 @@ struct WmmaConstantOpToNVVMLowering
       }
       cst = vecCst;
     }
-    Value matrixStruct = LLVM::PoisonOp::create(rewriter, loc, structType);
-    for (size_t i : llvm::seq(size_t(0), structType.getBody().size())) {
+    Value matrixStruct = LLVM::PoisonOp::create(rewriter, loc, type);
+    for (size_t i : llvm::seq(size_t(0), type.getBody().size())) {
       matrixStruct =
           LLVM::InsertValueOp::create(rewriter, loc, matrixStruct, cst, i);
     }
@@ -367,24 +354,10 @@ struct WmmaElementwiseOpToNVVMLowering
       return failure();
     Location loc = subgroupMmaElementwiseOp.getLoc();
     size_t numOperands = adaptor.getOperands().size();
-    Type destType = convertMMAToLLVMType(
+    LLVM::LLVMStructType destType = convertMMAToLLVMType(
         cast<gpu::MMAMatrixType>(subgroupMmaElementwiseOp.getType()));
-
-    // If the element is not a struct, it means it's a scalar f64.
-    LLVM::LLVMStructType structDestTy =
-        dyn_cast<LLVM::LLVMStructType>(destType);
-    if (!structDestTy) {
-      SmallVector<Value> operands;
-      for (auto operand : adaptor.getOperands()) {
-        operands.push_back(operand);
-      }
-      Value element = createScalarOp(
-          rewriter, loc, subgroupMmaElementwiseOp.getOpType(), operands);
-      rewriter.replaceOp(subgroupMmaElementwiseOp, element);
-      return success();
-    }
-    Value matrixStruct = LLVM::PoisonOp::create(rewriter, loc, structDestTy);
-    for (size_t i = 0, e = structDestTy.getBody().size(); i < e; ++i) {
+    Value matrixStruct = LLVM::PoisonOp::create(rewriter, loc, destType);
+    for (size_t i = 0, e = destType.getBody().size(); i < e; ++i) {
       SmallVector<Value> extractedOperands;
       for (size_t opIdx = 0; opIdx < numOperands; opIdx++) {
         extractedOperands.push_back(LLVM::ExtractValueOp::create(
@@ -404,18 +377,13 @@ struct WmmaElementwiseOpToNVVMLowering
 } // namespace
 
 /// Return the LLVMStructureType corresponding to the MMAMatrixType `type`.
-Type mlir::convertMMAToLLVMType(gpu::MMAMatrixType type) {
+LLVM::LLVMStructType mlir::convertMMAToLLVMType(gpu::MMAMatrixType type) {
   NVVM::MMAFrag frag = convertOperand(type.getOperand());
   NVVM::MMATypes eltType = getElementType(type);
   auto nRow = type.getShape()[0];
   auto nCol = type.getShape()[1];
   std::pair<Type, unsigned> typeInfo =
       NVVM::inferMMAType(eltType, frag, nRow, nCol, type.getContext());
-  // Special handling for f64 a and b fragments
-  Type f64Ty = Float64Type::get(type.getContext());
-  if (typeInfo.first == f64Ty && typeInfo.second == 1) {
-    return f64Ty;
-  }
   return LLVM::LLVMStructType::getLiteral(
       type.getContext(), SmallVector<Type, 8>(typeInfo.second, typeInfo.first));
 }

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -208,7 +208,7 @@ Type MMAMatrixType::getElementType() const { return getImpl()->elementType; }
 StringRef MMAMatrixType::getOperand() const { return getImpl()->getOperand(); }
 
 bool MMAMatrixType::isValidElementType(Type elementType) {
-  return elementType.isF16() || elementType.isF32() || elementType.isF64() ||
+  return elementType.isF16() || elementType.isF32() ||
          elementType.isUnsignedInteger(8) || elementType.isSignedInteger(8) ||
          elementType.isInteger(32);
 }
@@ -225,7 +225,7 @@ MMAMatrixType::verifyInvariants(function_ref<InFlightDiagnostic()> emitError,
 
   if (!MMAMatrixType::isValidElementType(elementType))
     return emitError()
-           << "MMAMatrixType elements must be SI8, UI8, I32, F16, F32, or F64";
+           << "MMAMatrixType elements must be SI8, UI8, I32, F16, or F32";
 
   return success();
 }

mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir

@@ -80,28 +80,6 @@ gpu.module @test_module {
 
 // -----
 
-gpu.module @test_module {
-
-  // CHECK-LABEL: func @gpu_wmma_f64_load_op() ->
-  // CHECK-SAME: f64
-  // CHECK32-LABEL: func @gpu_wmma_f64_load_op() ->
-  func.func @gpu_wmma_f64_load_op() -> (!gpu.mma_matrix<8x4xf64, "AOp">) {
-    %wg = memref.alloca() {alignment = 32} : memref<32x32xf64, 3>
-    %i = arith.constant 16 : index
-    %j = arith.constant 16 : index
-    %0 = gpu.subgroup_mma_load_matrix %wg[%i, %j] {leadDimension = 32 : index} : memref<32x32xf64, 3> -> !gpu.mma_matrix<8x4xf64, "AOp">
-    return %0 : !gpu.mma_matrix<8x4xf64, "AOp">
-    // CHECK: %[[MUL:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64
-    // CHECK: %[[ADD:.*]] = llvm.add %[[MUL]], %{{.*}} : i64
-    // CHECK: %[[GEP:.*]] = llvm.getelementptr %{{.*}}[%[[ADD]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f64
-    // CHECK: %[[C32_I32:.*]] = llvm.mlir.constant(32 : index) : i32
-    // CHECK: %[[LOAD:.*]] = nvvm.wmma.load %[[GEP]], %[[C32_I32]] {eltype = #nvvm.mma_type<f64>, frag = #nvvm.mma_frag<a>, k = 4 : i32, layout = #nvvm.mma_layout<row>, m = 8 : i32, n = 8 : i32} : (!llvm.ptr<3>) -> f64
-    // CHECK: llvm.return %[[LOAD]] : f64
-  }
-}
-
-// -----
-
 gpu.module @test_module {
 
   // CHECK-LABEL: func @gpu_wmma_store_op

mlir/test/Dialect/GPU/invalid.mlir

@@ -688,7 +688,7 @@ func.func @mmamatrix_operand_type(){
 func.func @mmamatrix_invalid_element_type(){
     %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
     %i = arith.constant 16 : index
-    // expected-error @+1 {{MMAMatrixType elements must be SI8, UI8, I32, F16, F32, or F64}}
+    // expected-error @+1 {{MMAMatrixType elements must be SI8, UI8, I32, F16, or F32}}
     %0 = gpu.subgroup_mma_load_matrix %wg[%i, %i] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xbf16, "AOp">
     return
 }
@@ -708,7 +708,7 @@ func.func @mmaLoadOp_identity_layout(){
 // -----
 
 func.func @mma_invalid_memref_type(%src: memref<32x4xvector<4x8xf32>>, %i: index) {
-    // expected-error @+1 {{operand #0 must be memref of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float or vector of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float values of ranks 1 values}}
+    // expected-error @+1 {{operand #0 must be memref of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float or vector of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float values of ranks 1 values}}
     %0 = gpu.subgroup_mma_load_matrix %src[%i, %i] {leadDimension = 4 : index} : memref<32x4xvector<4x8xf32>> -> !gpu.mma_matrix<16x16xf16, "AOp">
     return
 }