[AMD] Enable Pingpong by default on gfx950 arch (triton-lang#7697)

List of enabling conditions
- FP/BF16 GEMM with M,N>64 tilesize when num_stages=3 and num_warps=8
- GEMM using `dot_scaled` with M=N=256 tile size when num_stages=2 and
num_warps=8
- FA with num_stages=4
Only with using async_copy.

Author: jungpark-mlir

third_party/amd/backend/compiler.py

@@ -18,8 +18,9 @@ def get_min_dot_size(target: GPUTarget):
     return lambda lhs_type, rhs_type: (1, 1, 1)
 
 
-def is_pingpong_schedule_enabled(arch):
-    return (arch == "gfx942") if knobs.amd.use_block_pingpong is None else knobs.amd.use_block_pingpong
+def is_pingpong_schedule_enabled(arch, use_async_copy):
+    return (arch == "gfx942" or (arch == "gfx950" and use_async_copy is True)
+            ) if knobs.amd.use_block_pingpong is None else knobs.amd.use_block_pingpong
 
 
 def is_in_thread_transpose_enabled(arch):
@@ -233,7 +234,7 @@ def make_ttgir(mod, metadata, options):
         global_prefetch = knobs.amd.global_prefetch
         local_prefetch = knobs.amd.local_prefetch
         use_async_copy = knobs.amd.use_async_copy
-        use_block_pingpong = is_pingpong_schedule_enabled(options.arch)
+        use_block_pingpong = is_pingpong_schedule_enabled(options.arch, use_async_copy)
 
         amd.passes.ttgpuir.add_stream_pipeline(pm, options.num_stages, global_prefetch, local_prefetch, use_async_copy,
                                                use_block_pingpong)

third_party/amd/lib/TritonAMDGPUTransforms/BlockPingpong.cpp

@@ -764,6 +764,7 @@ Pingponger::transformTwoClusterWithLocalLoadAndAll(OpBuilder &builder,
       asyncWaitOp->erase();
     }
   }
+  assert(newAsyncWaitOp != nullptr);
 
   moveOpAndPredecessorsUpSameBlock(lLoadOps[0]);
   moveOpAndPredecessorsUpSameBlock(lLoadOps[1]);
@@ -918,21 +919,18 @@ void Pingponger::getDotPingponged() {
     auto aType = scaledDotOps[0].getA().getType();
     auto aShape = aType.getShape();
     auto elemWidth = aType.getElementTypeBitWidth();
-    int64_t tileSize = scaledDotShape[0] * scaledDotShape[1] * aShape[1];
 
-    // 256x256x256 (128xi8)
-    if (tileSize == 8388608 && aShape[0] == 256 && aShape[1] == 128 &&
+    // MxN = 256x256
+    if (scaledDotShape[0] == 256 && scaledDotShape[1] == 256 &&
         elemWidth == 8) {
-      kWidth = 16;
       if (transformTwoClusterWithAsyncAndAll(builder, scaledDotOps[0]->getLoc())
               .failed()) {
-        LDBG(
-            "Encountered failure when trying to execute the two-step ping pong "
-            "cluster transformation");
+        LDBG("Encountered failure when trying to execute the"
+             "TwoClusterWithAsyncAndAll transformation");
         return;
       }
+      addAsymmetricSyncToLoop(builder, loc);
     }
-    addAsymmetricSyncToLoop(builder, loc);
     return;
   } else if (scaledDotOps.size() == 1)
     return;
@@ -942,7 +940,6 @@ void Pingponger::getDotPingponged() {
   // Determine if we have a persistent GEMM. This will decide how we interpret
   // any memory operations that we find in conditionals.
   auto assumeNotTaken = isPersistentGemm(dotOps.size());
-
   // Compute tile size, kWidth, and mfma type.
   auto dotType = dotOps[0].getType();
   auto dotShape = dotType.getShape();
@@ -969,11 +966,11 @@ void Pingponger::getDotPingponged() {
       LDBG("Currently only support num_warp=8 for async PP");
       return;
     }
-    if (numStages > 2 && dotOps.size() == 1 && tileSize == mediumTile &&
-        aShape[1] == 32 && elemWidth == 16) {
+    if (numStages > 2 && dotOps.size() == 1 && dotShape[0] > 64 &&
+        dotShape[1] > 64 && (elemWidth == 16 || elemWidth == 8)) {
       if (transformTwoClusterWithLocalLoadAndAll(builder, loc).failed()) {
-        LDBG("Encountered failure when trying to execute the NS3 ping pong "
-             "cluster transformation");
+        LDBG("Encountered failure when trying to execute the "
+             "TwoClusterWithLocalLoadAndAll transformation");
         return;
       }
       addAsymmetricSyncToLoop(builder, loc);