num_worker_warps hint

Signed-off-by: Boyan Li <boyanl@nvidia.com>

Author: blinxt

src/cuda/tile/_bytecode/attribute.py

@@ -62,13 +62,17 @@ def encode_dictionary_untagged(items: Sequence[tuple[str, TaggedAttribute]],
 class EntryHints:
     num_cta_in_cga: Optional[int] = None
     occupancy: Optional[int] = None
+    num_worker_warps_per_cta: Optional[int] = None
 
     def as_dictionary(self) -> Dictionary:
         items = []
         if self.num_cta_in_cga is not None:
             items.append(("num_cta_in_cga", Integer.create_i32(self.num_cta_in_cga)))
         if self.occupancy is not None:
             items.append(("occupancy", Integer.create_i32(self.occupancy)))
+        if self.num_worker_warps_per_cta is not None:
+            items.append(("num_worker_warps_per_cta",
+                         Integer.create_i32(self.num_worker_warps_per_cta)))
         return Dictionary(items)
 
 

src/cuda/tile/_compiler_options.py

@@ -16,6 +16,7 @@ class CompilerOptions:
     num_ctas: None | int | ByTarget[int] = None
     occupancy: None | int | ByTarget[int] = None
     opt_level: int | ByTarget[int] = 3
+    num_worker_warps: None | int | ByTarget[int] = None
 
     def __post_init__(self):
         for field in dataclasses.fields(self):
@@ -59,3 +60,10 @@ def _validate_opt_level(opt_level: None | int):
     if opt_level is not None:
         if opt_level < 0 or opt_level > 3:
             raise ValueError(f'opt_level should be [0, 3], got {opt_level}')
+
+
+def _validate_num_worker_warps(num_worker_warps: None | int):
+    if num_worker_warps is not None:
+        if num_worker_warps not in (4, 8):
+            raise ValueError(f'num_worker_warps should be either 4 or 8,'
+                             f' got {num_worker_warps}')

src/cuda/tile/_execution.py

@@ -75,6 +75,16 @@ class kernel(TileDispatcher):
             Default: None (auto).
         occupancy: Expected number of active CTAs per SM, [1, 32]. Default: None (auto).
         opt_level: Optimization level [0, 3], default 3.
+        num_worker_warps: Number of warps in the CUDA core warp groups in a
+            warp-specialized kernel. The compiler may add warps
+            (e.g., for asynchronous memory transfers) that are not counted here.
+            This value does not represent the total warp count.
+            It's worth tuning when a warp-specialized kernel has high register pressure
+            that other approaches cannot resolve.
+            Normalization-style kernels with large tiles are the canonical cases.
+            Must be either 4 or 8.
+            Default: None (auto).
+            Since CTK 13.3. Ignored with a warning otherwise.
 
     Target-specific values for the compiler options above can be provided
     using a :py:class:`ByTarget` object.
@@ -92,7 +102,8 @@ def __init__(self,
                  /, *,
                  num_ctas: None | int | ByTarget[int] = None,
                  occupancy: None | int | ByTarget[int] = None,
-                 opt_level: None | int | ByTarget[int] = 3):
+                 opt_level: None | int | ByTarget[int] = 3,
+                 num_worker_warps: None | int | ByTarget[int] = None):
         if not isinstance(function, FunctionType):
             raise TypeError("`kernel` decorator must be applied to a Python function")
 
@@ -103,7 +114,8 @@ def __init__(self,
         compiler_options = CompilerOptions(
             num_ctas=num_ctas,
             occupancy=occupancy,
-            opt_level=opt_level
+            opt_level=opt_level,
+            num_worker_warps=num_worker_warps
         )
         super().__init__(ann_func.constant_parameter_mask, ann_func.int64_index_parameter_mask,
                          ann_func.int64_parameter_mask)

src/cuda/tile/_ir2bytecode.py

@@ -6,9 +6,11 @@
 import os
 from contextlib import contextmanager
 from typing import Dict, Tuple, Any, Optional
+import warnings
 
 from cuda.tile import _datatype as datatype
 from cuda.tile._bytecode.attribute import make_load_store_hints
+from cuda.tile._bytecode.version import BytecodeVersion
 from cuda.tile._datatype import get_signedness
 from cuda.tile import DType
 import cuda.tile._bytecode as bc
@@ -419,8 +421,18 @@ def generate_bytecode_for_kernel(func_body: Block,
                                  writer: bc.BytecodeWriter,
                                  anonymize_debug_attr: bool):
     target_options = compiler_options.specialize_for_target(sm_arch)
+    num_worker_warps = target_options.num_worker_warps
+    if num_worker_warps is not None and writer.version < BytecodeVersion.V_13_3:
+        warnings.warn(
+            f"num_worker_warps requires tileiras "
+            f"{BytecodeVersion.V_13_3.as_string()} or later; ignoring "
+            f"(current version is {writer.version.as_string()}).",
+        )
+        num_worker_warps = None
+
     entry_hints = bc.EntryHints(num_cta_in_cga=target_options.num_ctas,
-                                occupancy=target_options.occupancy)
+                                occupancy=target_options.occupancy,
+                                num_worker_warps_per_cta=num_worker_warps)
 
     param_type_ids = [typeid(writer.type_table, p.get_type()) for p in func_body.params]
     debug_attr_map = DebugAttrMap(writer.debug_attr_table, symbol, anonymize_debug_attr)

