Separated optimal DPAS layout calculation logic, exposed that for python, enabled all gemm benchmark cases for gluon

Author: mieshkiwrk

benchmarks/triton_kernels_benchmark/gemm_benchmark.py

@@ -21,6 +21,8 @@
 from triton_kernels_benchmark import xetla_kernel
 from triton_kernels_benchmark import cutlass_kernel
 
+from utils.dpas_layout_analyzer import calculate_optimal_warps_per_cta, calculate_optimal_rep_clusters
+
 
 def get_matmul_autotune_configs() -> List[triton.Config]:
     configs = [
@@ -178,32 +180,20 @@ def get_gluon_matmul_autotune_configs(base_configs_fn: Callable) -> List[triton.
             # Append additional meta parameters needed for gluon kernel
             # To determine prefetch distance and DPAS layout
             {**config.kwargs, 'NUM_STAGES': config.num_stages, 'NUM_WARPS': config.num_warps},
-            num_stages=config.num_stages,
-            num_warps=config.num_warps
-        )
-        for config in base_configs
+            num_stages=config.num_stages, num_warps=config.num_warps) for config in base_configs
     ]
 
 
 @gluon.constexpr_function
-def get_dpas_layout(num_warps: ttgl.constexpr) -> ttgl.constexpr:
-    # TODO: return same DPAS layout as calculated by passes for triton
-    warps_per_cta = [2, 2]
-    if num_warps == 16:
-        warps_per_cta = [4, 4]
-    if num_warps == 32:
-        warps_per_cta = [4, 8]
-    elif num_warps == 64:
-        warps_per_cta = [8, 8]
+def get_dpas_layout(num_warps: ttgl.constexpr, m_shape: ttgl.constexpr, n_shape: ttgl.constexpr,
+                    k_shape: ttgl.constexpr) -> ttgl.constexpr:
+    threads_per_warp = 16
+    warps_per_cta = calculate_optimal_warps_per_cta(num_warps, m_shape, n_shape)
+
     return IntelDPASLayout(
-        repeatCount=8,
-        systolic_depth=8,
-        execution_size=16,
-        ops_per_chan=2,
-        warps_per_cta=warps_per_cta,
-        rep_cluster=[4, 2],
-        threads_per_warp=16
-   )
+        repeatCount=8, systolic_depth=8, execution_size=16, ops_per_chan=2, warps_per_cta=warps_per_cta,
+        rep_cluster=calculate_optimal_rep_clusters(m_shape, n_shape, k_shape, threads_per_warp,
+                                                   warps_per_cta), threads_per_warp=threads_per_warp)
 
 
 @triton.autotune(
@@ -217,16 +207,14 @@ def gluon_matmul_kernel_dpas_tensor_desc(
         # Matrix dimensions
         M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr,
         # Stride variables
-        stride_am: ttgl.constexpr, stride_ak: ttgl.constexpr,
-        stride_bk: ttgl.constexpr, stride_bn: ttgl.constexpr,
+        stride_am: ttgl.constexpr, stride_ak: ttgl.constexpr, stride_bk: ttgl.constexpr, stride_bn: ttgl.constexpr,
         stride_cm: ttgl.constexpr, stride_cn: ttgl.constexpr,
         # Meta parameters
         BLOCK_SIZE_M: ttgl.constexpr, BLOCK_SIZE_N: ttgl.constexpr, BLOCK_SIZE_K: ttgl.constexpr,
         GROUP_SIZE_M: ttgl.constexpr,
         # Gluon meta parameters
         NUM_STAGES: ttgl.constexpr, NUM_WARPS: ttgl.constexpr):
-    layout: ttgl.constexpr = get_dpas_layout(NUM_WARPS)
-
+    layout: ttgl.constexpr = get_dpas_layout(NUM_WARPS, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K)
 
     lhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=0, k_width=1)
     rhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=1, k_width=2)
