implement launch

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

python/triton_tvm_ffi/driver.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 from ctypes import c_void_p
-from typing import Mapping, Tuple, Type
+from typing import List, Mapping, Tuple, Type
 from triton.backends.nvidia.driver import CudaDriver
-from .utils import get_device_properties, load_binary
+from .utils import get_device_properties, launch, load_binary
 
 
 class TVMFFIUtils(object):
@@ -49,10 +49,84 @@ class TVMFFIUtils(object):
         )
 
 
+class TVMLauncher(object):
+    def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
+        super().__init__(*args, **kwargs)
+
+        self.mask: List[bool] = [annotation != "constexpr" for annotation in src.signature.values()]
+        self.num_ctas = getattr(metadata, "num_ctas", 1)
+        self.launch = launch
+        self.global_scratch_size = metadata.global_scratch_size
+        self.global_scratch_align = metadata.global_scratch_align
+        self.profile_scratch_size = metadata.profile_scratch_size
+        self.profile_scratch_align = metadata.profile_scratch_align
+        self.launch_cooperative_grid = metadata.launch_cooperative_grid
+        self.launch_pdl = metadata.launch_pdl
+
+        # We assume the global Triton allocator is not enabled: `_allocator` must be a NullAllocator.
+        # This module depends on NullAllocator behavior; ensure no other code replaces the allocator.
+        from triton.runtime._allocation import _allocator, NullAllocator
+
+        assert isinstance(_allocator.get(), NullAllocator)
+
+    def __call__(
+        self,
+        gridX,
+        gridY,
+        gridZ,
+        stream,
+        function,
+        kernel_metadata,
+        launch_metadata,
+        launch_enter_hook,
+        launch_exit_hook,
+        *args,
+    ):
+        from triton.runtime import _allocation
+
+        def allocate_scratch(size, align, allocator):
+            if size > 0:
+                grid_size = gridX * gridY * gridZ
+                alloc_size = grid_size * self.num_ctas * size
+                alloc_fn = allocator.get()
+                return alloc_fn(alloc_size, align, stream)
+            return None
+
+        global_scratch = allocate_scratch(
+            self.global_scratch_size, self.global_scratch_align, _allocation._allocator
+        )
+        profile_scratch = allocate_scratch(
+            self.profile_scratch_size,
+            self.profile_scratch_align,
+            _allocation._profile_allocator,
+        )
+        assert not self.launch_cooperative_grid
+        assert not self.launch_pdl
+        assert len(self.mask) == len(args)
+        args = [arg for arg, m in zip(args, self.mask) if m]
+        return launch(
+            gridX,
+            gridY,
+            gridZ,
+            stream,
+            function,
+            kernel_metadata,
+            launch_metadata,
+            launch_enter_hook,
+            launch_exit_hook,
+            self.launch_cooperative_grid,
+            self.launch_pdl,
+            global_scratch,
+            profile_scratch,
+            *args,
+        )
+
+
 class TVMFFIDriver(CudaDriver):
     def __init__(self, *args, **kwargs) -> TVMFFIDriver:
         super().__init__(*args, **kwargs)
         self.utils: TVMFFIUtils = TVMFFIUtils()
+        self.launcher_cls: Type[TVMLauncher] = TVMLauncher
 
 
 del CudaDriver

python/triton_tvm_ffi/utils/_ffi_api.py

@@ -7,6 +7,7 @@ from tvm_ffi.libinfo import load_lib_module as _FFI_LOAD_LIB
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from collections.abc import Mapping
+    from typing import Any
 # isort: on
 # fmt: on
 # tvm-ffi-stubgen(end)
@@ -17,6 +18,7 @@ LIB = _FFI_LOAD_LIB("triton_tvm_ffi", "utils")
 _FFI_INIT_FUNC("triton_tvm_ffi.utils", __name__)
 if TYPE_CHECKING:
     def get_device_properties(_0: int, /) -> Mapping[str, int]: ...
+    def launch(*args: Any) -> Any: ...
     def load_binary(_0: str, _1: bytes, _2: int, _3: int, /) -> tuple[int, int, int, int, int]: ...
 # fmt: on
 # tvm-ffi-stubgen(end)
@@ -25,6 +27,7 @@ __all__ = [
     # tvm-ffi-stubgen(begin): __all__
     "LIB",
     "get_device_properties",
+    "launch",
     "load_binary",
     # tvm-ffi-stubgen(end)
 ]