support GetDeviceProperties

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

python/triton_tvm_ffi/driver.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from typing import Type
+from typing import Mapping, Type
 from triton.backends.nvidia.driver import CudaDriver
+from .utils import get_device_properties
 
 
 class TVMFFIUtils(object):
@@ -19,8 +20,8 @@ class TVMFFIUtils(object):
     def load_binary(self, *args, **kwargs):
         return self._utils.load_binary(*args, **kwargs)
 
-    def get_device_properties(self, *args, **kwargs):
-        return self._utils.get_device_properties(*args, **kwargs)
+    def get_device_properties(self, device_id: int) -> Mapping[str, int]:
+        return get_device_properties(device_id)
 
     def cuOccupancyMaxActiveClusters(self, *args, **kwargs):
         return self._utils.cuOccupancyMaxActiveClusters(*args, **kwargs)

python/triton_tvm_ffi/utils/_ffi_api.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 from tvm_ffi import init_ffi_api as _FFI_INIT_FUNC
 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
 # isort: on
 # fmt: on
 # tvm-ffi-stubgen(end)
@@ -14,13 +16,13 @@ LIB = _FFI_LOAD_LIB("triton_tvm_ffi", "utils")
 # fmt: off
 _FFI_INIT_FUNC("triton_tvm_ffi.utils", __name__)
 if TYPE_CHECKING:
-    def hello() -> None: ...
+    def get_device_properties(_0: int, /) -> Mapping[str, int]: ...
 # fmt: on
 # tvm-ffi-stubgen(end)
 
 __all__ = [
     # tvm-ffi-stubgen(begin): __all__
     "LIB",
-    "hello",
+    "get_device_properties",
     # tvm-ffi-stubgen(end)
 ]