[MLIR][AMDGPU] Changes from the review.

Author: pcf000

mlir/include/mlir/Dialect/AMDGPU/Utils/Chipset.h

@@ -47,15 +47,16 @@ struct Chipset {
   DEFINE_COMP_OPERATOR(>)
   DEFINE_COMP_OPERATOR(>=)
 #undef DEFINE_COMP_OPERATOR
-
-  bool isGfx940() const {
-    return majorVersion == 9 && minorVersion >= 4 && minorVersion < 5;
-  }
-  bool hasOcpFp8() const {
-    return (majorVersion == 9 && minorVersion >= 5) || majorVersion >= 12;
-  }
 };
 
+inline bool isGfx940Series(const Chipset &chipset) {
+  return chipset.majorVersion == 9 && chipset.minorVersion == 4;
+}
+inline bool hasOcpFp8(const Chipset &chipset) {
+  return (chipset.majorVersion == 9 && chipset.minorVersion >= 5) ||
+         chipset.majorVersion >= 12;
+}
+
 } // namespace mlir::amdgpu
 
 #endif

mlir/include/mlir/IR/Types.h

@@ -132,6 +132,9 @@ class Type {
   bool isF64() const;
   bool isF80() const;
   bool isF128() const;
+  /// Return true if this is an float type (with the specified width).
+  bool isFloat() const;
+  bool isFloat(unsigned width) const;
 
   /// Return true if this is an integer type (with the specified width).
   bool isInteger() const;

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -471,15 +471,15 @@ static void wmmaPushOutputOperand(ConversionPatternRewriter &rewriter,
 /// Return true if `type` is the E5M2 variant of an 8-bit float that is
 /// supported by the `_bf8` instructions on the given `chipset`.
 static bool isNativeBf8(Chipset chipset, Type type) {
-  return (chipset.isGfx940() && isa<Float8E5M2FNUZType>(type)) ||
-         (chipset.hasOcpFp8() && isa<Float8E5M2Type>(type));
+  return (isGfx940Series(chipset) && isa<Float8E5M2FNUZType>(type)) ||
+         (hasOcpFp8(chipset) && isa<Float8E5M2Type>(type));
 }
 
 /// Return true if `type` is the E4M3FN variant of an 8-bit float that is
 /// supported by the `_fp8` instructions on the given `chipset`.
 static bool isNativeFp8(Chipset chipset, Type type) {
-  return (chipset.isGfx940() && isa<Float8E4M3FNUZType>(type)) ||
-         (chipset.hasOcpFp8() && isa<Float8E4M3FNType>(type));
+  return (isGfx940Series(chipset) && isa<Float8E4M3FNUZType>(type)) ||
+         (hasOcpFp8(chipset) && isa<Float8E4M3FNType>(type));
 }
 
 /// Return the `rocdl` intrinsic corresponding to a MFMA operation `mfma`
@@ -787,7 +787,7 @@ LogicalResult ExtPackedFp8OpLowering::matchAndRewrite(
     ExtPackedFp8Op op, ExtPackedFp8OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   Location loc = op.getLoc();
-  if (!(chipset.isGfx940() || chipset.hasOcpFp8()))
+  if (!(isGfx940Series(chipset) || hasOcpFp8(chipset)))
     return rewriter.notifyMatchFailure(
         loc, "Fp8 conversion instructions are not available on target "
              "architecture and their emulation is not implemented");
@@ -831,7 +831,7 @@ LogicalResult PackedTrunc2xFp8OpLowering::matchAndRewrite(
     PackedTrunc2xFp8Op op, PackedTrunc2xFp8OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   Location loc = op.getLoc();
-  if (!(chipset.isGfx940() || chipset.hasOcpFp8()))
+  if (!(isGfx940Series(chipset) || hasOcpFp8(chipset)))
     return rewriter.notifyMatchFailure(
         loc, "Fp8 conversion instructions are not available on target "
              "architecture and their emulation is not implemented");
@@ -868,7 +868,7 @@ LogicalResult PackedStochRoundFp8OpLowering::matchAndRewrite(
     PackedStochRoundFp8Op op, PackedStochRoundFp8OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   Location loc = op.getLoc();
-  if (!(chipset.isGfx940() || chipset.hasOcpFp8()))
+  if (!(isGfx940Series(chipset) || hasOcpFp8(chipset)))
     return rewriter.notifyMatchFailure(
         loc, "Fp8 conversion instructions are not available on target "
              "architecture and their emulation is not implemented");

mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp

@@ -72,11 +72,11 @@ struct TruncfToFloat16RewritePattern final
 
 } // end namespace
 
-static LogicalResult isSupportedFp8(Type elementType, Chipset chipset) {
-  if (chipset.isGfx940())
-    return success(isa<Float8E5M2FNUZType, Float8E4M3FNUZType>(elementType));
-  if (chipset.hasOcpFp8())
-    return success(isa<Float8E5M2Type, Float8E4M3FNType>(elementType));
+static LogicalResult isSupportedF8(Type elementType, Chipset chipset) {
+  if (isGfx940Series(chipset))
+    return success(isa<Float8E4M3FNUZType, Float8E5M2FNUZType>(elementType));
+  if (hasOcpFp8(chipset))
+    return success(isa<Float8E4M3FNType, Float8E5M2Type>(elementType));
   return failure();
 }
 
@@ -98,7 +98,7 @@ LogicalResult ExtFOnFloat8RewritePattern::match(arith::ExtFOp op) const {
       return failure();
     inType = inVecType.getElementType();
   }
-  return isSupportedFp8(inType, chipset);
+  return isSupportedF8(inType, chipset);
 }
 
 void ExtFOnFloat8RewritePattern::rewrite(arith::ExtFOp op,
@@ -232,7 +232,7 @@ LogicalResult TruncFToFloat8RewritePattern::match(arith::TruncFOp op) const {
     // Conversion between 8-bit floats is not supported with truncation enabled.
     return failure();
 
-  return isSupportedFp8(outType, chipset);
+  return isSupportedF8(outType, chipset);
 }
 
 void TruncFToFloat8RewritePattern::rewrite(arith::TruncFOp op,
@@ -397,7 +397,7 @@ void ArithToAMDGPUConversionPass::runOnOperation() {
   }
 
   bool convertFP8Arithmetic =
-      maybeChipset->isGfx940() || maybeChipset->hasOcpFp8();
+      isGfx940Series(*maybeChipset) || hasOcpFp8(*maybeChipset);
   arith::populateArithToAMDGPUConversionPatterns(
       patterns, convertFP8Arithmetic, saturateFP8Truncf, allowPackedF16Rtz,
       *maybeChipset);

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

@@ -272,16 +272,14 @@ LogicalResult MFMAOp::verify() {
   }
 
   Type sourceBType = getSourceB().getType();
-  if (isa<Float8E5M2FNUZType, Float8E4M3FNUZType, Float8E5M2Type,
-          Float8E4M3FNType>(sourceElem)) {
+  if (sourceElem.isFloat(8)) {
     int64_t sourceBLen = 1;
     Type sourceBElem = sourceBType;
     if (auto sourceBVector = llvm::dyn_cast<VectorType>(sourceBType)) {
       sourceBLen = sourceBVector.getNumElements();
       sourceBElem = sourceBVector.getElementType();
     }
-    if (!isa<Float8E5M2FNUZType, Float8E4M3FNUZType, Float8E5M2Type,
-             Float8E4M3FNType>(sourceBElem))
+    if (!sourceBElem.isFloat(8))
       return emitOpError("expected both source operands to have f8 elements");
     if (sourceLen != sourceBLen)
       return emitOpError(

mlir/lib/IR/Types.cpp

@@ -42,6 +42,15 @@ bool Type::isF64() const { return llvm::isa<Float64Type>(*this); }
 bool Type::isF80() const { return llvm::isa<Float80Type>(*this); }
 bool Type::isF128() const { return llvm::isa<Float128Type>(*this); }
 
+bool Type::isFloat() const { return llvm::isa<FloatType>(*this); }
+
+/// Return true if this is an integer type with the specified width.
+bool Type::isFloat(unsigned width) const {
+  if (auto fltTy = llvm::dyn_cast<FloatType>(*this))
+    return fltTy.getWidth() == width;
+  return false;
+}
+
 bool Type::isIndex() const { return llvm::isa<IndexType>(*this); }
 
 bool Type::isInteger() const { return llvm::isa<IntegerType>(*this); }