[AMD] Properly categorize gfx950 in ISA family (triton-lang#5873)

While also here, audit features and enable them for gfx950 when
sutiable.

Author: antiagainst

third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h

@@ -11,6 +11,7 @@ enum class ISAFamily {
   CDNA1,
   CDNA2,
   CDNA3,
+  CDNA4,
   RDNA1,
   RDNA2,
   RDNA3,

third_party/amd/lib/TritonAMDGPUToLLVM/BufferOpsEmitter.cpp

@@ -2,15 +2,11 @@
 #include "TargetInfo.h"
 #include "Utility.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
-#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/IR/PatternMatch.h"
-#include "triton/Conversion/TritonGPUToLLVM/TypeConverter.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
 #include "BufferOpsEmitter.h"
 
-using mlir::triton::gpu::appendOrGetExternFuncOp;
-using mlir::triton::gpu::getFunctionType;
 using namespace triton::AMD;
 
 namespace {
@@ -66,7 +62,8 @@ Value BufferEmitter::createResourceDescriptor(Value basePtr,
   }
 
   Value stride = b.int_val(16, 0);
-  if (targetInfo.getISAFamily() == ISAFamily::CDNA3) {
+  if (llvm::is_contained({ISAFamily::CDNA3, ISAFamily::CDNA4},
+                         targetInfo.getISAFamily())) {
     if (blockStride) { // TODO: BufferAtomicRMWOp is unsupported
       Value enableSwizzle = b.int_val(16, 16384);
       Value mask14b = b.int_val(16, 16383);

third_party/amd/lib/TritonAMDGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -23,6 +23,7 @@ using ::mlir::LLVM::getSharedMemoryBase;
 using ::mlir::LLVM::AMD::getVectorSize;
 using ::mlir::LLVM::AMD::llLoad;
 using ::mlir::LLVM::AMD::llStore;
+using ::mlir::triton::AMD::ISAFamily;
 using ::mlir::triton::gpu::getTotalElemsPerThread;
 
 namespace {
@@ -408,25 +409,18 @@ struct AsyncCopyGlobalToLocalOpConversion
 
   bool supportsLoadWidth(unsigned bits,
                          const AMD::TargetInfo &targetInfo) const {
-    llvm::SmallSetVector<unsigned, 10> supportedWidths;
-    using mlir::triton::AMD::ISAFamily;
     switch (targetInfo.getISAFamily()) {
     case ISAFamily::CDNA1:
     case ISAFamily::CDNA2:
     case ISAFamily::CDNA3:
-      supportedWidths.insert(8);
-      supportedWidths.insert(16);
-      supportedWidths.insert(32);
-      if (targetInfo.getGPUKind() == llvm::AMDGPU::GPUKind::GK_GFX950) {
-        supportedWidths.insert(96);
-        supportedWidths.insert(128);
-      }
-      break;
+      return llvm::is_contained({32, 16, 8}, bits);
+    case ISAFamily::CDNA4:
+      return llvm::is_contained({128, 96, 32, 16, 8}, bits);
     default:
-      return false;
+      break;
     }
 
-    return supportedWidths.contains(bits);
+    return false;
   }
 
   LogicalResult
@@ -1120,10 +1114,17 @@ struct AtomicCASOpConversion
   }
 };
 
-bool supportsGlobalAtomicF16PackedAndDpp(triton::AMD::ISAFamily isaFamily) {
-  return isaFamily == triton::AMD::ISAFamily::CDNA1 ||
-         isaFamily == triton::AMD::ISAFamily::CDNA2 ||
-         isaFamily == triton::AMD::ISAFamily::CDNA3;
+bool supportsGlobalAtomicF16PackedAndDpp(ISAFamily isaFamily) {
+  switch (isaFamily) {
+  case ISAFamily::CDNA1:
+  case ISAFamily::CDNA2:
+  case ISAFamily::CDNA3:
+  case ISAFamily::CDNA4:
+    return true;
+  default:
+    break;
+  }
+  return false;
 }
 
 Value generateI32DppMove(PatternRewriter &rewriter, Value val, int dppCtrl) {
@@ -1284,11 +1285,12 @@ struct AtomicRMWOpConversion
     int numElems = 1;
     Type packF16Ty = vec_ty(valueElemTy, 2);
 
-    // CDNA3 arch allows to accelerate its atomics with LDS reduction algorithm,
-    // which is only applicable for atomics with no return. Otherwise we have to
-    // deal with an additional overhead.
+    // CDNA3/CDNA4 arch allows to accelerate its atomics with LDS reduction
+    // algorithm, which is only applicable for atomics with no return. Otherwise
+    // we have to deal with an additional overhead.
     bool enableIntraWaveReduce =
-        targetInfo.getISAFamily() == triton::AMD::ISAFamily::CDNA3 &&
+        llvm::is_contained({ISAFamily::CDNA3, ISAFamily::CDNA4},
+                           targetInfo.getISAFamily()) &&
         tensorTy && opResult.use_empty();
 
     // TODO: support data types less than 32 bits
@@ -1648,17 +1650,15 @@ struct AsyncWaitOpConversion : public ConvertOpToLLVMPattern<AsyncWaitOp> {
   LogicalResult
   matchAndRewrite(AsyncWaitOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-
-    using mlir::triton::AMD::ISAFamily;
-
     switch (targetInfo.getISAFamily()) {
     case ISAFamily::CDNA1:
     case ISAFamily::CDNA2:
     case ISAFamily::CDNA3:
+    case ISAFamily::CDNA4:
       break;
     default:
       return rewriter.notifyMatchFailure(
-          op, "Only supported on target architecture");
+          op, "Only supported on CDNA target architecture");
     }
 
     auto loc = op->getLoc();

third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp

@@ -1,6 +1,7 @@
 #include "TargetInfo.h"
 #include "SchedInstructions.h"
 #include "TritonAMDGPUToLLVM/GCNAsmFormat.h"
+#include "TritonAMDGPUToLLVM/TargetUtils.h"
 #include "Utility.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -193,8 +194,9 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
   if (numLaneToReduce != 64)
     return false;
 
-  if (auto family = getISAFamily();
-      family != ISAFamily::CDNA3 && family != ISAFamily::CDNA2) {
+  if (!llvm::is_contained(
+          {ISAFamily::CDNA2, ISAFamily::CDNA3, ISAFamily::CDNA4},
+          getISAFamily())) {
     return false;
   }
 

third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp

@@ -12,6 +12,7 @@ ISAFamily deduceISAFamily(llvm::StringRef arch) {
   // CDNA ISA cases
   switch (kind) {
   case llvm::AMDGPU::GK_GFX950:
+    return ISAFamily::CDNA4;
   case llvm::AMDGPU::GK_GFX942:
   case llvm::AMDGPU::GK_GFX941:
   case llvm::AMDGPU::GK_GFX940:
@@ -40,6 +41,7 @@ bool supportsVDot(llvm::StringRef arch) {
   case AMD::ISAFamily::CDNA1:
   case AMD::ISAFamily::CDNA2:
   case AMD::ISAFamily::CDNA3:
+  case AMD::ISAFamily::CDNA4:
   case AMD::ISAFamily::RDNA2:
   case AMD::ISAFamily::RDNA3:
     return true;

third_party/amd/lib/TritonAMDGPUToLLVM/Utility.cpp

@@ -145,9 +145,11 @@ static Value shuffleCommonImpl(Location loc, RewriterBase &rewriter,
       Value offset = b.i32_val(0x401F);
       return rewriter.create<ROCDL::DsSwizzleOp>(loc, valType, val, offset);
     } else {
-      if (isaFamily != ISAFamily::CDNA2 && isaFamily != ISAFamily::CDNA3) {
-        // DPP is only supportted for CDNA2 and CDNA3 right now, so we fallback
-        // to ds_swizzle for other archs.
+      if (!llvm::is_contained(
+              {ISAFamily::CDNA2, ISAFamily::CDNA3, ISAFamily::CDNA4},
+              isaFamily)) {
+        // DPP is only supported for CDNA2/CDNA3/CDNA4 right now, so we fallback
+        // to ds_swizzle for other architectures.
         //
         // This map facilates the butterfly shuffle pattern for a stride less
         // than 16. The pattern stride is the key of the map.

third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp

@@ -24,6 +24,8 @@ int getMfmaVersion(ISAFamily isaFamily) {
     return 2;
   case ISAFamily::CDNA3:
     return 3;
+  case ISAFamily::CDNA4:
+    return 4;
   default:
     break;
   }

third_party/amd/python/triton_amd.cc

@@ -273,6 +273,7 @@ void init_triton_amd(py::module &&m) {
   m.def("has_matrix_core_feature", [](const std::string &arch) {
     using mlir::triton::AMD::ISAFamily;
     switch (mlir::triton::AMD::deduceISAFamily(arch)) {
+    case ISAFamily::CDNA4:
     case ISAFamily::CDNA3:
     case ISAFamily::CDNA2:
     case ISAFamily::CDNA1: