mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-07-01 08:51:56 +08:00
53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Mapping, Type
|
|
from triton.backends.nvidia.driver import CudaDriver
|
|
from .utils import get_device_properties
|
|
|
|
|
|
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)
|
|
from triton.backends.nvidia.driver import CudaUtils
|
|
|
|
self._utils: CudaUtils = CudaUtils()
|
|
|
|
def load_binary(self, *args, **kwargs):
|
|
return self._utils.load_binary(*args, **kwargs)
|
|
|
|
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):
|
|
return self._utils.fill_tma_descriptor(*args, **kwargs)
|
|
|
|
def launch(self, *args, **kwargs):
|
|
return self._utils.launch(*args, **kwargs)
|
|
|
|
def build_signature_metadata(self, *args, **kwargs):
|
|
return self._utils.build_signature_metadata(*args, **kwargs)
|
|
|
|
|
|
class TVMFFIDriver(CudaDriver):
|
|
def __init__(self, *args, **kwargs) -> TVMFFIDriver:
|
|
super().__init__(*args, **kwargs)
|
|
self.utils: TVMFFIUtils = TVMFFIUtils()
|
|
|
|
|
|
del CudaDriver
|