@@ -241,19 +229,19 @@ def gluon_matmul_kernel_dpas_tensor_desc(
     pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
     pid_n = (pid % num_pid_in_group) // group_size_m
 
-    a_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(a_ptr, (M, K), (stride_am, stride_ak), (BLOCK_SIZE_M, BLOCK_SIZE_K),
-                                                      lhs_layout)
-    b_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(b_ptr, (K, N), (stride_bk, stride_bn), (BLOCK_SIZE_K, BLOCK_SIZE_N),
-                                                      rhs_layout)
-    c_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(c_ptr, (M, N), (stride_cm, stride_cn), (BLOCK_SIZE_M, BLOCK_SIZE_N), layout)
+    a_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(a_ptr, (M, K), (stride_am, stride_ak),
+                                                      (BLOCK_SIZE_M, BLOCK_SIZE_K), lhs_layout)
+    b_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(b_ptr, (K, N), (stride_bk, stride_bn),
+                                                      (BLOCK_SIZE_K, BLOCK_SIZE_N), rhs_layout)
+    c_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(c_ptr, (M, N), (stride_cm, stride_cn),
+                                                      (BLOCK_SIZE_M, BLOCK_SIZE_N), layout)
 
     # Clear accumulator
     zero_tensor = ttgl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=ttgl.float32, layout=layout)
     c_desc.store_2d([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], zero_tensor)
 
     accumulator = c_desc.load_2d([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N])
 
-
     # Prefetch first blocks for A and B matrices (pre-loop prefetches)
     for i in range(NUM_STAGES):
         if i * BLOCK_SIZE_K < K:
@@ -286,15 +274,15 @@ def gluon_matmul_kernel_dpas_tensor_desc_batched(
         # Matrix dimensions
         B: ttgl.constexpr, M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr,
         # Stride variables
-        stride_az: ttgl.constexpr, stride_am: ttgl.constexpr, stride_ak: ttgl.constexpr,
-        stride_bz: ttgl.constexpr, stride_bk: ttgl.constexpr, stride_bn: ttgl.constexpr,
-        stride_cz: ttgl.constexpr, stride_cm: ttgl.constexpr, stride_cn: ttgl.constexpr,
+        stride_az: ttgl.constexpr, stride_am: ttgl.constexpr, stride_ak: ttgl.constexpr, stride_bz: ttgl.constexpr,
+        stride_bk: ttgl.constexpr, stride_bn: ttgl.constexpr, stride_cz: ttgl.constexpr, stride_cm: ttgl.constexpr,
+        stride_cn: ttgl.constexpr,
         # Meta parameters
         BLOCK_SIZE_M: ttgl.constexpr, BLOCK_SIZE_N: ttgl.constexpr, BLOCK_SIZE_K: ttgl.constexpr,
         GROUP_SIZE_M: ttgl.constexpr,
         # Gluon meta parameters
         NUM_STAGES: ttgl.constexpr, NUM_WARPS: ttgl.constexpr):
-    layout: ttgl.constexpr = get_dpas_layout(NUM_WARPS)
+    layout: ttgl.constexpr = get_dpas_layout(NUM_WARPS, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K)
 
     lhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=0, k_width=1)
     rhs_layout: ttgl.constexpr = ttgl.DotOperandLayout(parent=layout, operand_index=1, k_width=2)
@@ -315,18 +303,12 @@ def gluon_matmul_kernel_dpas_tensor_desc_batched(
     offset_b = bid.to(ttgl.int64) * stride_bz
     offset_c = bid.to(ttgl.int64) * stride_cz
 
-    a_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(
-        a_ptr + offset_a, (M, K), (stride_am, stride_ak),
-        (BLOCK_SIZE_M, BLOCK_SIZE_K), lhs_layout
-    )
-    b_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(
-        b_ptr + offset_b, (K, N), (stride_bk, stride_bn),
-        (BLOCK_SIZE_K, BLOCK_SIZE_N), rhs_layout
-    )
-    c_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(
-        c_ptr + offset_c, (M, N), (stride_cm, stride_cn),
-        (BLOCK_SIZE_M, BLOCK_SIZE_N), layout
-    )
+    a_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(a_ptr + offset_a, (M, K), (stride_am, stride_ak),
+                                                      (BLOCK_SIZE_M, BLOCK_SIZE_K), lhs_layout)
+    b_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(b_ptr + offset_b, (K, N), (stride_bk, stride_bn),
+                                                      (BLOCK_SIZE_K, BLOCK_SIZE_N), rhs_layout)
+    c_desc = ttgl.intel.xpu.xe.make_tensor_descriptor(c_ptr + offset_c, (M, N), (stride_cm, stride_cn),
+                                                      (BLOCK_SIZE_M, BLOCK_SIZE_N), layout)
 
     # Clear accumulator
     zero_tensor = ttgl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=ttgl.float32, layout=layout)
