Set SM and kpack according to arch and early return if no streamK tiles

Author: zhanglx13

python/perf-kernels/streamk/03-matrix-multiplication-stream-k.py

@@ -9,7 +9,25 @@
 torch.manual_seed(123)
 random.seed(123)
 
-total_sm = 304
+
+def is_hip_cdna3():
+    target = triton.runtime.driver.active.get_current_target()
+    return target.backend == 'hip' and target.arch == 'gfx942'
+
+
+def is_hip_cdna4():
+    target = triton.runtime.driver.active.get_current_target()
+    return target.backend == 'hip' and target.arch == 'gfx950'
+
+
+if is_hip_cdna3():
+    total_sm = 304
+elif is_hip_cdna4():
+    total_sm = 256
+else:
+    print("Unknown target")
+    exit(0)
+
 print(f"total SMs: {total_sm}")
 
 
@@ -171,7 +189,14 @@ def forward(ctx, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, bias: torch.
 num_warps = 8
 waves_per_eu = 0
 mfmaInstrSize = 16
-kpack = 2
+
+if is_hip_cdna3():
+    kpack = 2
+elif is_hip_cdna4():
+    kpack = 1
+else:
+    print("Unknown target")
+    exit(0)
 
 ##for total_sm in range(1, 305):
 ##    print(f"{total_sm=}")

python/perf-kernels/streamk/streamk_kernel.py

@@ -105,6 +105,9 @@ def streamk_gemm(
         mask = (rm < M)[:, None] & (rn < N)[None, :]
         tl.store(C_, c, mask=mask)
 
+    if STREAMK_TILES == 0:
+        return
+
     tl.assume(pid >= 0)
     total_streamk_iters = STREAMK_TILES * iters_per_tile
     streamk_iters_pcu = total_streamk_iters // NUM_SMS