put utils in C++ side

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

python/triton_tvm_ffi/_ffi_api.py

@@ -6,7 +6,6 @@ from tvm_ffi import Object as _ffi_Object, init_ffi_api as _FFI_INIT_FUNC, regis
 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 tvm_ffi import Object
     from typing import Any
 # isort: on
@@ -19,15 +18,11 @@ LIB = _FFI_LOAD_LIB("triton_tvm_ffi", "triton_tvm_ffi")
 # fmt: off
 _FFI_INIT_FUNC("triton_tvm_ffi", __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]: ...
     def string_to_type(_0: str, /) -> int | None: ...
     def type_to_string(_0: int, /) -> str: ...
 # fmt: on
 # tvm-ffi-stubgen(end)
 
-
 # tvm-ffi-stubgen(import-object): tvm_ffi.register_object;False;_FFI_REG_OBJ
 # tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object
 @_FFI_REG_OBJ("triton_tvm_ffi.TypedValue")
@@ -40,14 +35,10 @@ class TypedValue(_ffi_Object):
     # fmt: on
     # tvm-ffi-stubgen(end)
 
-
 __all__ = [
     # tvm-ffi-stubgen(begin): __all__
     "LIB",
     "TypedValue",
-    "get_device_properties",
-    "launch",
-    "load_binary",
     "string_to_type",
     "type_to_string",
     # tvm-ffi-stubgen(end)

python/triton_tvm_ffi/driver.py

@@ -1,67 +1,23 @@
 from __future__ import annotations
 
-from ctypes import c_void_p
-from typing import Any, List, Mapping, Optional, Tuple, Type
+from typing import Any, List, Optional, Type
 from triton.backends.nvidia.driver import CudaDriver
-from . import TypedValue, get_device_properties, launch, load_binary, string_to_type
-
-
-class TVMFFIUtils(object):
-    def __new__(cls: Type[TVMFFIUtils]) -> TVMFFIUtils:
-        if not hasattr(cls, "instance"):
-            cls.instance = super().__new__(cls)
-        return cls.instance
-
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-
-    def load_binary(
-        self, name: str, data: bytes, shared: int, device: int
-    ) -> Tuple[c_void_p, c_void_p, int, int, int]:
-        return load_binary(name, data, shared, device)
-
-    def get_device_properties(self, device_id: int) -> Mapping[str, int]:
-        return get_device_properties(device_id)
-
-    def cuOccupancyMaxActiveClusters(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"cuOccupancyMaxActiveClusters isn\'t expected to be invoked"'
-        )
-
-    def set_printf_fifo_size(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"set_printf_fifo_size" isn\'t expected to be invoked'
-        )
-
-    def fill_tma_descriptor(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"fill_tma_descriptor" hasn\'t been supported for Hopper'
-        )
-
-    def launch(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"launch" is introduced in triton after commit d2b3925410689155e0f6028e8554bba972989348, which is still not supported yed'
-        )
-
-    def build_signature_metadata(self, *args, **kwargs):
-        raise NotImplementedError(
-            '"launch" is introduced in triton after commit d2b3925410689155e0f6028e8554bba972989348, which is still not supported yed'
-        )
+from . import TypedValue, utils, string_to_type
 
 
 class TVMLauncher(object):
     def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
         super().__init__(*args, **kwargs)
 
-        self.signature = 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
+        self.signature: List[str] = src.signature.values()
+        self.num_ctas: int = getattr(metadata, "num_ctas", 1)
+        self.launch = utils.launch
+        self.global_scratch_size: int = metadata.global_scratch_size
+        self.global_scratch_align: int = metadata.global_scratch_align
+        self.profile_scratch_size: int = metadata.profile_scratch_size
+        self.profile_scratch_align: int = metadata.profile_scratch_align
+        self.launch_cooperative_grid: bool = metadata.launch_cooperative_grid
+        self.launch_pdl: bool = metadata.launch_pdl
 
     def __call__(
         self,
@@ -105,7 +61,7 @@ class TVMLauncher(object):
 
         args = [canonicalize(arg, sig) for arg, sig in zip(args, self.signature)]
 
-        return launch(
+        return self.launch(
             gridX,
             gridY,
             gridZ,
@@ -126,7 +82,7 @@ class TVMLauncher(object):
 class TVMFFIDriver(CudaDriver):
     def __init__(self, *args, **kwargs) -> TVMFFIDriver:
         super().__init__(*args, **kwargs)
-        self.utils: TVMFFIUtils = TVMFFIUtils()
+        self.utils = utils
         self.launcher_cls: Type[TVMLauncher] = TVMLauncher
 
 

python/triton_tvm_ffi/utils.py

@@ -0,0 +1,38 @@
+# tvm-ffi-stubgen(begin): import-section
+# fmt: off
+# isort: off
+from __future__ import annotations
+from tvm_ffi import init_ffi_api as _FFI_INIT_FUNC
+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)
+
+# tvm-ffi-stubgen(begin): global/triton_tvm_ffi.utils
+# fmt: off
+_FFI_INIT_FUNC("triton_tvm_ffi.utils", __name__)
+if TYPE_CHECKING:
+    def build_signature_metadata(*args: Any) -> Any: ...
+    def cuOccupancyMaxActiveClusters(*args: Any) -> Any: ...
+    def fill_tma_descriptor(*args: Any) -> Any: ...
+    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]: ...
+    def set_printf_fifo_size(*args: Any) -> Any: ...
+# fmt: on
+# tvm-ffi-stubgen(end)
+
+__all__ = [
+    # tvm-ffi-stubgen(begin): __all__
+    "build_signature_metadata",
+    "cuOccupancyMaxActiveClusters",
+    "fill_tma_descriptor",
+    "get_device_properties",
+    "launch",
+    "load_binary",
+    "set_printf_fifo_size",
+    # tvm-ffi-stubgen(end)
+]