test/test_bytecode_version_compat.py

@@ -13,8 +13,7 @@
 from cuda.tile.compilation import CallingConvention, KernelSignature
 
 
-def compile_with_version(pyfunc, args, version: str):
-    kernel = ct.kernel(pyfunc)
+def compile_with_version(kernel, args, version: str):
     cconv = CallingConvention.cutile_python_v1()
     sig = KernelSignature.from_kernel_args(kernel, args, cconv)
     ct.compilation.export_kernel(kernel, [sig], output_file=BytesIO(),
@@ -27,6 +26,7 @@ def tensor(dtype=torch.float32):
 
 
 def test_atan2_requires_13_2():
+    @ct.kernel
     def kernel(x, y, z):
         tx = ct.load(x, 0, shape=64)
         ty = ct.load(y, 0, shape=64)
@@ -37,6 +37,7 @@ def kernel(x, y, z):
 
 
 def test_tanh_rounding_mode_requires_13_2():
+    @ct.kernel
     def kernel(x, y):
         tx = ct.load(x, 0, shape=64)
         ct.store(y, 0, tile=ct.tanh(tx, rounding_mode=RoundingMode.APPROX))
@@ -47,6 +48,7 @@ def kernel(x, y):
 
 
 def test_tanh_without_rounding_mode_works_on_13_1():
+    @ct.kernel
     def kernel(x, y):
         tx = ct.load(x, 0, shape=64)
         ct.store(y, 0, tile=ct.tanh(tx))
@@ -56,6 +58,7 @@ def kernel(x, y):
 
 
 def test_exp_rounding_mode_requires_13_3():
+    @ct.kernel
     def kernel(x, y):
         tx = ct.load(x, 0, shape=64)
         ct.store(y, 0, tile=ct.exp(tx, rounding_mode=RoundingMode.APPROX))
@@ -66,9 +69,21 @@ def kernel(x, y):
 
 
 def test_exp_without_rounding_mode_works_on_13_1():
+    @ct.kernel
     def kernel(x, y):
         tx = ct.load(x, 0, shape=64)
         ct.store(y, 0, tile=ct.exp(tx))
 
     # Should not raise version error
     compile_with_version(kernel, (tensor(), tensor()), "13.1")
+
+
+def test_num_worker_warps_warns_below_13_3():
+    @ct.kernel(num_worker_warps=8)
+    def kernel(x, y):
+        tx = ct.load(x, 0, shape=64)
+        ct.store(y, 0, tile=tx)
+
+    with pytest.warns(UserWarning,
+                      match=r"num_worker_warps requires tileiras 13\.3"):
+        compile_with_version(kernel, (tensor(), tensor()), "13.1")

test/test_compiler_options.py

@@ -11,3 +11,22 @@ def test_invalid_target_name():
     err = r"Invalid GPU architecture name: sm100, expected sm_<major><minor>"
     with pytest.raises(ValueError, match=err):
         ct.ByTarget(sm100=4)
+
+
+def _dummy():
+    pass
+
+
+@pytest.mark.parametrize("value", [None, 4, 8])
+def test_num_worker_warps_accepts_valid(value):
+    ct.kernel(_dummy, num_worker_warps=value)
+
+
+@pytest.mark.parametrize("value", [3, 7, 10])
+def test_num_worker_warps_rejects_invalid(value):
+    with pytest.raises(ValueError, match="num_worker_warps should be either 4 or 8"):
+        ct.kernel(_dummy, num_worker_warps=value)
+
+
+def test_num_worker_warps_accepts_by_target():
+    ct.kernel(_dummy, num_worker_warps=ct.ByTarget(sm_100=8, default=4))