@@ -386,20 +368,12 @@ def matmul(
             triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
             B,
         )
-        matmul_kernel_batched[grid](
-            a, b, c,  #
-            B, M, N, K,  #
-            a.stride(0), a.stride(a_major), a.stride(a_minor),  #
-            b.stride(0), b.stride(b_minor), b.stride(b_major),  #
-            c.stride(0), c.stride(1), c.stride(2))
+        matmul_kernel_batched[grid](a, b, c, B, M, N, K, a.stride(0), a.stride(a_major), a.stride(a_minor), b.stride(0),
+                                    b.stride(b_minor), b.stride(b_major), c.stride(0), c.stride(1), c.stride(2))
     elif len(a.shape) == 2 and len(b.shape) == 2:
         grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
-        matmul_kernel[grid](
-            a, b, c,  #
-            M, N, K,  #
-            a.stride(a_major), a.stride(a_minor),  #
-            b.stride(b_minor), b.stride(b_major),  #
-            c.stride(0), c.stride(1))
+        matmul_kernel[grid](a, b, c, M, N, K, a.stride(a_major), a.stride(a_minor), b.stride(b_minor),
+                            b.stride(b_major), c.stride(0), c.stride(1))
     else:
         assert False, 'Input matrixs dimensions mismatch'
     return c
@@ -459,7 +433,7 @@ def get_shapes(B, M, N, K, transpose_a, transpose_b):
     [4, 32768, 4096, 128],
     [32, 4096, 128, 4096],
     [4096, 8, 128, 16384],
-    # [4096, 8, 16384, 128], # TODO: mismatches for gluon
+    [4096, 8, 16384, 128],
 ]
 
 DEVICE_NAME = torch.xpu.get_device_name()
