[AMD][gfx1250] Add Assumptions and Fix Predicate in MXGEMM Kernel (#9285)

This PR
- added assumptions to loop boundary
- fix predicate to eliminate readfirstlane instrs
- moved `static_profile` and `composition` to a shared location

Author: knwng

third_party/amd/python/examples/gluon/gfx1250_utils.py

@@ -0,0 +1,37 @@
+import re
+
+
+def composition(cls):
+    """ A decorator lets aggregate type to directly access attributes from its aggregate member. """
+
+    def __getattr__(self, name):
+        if name in self.__dict__:
+            return object.__getattribute__(self, name)
+        for member in self.__dict__.values():
+            if getattr(member, "__triton_aggregate__", False) and not hasattr(member, name):
+                continue
+            return getattr(member, name)
+        raise AttributeError(f"{type(self).__name__} object has no attribute '{name}'")
+
+    cls.__getattr__ = __getattr__
+    return cls
+
+
+def static_profile(kernel):
+    amdgcn = kernel.asm['amdgcn']
+
+    sgpr_count = int(re.search(r'\.sgpr_count:\s+(\d+)', amdgcn).group(1))
+    sgpr_spill_count = int(re.search(r'\.sgpr_spill_count:\s+(\d+)', amdgcn).group(1))
+    vgpr_count = int(re.search(r'\.vgpr_count:\s+(\d+)', amdgcn).group(1))
+    vgpr_spill_count = int(re.search(r'\.vgpr_spill_count:\s+(\d+)', amdgcn).group(1))
+    scratch_size = int(re.search(r';\s+ScratchSize:\s+(\d+)', amdgcn).group(1))
+    code_len_in_byte = int(re.search(r';\s+codeLenInByte\s+=\s+(\d+)', amdgcn).group(1))
+    occupancy = int(re.search(r';\s+Occupancy:\s+(\d+)', amdgcn).group(1))
+
+    print(f"- sgpr_count: {sgpr_count}\n"
+          f"- sgpr_spill_count: {sgpr_spill_count}\n"
+          f"- vgpr_count: {vgpr_count}\n"
+          f"- vgpr_spill_count: {vgpr_spill_count}\n"
+          f"- scratch_size: {scratch_size}\n"
+          f"- code_len_in_byte: {code_len_in_byte}\n"
+          f"- occupancy: {occupancy}\n")

third_party/amd/python/examples/gluon/mxfp_fa_gfx1250.py

@@ -27,27 +27,17 @@
 from triton.experimental.gluon.language.amd.gfx1250 import buffer_load, buffer_store
 from triton.experimental.gluon.language.amd.gfx1250 import async_copy as cp
 
+# Handle imports for both pytest (module context) and direct execution
+try:
+    from .gfx1250_utils import static_profile, composition
+except ImportError:
+    from gfx1250_utils import static_profile, composition
+
 # ===-----------------------------------------------------------------------===#
 # Kernel Utilities
 # ===-----------------------------------------------------------------------===#
 
 
-def composition(cls):
-    """ A decorator lets aggregate type to directly access attributes from its aggregate member. """
-
-    def __getattr__(self, name):
-        if name in self.__dict__:
-            return object.__getattribute__(self, name)
-        for member in self.__dict__.values():
-            if getattr(member, "__triton_aggregate__", False) and not hasattr(member, name):
-                continue
-            return getattr(member, name)
-        raise AttributeError(f"{type(self).__name__} object has no attribute '{name}'")
-
-    cls.__getattr__ = __getattr__
-    return cls
-
-
 @gluon.constexpr_function
 def get_padded_shared_layout(shape, transposed=False):
     """ Get a padded shared layout without back conflict for a given tensor shape. """
@@ -1840,26 +1830,6 @@ def create_global_scale(dtype: str):
     return scale, scale_ref
 
 
-def static_profile(kernel):
-    amdgcn = kernel.asm['amdgcn']
-
-    sgpr_count = int(re.search(r'\.sgpr_count:\s+(\d+)', amdgcn).group(1))
-    sgpr_spill_count = int(re.search(r'\.sgpr_spill_count:\s+(\d+)', amdgcn).group(1))
-    vgpr_count = int(re.search(r'\.vgpr_count:\s+(\d+)', amdgcn).group(1))
-    vgpr_spill_count = int(re.search(r'\.vgpr_spill_count:\s+(\d+)', amdgcn).group(1))
-    scratch_size = int(re.search(r';\s+ScratchSize:\s+(\d+)', amdgcn).group(1))
-    code_len_in_byte = int(re.search(r';\s+codeLenInByte\s+=\s+(\d+)', amdgcn).group(1))
-    occupancy = int(re.search(r';\s+Occupancy:\s+(\d+)', amdgcn).group(1))
-
-    print(f"- sgpr_count: {sgpr_count}\n"
-          f"- sgpr_spill_count: {sgpr_spill_count}\n"
-          f"- vgpr_count: {vgpr_count}\n"
-          f"- vgpr_spill_count: {vgpr_spill_count}\n"
-          f"- scratch_size: {scratch_size}\n"
-          f"- code_len_in_byte: {code_len_in_byte}\n"
-          f"- occupancy: {occupancy}\n")
-
-
 def get_source_mapping(block_scaling, subtile, pipelined, amdgcn):
     """
     Create a mapping from amdgcn assembly to source code lines:

third_party/amd/python/examples/gluon/mxfp_gemm_gfx1250.py

@@ -4,7 +4,6 @@
 # Needed for internal dev flow for now; will remove later
 hip.hip.hipInit(0)
 
-import re
 import torch
 import pytest
 import triton
@@ -15,25 +14,11 @@
 from triton.language.core import _aggregate as aggregate
 from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
 
-
-def static_profile(kernel):
-    amdgcn = kernel.asm['amdgcn']
-
-    sgpr_count = int(re.search(r'\.sgpr_count:\s+(\d+)', amdgcn).group(1))
-    sgpr_spill_count = int(re.search(r'\.sgpr_spill_count:\s+(\d+)', amdgcn).group(1))
-    vgpr_count = int(re.search(r'\.vgpr_count:\s+(\d+)', amdgcn).group(1))
-    vgpr_spill_count = int(re.search(r'\.vgpr_spill_count:\s+(\d+)', amdgcn).group(1))
-    scratch_size = int(re.search(r';\s+ScratchSize:\s+(\d+)', amdgcn).group(1))
-    code_len_in_byte = int(re.search(r';\s+codeLenInByte\s+=\s+(\d+)', amdgcn).group(1))
-    occupancy = int(re.search(r';\s+Occupancy:\s+(\d+)', amdgcn).group(1))
-
-    print(f"- sgpr_count: {sgpr_count}\n"
-          f"- sgpr_spill_count: {sgpr_spill_count}\n"
-          f"- vgpr_count: {vgpr_count}\n"
-          f"- vgpr_spill_count: {vgpr_spill_count}\n"
-          f"- scratch_size: {scratch_size}\n"
-          f"- code_len_in_byte: {code_len_in_byte}\n"
-          f"- occupancy: {occupancy}\n")
+# Handle imports for both pytest (module context) and direct execution
+try:
+    from .gfx1250_utils import static_profile
+except ImportError:
+    from gfx1250_utils import static_profile
 
 
 @gluon.constexpr_function
@@ -78,7 +63,6 @@ class MXFPGEMMConfig:
     BLOCK_M_PRESHUFFLED: gl.constexpr
     BLOCK_N_PRESHUFFLED: gl.constexpr
     BLOCK_K_SCALE_PRESHUFFLED: gl.constexpr
-    tiles_per_warp: gl.constexpr
     SCALE_BLOCK: gl.constexpr
     ASYNC_COPY_SCALE: gl.constexpr
 
@@ -116,8 +100,6 @@ def __init__(self, BLOCK_M, BLOCK_N, BLOCK_K, DTYPE_A, DTYPE_B, SCALE_BLOCK, NUM
             reg_bases: gl.constexpr = []
             warp_bases: gl.constexpr = [[0, 1], [1, 0]]
 
-        self.tiles_per_warp = gl.constexpr([2, 2] if SCALE_PRESHUFFLE else [1, 1])
-
         self.BLOCK_M_PRESHUFFLED = gl.constexpr(BLOCK_M // self.PRESHUFFLE_FACTOR)
         self.BLOCK_N_PRESHUFFLED = gl.constexpr(BLOCK_N // self.PRESHUFFLE_FACTOR)
         self.BLOCK_K_SCALE_PRESHUFFLED = gl.constexpr(BLOCK_K_SCALE * self.PRESHUFFLE_FACTOR)
@@ -280,7 +262,7 @@ def initialize(cfg: MXFPGEMMConfig, a_desc, b_desc, a_scale_desc, b_scale_desc,
                                         a_scale_desc, b_scale_desc, c_ptr, c_offs, c_mask)
 
     @gluon.jit
-    def issue_loads(self, load_idx, pred=True):
+    def issue_loads(self, load_idx, pred=1):
         cfg = self.cfg
         NUM_SUBTILES_K = cfg.NUM_SUBTILES[2]
         BLOCK_K_PACKED_A: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A // NUM_SUBTILES_K
@@ -359,9 +341,12 @@ def pipeline(self, K):
 
         accumulator = gl.zeros((cfg.BLOCK_M, cfg.BLOCK_N), dtype=gl.float32, layout=self.cfg.acc_layout)
         loop_ub = gl.cdiv(K, cfg.BLOCK_K)
+        gl.assume(loop_ub > 0)
         epilogue_lb = loop_ub - (cfg.NUM_BUFFERS - 1)
         for i in range(0, loop_ub):
-            load_idx = self.issue_loads(load_idx, pred=(i < epilogue_lb))
+            pred = i - epilogue_lb
+            pred = (pred >> 31) & 1
+            load_idx = self.issue_loads(load_idx, pred=pred)
 
             gl.amd.gfx1250.tdm.async_wait((cfg.NUM_BUFFERS - 1) * self.cfg.NUM_LOADS_IN_BATCH)
 
@@ -554,7 +539,7 @@ def issue_local_load_b(self, wmma_idx, b_buffer, b_scale_buffer):
         return b, scale_b
 
     @gluon.jit
-    def issue_load_a(self, load_idx, a_buffer, a_scale_buffer, pred=True):
+    def issue_load_a(self, load_idx, a_buffer, a_scale_buffer, pred=1):
         cfg = self.cfg
         NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
         BLOCK_K: gl.constexpr = cfg.BLOCK_K // cfg.DIV_FACTOR_A // NUM_SUBTILES_K
@@ -574,7 +559,7 @@ def issue_load_a(self, load_idx, a_buffer, a_scale_buffer, pred=True):
         return load_idx + 1
 
     @gluon.jit
-    def issue_load_b(self, load_idx, b_buffer, b_scale_buffer, pred=True):
+    def issue_load_b(self, load_idx, b_buffer, b_scale_buffer, pred=1):
         cfg = self.cfg
         NUM_SUBTILES_N: gl.constexpr = cfg.NUM_SUBTILES[1]
         NUM_SUBTILES_K: gl.constexpr = cfg.NUM_SUBTILES[2]
@@ -642,9 +627,11 @@ def pipeline(self, K):
                       layout=cfg.acc_layout)
 
         loop_ub = gl.cdiv(K, cfg.BLOCK_K)
+        gl.assume(loop_ub > 0)
         epilogue_lb = loop_ub - (cfg.NUM_BUFFERS - 1)
         for i in range(0, loop_ub):
-            pred = (i < epilogue_lb)
+            pred = i - epilogue_lb
+            pred = (pred >> 31) & 1
 
             # iter i + 1
             load_a_idx = self.issue_load_a(load_a_idx, self.a_buffer0, self.a_scale_buffer0, pred=pred)