PR Review Round 1

Signed-off-by: Muzammiluddin Syed <[email protected]>

Author: Muzammiluddin-Syed-ECE

mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td

@@ -804,7 +804,7 @@ def AMDGPU_GatherToLDSOp :
                    TypeAttr:$transferType
                    )>,
     Results<(outs)> {
-  let summary = "MLIR wrapper for CDNA mfma instructions";
+  let summary = "MLIR wrapper for CDNA Gather to LDS instructions";
   let description = [{
     The `amdgpu.global_load` op is a wrapper around the `global_load_lds` instructions.
 
@@ -845,7 +845,7 @@ def AMDGPU_ScaledMFMAOp :
                    I32Attr:$opselA,
                    I32Attr:$opselB)>,
     Results<(outs MFMAOutTypes: $destD)> {
-  let summary = "MLIR wrapper for CDNA mfma instructions";
+  let summary = "MLIR wrapper for CDNA scaled mfma instructions";
   let description = [{
     The `amdgpu.scaled_mfma` op is an MLIR wrapper around intrinsics
     for various scaled versions of `mfma` instructions in the CDNA architecture, which perform
@@ -857,7 +857,7 @@ def AMDGPU_ScaledMFMAOp :
 
     Note, this wrapper allows specifying `vector<4Kxi8>` arguments to MFMA
     intrinsics that take an integer type of width `4K`. For example,
-    one can provide a vector<4xi8> as an argument to an MFMA instruction that
+    one can provide a `vector<4xi8>` as an argument to an MFMA instruction that
     logically takes 4 i8s but whose intrinsics are specified to take an i32.
     In these cases, the bytes in the vector will be concatenated in little-endian
     order (that is, v[0] will go to arg[7:0], v[1] to arg[15:8] and so on).
@@ -868,7 +868,7 @@ def AMDGPU_ScaledMFMAOp :
     size. 
     - `amdgpu.scaled_mfma` does not support broadcasting. So, `cbsz`, `abid`, and `blgp` 
     are omitted from this wrapper.
-    - The negateA, negateB, and negateC flags in `amdgpu.mfma` are only supported for 
+    - The `negateA`, `negateB`, and `negateC` flags in `amdgpu.mfma` are only supported for 
     double-precision operations on gfx94x and so are not included here. 
   }];
   let assemblyFormat = [{

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -23,6 +23,7 @@
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Casting.h"
 #include <optional>
 
 namespace mlir {
@@ -826,19 +827,20 @@ mfmaOpToScaledIntrinsic(Type aType, Type bType, Type destType, uint32_t m,
 }
 
 static std::optional<std::tuple<StringRef, uint32_t, uint32_t>>
-mfmaOpToScaledIntrinsic(MFMAOp mfma, Chipset chipset) {
-  return mfmaOpToScaledIntrinsic(
-      mfma.getSourceA().getType(), mfma.getSourceB().getType(),
-      mfma.getDestC().getType(), mfma.getM(), mfma.getN(), mfma.getK(),
-      mfma.getBlocks(), chipset);
-}
-
-static std::optional<std::tuple<StringRef, uint32_t, uint32_t>>
-mfmaOpToScaledIntrinsic(ScaledMFMAOp smfma, Chipset chipset) {
-  return mfmaOpToScaledIntrinsic(smfma.getSourceA().getType(),
-                                 smfma.getSourceB().getType(),
-                                 smfma.getDestC().getType(), smfma.getM(),
-                                 smfma.getN(), smfma.getK(), 1u, chipset);
+mfmaOpToScaledIntrinsic(Operation *op, Chipset chipset) {
+  if (auto mfma = llvm::dyn_cast_or_null<MFMAOp>(op)) {
+    return mfmaOpToScaledIntrinsic(
+        mfma.getSourceA().getType(), mfma.getSourceB().getType(),
+        mfma.getDestC().getType(), mfma.getM(), mfma.getN(), mfma.getK(),
+        mfma.getBlocks(), chipset);
+  }
+  if (auto smfma = llvm::dyn_cast_or_null<ScaledMFMAOp>(op)) {
+    return mfmaOpToScaledIntrinsic(smfma.getSourceA().getType(),
+                                   smfma.getSourceB().getType(),
+                                   smfma.getDestC().getType(), smfma.getM(),
+                                   smfma.getN(), smfma.getK(), 1u, chipset);
+  }
+  return std::nullopt;
 }
 
 /// Return the `rocdl` intrinsic corresponding to a WMMA operation `wmma`
@@ -964,7 +966,7 @@ struct MFMAOpLowering : public ConvertOpToLLVMPattern<MFMAOp> {
 
 struct ScaledMFMAOpLowering : public ConvertOpToLLVMPattern<ScaledMFMAOp> {
   ScaledMFMAOpLowering(const LLVMTypeConverter &converter, Chipset chipset)
-      : ConvertOpToLLVMPattern<ScaledMFMAOp>(converter), chipset(chipset) {}
+      : ConvertOpToLLVMPattern(converter), chipset(chipset) {}
 
   Chipset chipset;
 
@@ -986,7 +988,7 @@ struct ScaledMFMAOpLowering : public ConvertOpToLLVMPattern<ScaledMFMAOp> {
       return op.emitOpError(
           "no intrinsic matching Scaled MFMA size on given chipset");
 
-    StringRef intrinsicName = std::get<0>(*maybeScaledIntrinsic);
+    auto [intrinsicName, aTypeCode, bTypeCode] = *maybeScaledIntrinsic;
     OperationState loweredOp(loc, intrinsicName);
     loweredOp.addTypes(intrinsicOutType);
     loweredOp.addOperands(
@@ -997,7 +999,6 @@ struct ScaledMFMAOpLowering : public ConvertOpToLLVMPattern<ScaledMFMAOp> {
     Value scaleB = createI32Constant(rewriter, loc, adaptor.getScaleB());
     Value opselA = createI32Constant(rewriter, loc, adaptor.getOpselA());
     Value opselB = createI32Constant(rewriter, loc, adaptor.getOpselB());
-    auto [_scaledName, aTypeCode, bTypeCode] = *maybeScaledIntrinsic;
     loweredOp.addOperands({createI32Constant(rewriter, loc, aTypeCode),
                            createI32Constant(rewriter, loc, bTypeCode),
                            /*scale A byte=*/opselA, /*scale A=*/scaleA,

mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp

@@ -507,44 +507,35 @@ LogicalResult GatherToLDSOp::verify() {
 }
 
 LogicalResult ScaledMFMAOp::verify() {
-  unsigned opselA = getOpselA();
-  unsigned opselB = getOpselB();
-
-  opselA >>= 8;
-  opselB >>= 8;
+  unsigned opselA = getOpselA() >> 8;
+  unsigned opselB = getOpselB() >> 8;
 
   if (opselA != 0)
-    return emitOpError("Opsel A must be a zero extended 8 bit value.");
+    return emitOpError("Opsel A must be a zero extended 8 bit value");
 
   if (opselB != 0)
-    return emitOpError("Opsel B must be a zero extended 8 bit value.");
-
-  auto validType = [&](Type mlirElemType) {
-    return llvm::TypeSwitch<Type, bool>(mlirElemType)
-        .Case([](Float8E4M3FNType) { return true; })
-        .Case([](Float8E5M2Type) { return true; })
-        .Case([](Float6E2M3FNType) { return true; })
-        .Case([](Float6E3M2FNType) { return true; })
-        .Case([](Float4E2M1FNType) { return true; })
-        .Default([](Type) { return false; });
-  };
+    return emitOpError("Opsel B must be a zero extended 8 bit value");
+
+  auto isValidType =
+      llvm::IsaPred<Float8E4M3FNType, Float8E5M2Type, Float6E2M3FNType,
+                    Float6E3M2FNType, Float4E2M1FNType>;
 
   Type aType = getSourceA().getType();
   Type bType = getSourceB().getType();
   aType = getElementTypeOrSelf(aType);
   bType = getElementTypeOrSelf(bType);
-  if (!validType(aType))
-    return emitOpError("Source A must be of element type fp4, fp6 or fp8.");
-  if (!validType(bType))
-    return emitOpError("Source B must be of element type fp4, fp6 or fp8.");
+  if (!isValidType(aType))
+    return emitOpError("Source A must be of element type fp4, fp6 or fp8");
+  if (!isValidType(bType))
+    return emitOpError("Source B must be of element type fp4, fp6 or fp8");
 
   unsigned m = getM();
   unsigned n = getN();
   unsigned k = getK();
   bool tileConfig1 = (m == n && n == 32 && k == 64);
   bool tileConfig2 = (m == n && n == 16 && k == 128);
   if (!tileConfig1 && !tileConfig2)
-    return emitOpError("Invalid tile size for scaled mfma.");
+    return emitOpError("Invalid tile size for scaled mfma");
 
   return success();
 }