@@ -498,13 +472,13 @@ def get_benchmark(
     supported_providers = {
         'gluon': 'Gluon',
         'triton': 'Triton',
-        'onednn': 'OneDNN',
+        #'onednn': 'OneDNN',
     }
     # use_cutlass
-    if not (transpose_a or transpose_b):
-        if torch.xpu.get_device_name() != 'Intel(R) Arc(TM) Graphics':
-            # FIXME: enable cutlass on LNL
-            supported_providers['cutlass'] = 'CUTLASS'
+    # if not (transpose_a or transpose_b):
+    #     if torch.xpu.get_device_name() != 'Intel(R) Arc(TM) Graphics':
+    #         # FIXME: enable cutlass on LNL
+    #         supported_providers['cutlass'] = 'CUTLASS'
     providers = benchmark_suite.filter_providers(supported_providers, providers_filter)
 
     # Benchmark Performance

benchmarks/triton_kernels_benchmark/utils/dpas_layout_analyzer.py

@@ -0,0 +1,50 @@
+from functools import wraps
+from triton._C.libtriton import intel
+
+from triton.experimental.gluon.language.intel.xpu.xe import get_dpas_capabilities
+from triton.language.core import TRITON_BUILTIN
+
+
+def allow_in_kernel(fn):
+    """Mark a function as a builtin."""
+    assert callable(fn)
+
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        return fn(*args, **kwargs)
+
+    setattr(wrapper, TRITON_BUILTIN, True)
+
+    return wrapper
+
+
+@allow_in_kernel
+def calculate_optimal_warps_per_cta(num_warps, m_shape, n_shape):
+    ret_shape = [m_shape, n_shape]
+    dpas_cap = get_dpas_capabilities()
+    return intel.calculate_warps_per_tile(capRepeatCount=dpas_cap['repeatCount'],
+                                          capExecutionSize=dpas_cap['executionSize'], shape=ret_shape,
+                                          numWarps=num_warps)
+
+
+@allow_in_kernel
+def calculate_optimal_rep_clusters(block_m, block_n, block_k, threads_per_warp, warps_per_cta):
+    dtype_bitwidth = 16  # bf16  TODO: auto detect
+    is_fp8 = dtype_bitwidth == 8
+    dpas_cap = get_dpas_capabilities()
+    cap_repeat_count = dpas_cap['repeatCount']
+    cap_systolic_depth = dpas_cap['systolicDepth']
+    cap_execution_size = dpas_cap['executionSize']
+    ops_per_chan = int(dpas_cap['opsChanBitWidths'] / dtype_bitwidth)
+
+    ret_shape = [block_m, block_n]
+    a_shape = [block_m, block_k]
+    b_shape = [block_k, block_n]
+
+    rep_cluster = intel.calculate_rep_cluster(cap_repeat_count=cap_repeat_count, cap_systolic_depth=cap_systolic_depth,
+                                              cap_execution_size=cap_execution_size, ops_per_chan=ops_per_chan,
+                                              ret_shape=ret_shape, threads_per_warp=threads_per_warp,
+                                              a_bitwidth=dtype_bitwidth, is_fp8=is_fp8, a_shape=a_shape,
+                                              b_shape=b_shape, warps_per_tile=warps_per_cta)
+
+    return rep_cluster

python/triton/experimental/gluon/language/intel/xpu/xe.py

@@ -2,19 +2,42 @@
 
 from typing import List, Tuple, Sequence
 from dataclasses import dataclass
+from functools import cache
 
 import triton.experimental.gluon.language._core as ttgl
 from triton.experimental.gluon.language._layouts import DotOperandLayout
 from triton.experimental.gluon.language.intel._layouts import IntelDPASLayout
 from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr
 from triton.language.core import ir, constexpr, tensor_descriptor_base, block_type, tensor, tuple
 
-# load_tensor_descriptor = builtin(tl_core.load_tensor_descriptor)
-# store_tensor_descriptor = builtin(tl_core.store_tensor_descriptor)
-
 __all__ = ["make_tensor_descriptor", "dot_fma"]
 
 
+@cache
+def get_dpas_capabilities():
+    from triton.backends.intel.driver import XPUDriver
+
+    driver = XPUDriver()
+    target = driver.get_current_target()
+    properties = target.arch
+
+    # like annotate_module in passes
+    dpas_cap = {
+        "systolicDepth": 8,
+        "repeatCount": 8,
+        "executionSize": min(properties.get("sub_group_sizes", [16])),
+        "opsChanBitWidths": 32,
+        "has_subgroup_2d_block_io": properties.get("has_subgroup_2d_block_io", False),
+    }
+
+    return dpas_cap
+
+
+def is_2d_block_supported():
+    capabilities = get_dpas_capabilities()
+    return capabilities["has_subgroup_2d_block_io"]
+
+
 class tensor_descriptor(tensor_descriptor_base):
     """A descriptor representing a tensor in global memory."""
 
@@ -36,68 +59,55 @@ def _flatten_ir(self, handles: List[ir.value]) -> None:
         self.shape._flatten_ir(handles)
         self.strides._flatten_ir(handles)
 
+    def mark_2d_block_attribute(self, op, order, _semantic):
+        if order not in ('row_major', 'column_major'):
+            raise ValueError("Only row_major/column_major order is supported for 2d block")
+
+        attr = _semantic.builder.get_string_attr(order)
+        op.set_attr("ttig.block_io", attr)
+
     @builtin
     def load(self, offsets: Sequence[constexpr | tensor], _semantic=None) -> tensor:
         return _semantic.descriptor_load(self, offsets, "", "")
 
-    def load_2d(self, offsets: Sequence[constexpr | tensor], is_2d_block=False, _semantic=None) -> tensor:
-        # TODO: MaterializeBlockPointers.cpp
-        # Add 2d_block_io parameter + validation to set proper attribute
-        # Validation: (?)
-        #   > 2 dims
-        #   > stride 16 bytes aligned
-        #   and others
+    @builtin
+    def load_2d(self, offsets: Sequence[constexpr | tensor], order: str = "row_major", _semantic=None) -> tensor:
+        if not is_2d_block_supported():
+            raise ValueError("2d block functionality is not supported for this hardware")
 
         op = _semantic.descriptor_load(self, offsets, "", "")
-
-        # TODO: proper handling like below test example
-        # Option to set row/column major and other params
-        attr = _semantic.builder.get_string_attr("row_major")
-        op.handle.set_attr("ttig.block_io", attr)
-
+        self.mark_2d_block_attribute(op.handle, order, _semantic)
         return op
 
     @builtin
     def store(self, offsets: Sequence[constexpr | tensor], value: tensor, _semantic=None) -> tensor:
         return _semantic.descriptor_store(self, value, offsets)
 
     @builtin
-    def store_2d(self, offsets: Sequence[constexpr | tensor], value: tensor, _semantic=None) -> tensor:
-        op = _semantic.descriptor_store(self, value, offsets)
-
-        attr = _semantic.builder.get_string_attr("row_major")
-        op.handle.set_attr("ttig.block_io", attr)
+    def store_2d(self, offsets: Sequence[constexpr | tensor], value: tensor, order: str = "row_major",
+                 _semantic=None) -> tensor:
+        if not is_2d_block_supported():
+            raise ValueError("2d block functionality is not supported for this hardware")
 
+        op = _semantic.descriptor_store(self, value, offsets)
+        self.mark_2d_block_attribute(op.handle, order, _semantic)
         return op
 
     @builtin
-    def prefetch(self, offsets: Sequence[constexpr | tensor], mask=None, cache=None, evict=None, is_volatile=False, _semantic=None):
+    def prefetch(self, offsets: Sequence[constexpr | tensor], _semantic=None):
         ptr_handle = self.handle
         offsets_handles = [offset.handle if hasattr(offset, 'handle') else offset for offset in offsets]
         return _semantic.builder.create_prefetch(ptr_handle, offsets_handles, False)
 
     @builtin
-    def prefetch_2d(self, offsets: Sequence[constexpr | tensor], mask=None, cache=None, evict=None, is_volatile=False, _semantic=None):
-        # TODO: handle other ttig.prefetch params
-        # ptr is just temporary, support for tensor descriptor is needed
-        # calculate offsets like tt.advance
-        # maybe add support for mask, seems optional
-        # also 2d block attr and others
-        #return _semantic.builder.create_prefetch(ptr.handle, False)
-        """
-        pyton/triton/language/semantic.py @ load:1077 (TritonSemantic)
-        cache_modifier: str, eviction_policy: str
-        cache = self._str_to_load_cache_modifier(cache_modifier)
-        eviction = self._str_to_eviction_policy(eviction_policy)
-        """
+    def prefetch_2d(self, offsets: Sequence[constexpr | tensor], order: str = "row_major", _semantic=None):
+        if not is_2d_block_supported():
+            raise ValueError("2d block functionality is not supported for this hardware")
 
         ptr_handle = self.handle
         offsets_handles = [offset.handle if hasattr(offset, 'handle') else offset for offset in offsets]
         op = _semantic.builder.create_prefetch(ptr_handle, offsets_handles, False)
-
-        attr = _semantic.builder.get_string_attr("row_major")
-        op.set_attr("ttig.block_io", attr)
-
+        self.mark_2d_block_attribute(op, order, _semantic)
         return op