mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-07-01 08:51:56 +08:00
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Callable, Final, List, Sequence, Type, Union
|
|
from triton.backends.nvidia.driver import CudaDriver
|
|
from triton.runtime import _allocation
|
|
from . import TVMFFILauncherImpl, utils, string_to_type
|
|
|
|
|
|
class TVMLauncher(object):
|
|
def __init__(self, src, metadata, *args, **kwargs) -> TVMLauncher:
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.signature: List[str] = [*src.signature.values()]
|
|
self.num_ctas: Final[int] = getattr(metadata, "num_ctas", 1)
|
|
self.global_scratch_size: Final[int] = metadata.global_scratch_size
|
|
self.global_scratch_align: Final[int] = metadata.global_scratch_align
|
|
self.profile_scratch_size: Final[int] = metadata.profile_scratch_size
|
|
self.profile_scratch_align: Final[int] = metadata.profile_scratch_align
|
|
self.launch_cooperative_grid: Final[bool] = metadata.launch_cooperative_grid
|
|
self.launch_pdl: Final[bool] = metadata.launch_pdl
|
|
self.impl: TVMFFILauncherImpl = TVMFFILauncherImpl(
|
|
[string_to_type(t) for t in self.signature],
|
|
self.launch_cooperative_grid,
|
|
self.launch_pdl,
|
|
)
|
|
self.launch: Callable[
|
|
[
|
|
int,
|
|
int,
|
|
int,
|
|
int,
|
|
int,
|
|
tuple[int, int, int],
|
|
object,
|
|
object,
|
|
object,
|
|
object,
|
|
object,
|
|
Sequence[Union[Any]],
|
|
]
|
|
] = self.impl.launch
|
|
|
|
def __call__(
|
|
self,
|
|
gridX,
|
|
gridY,
|
|
gridZ,
|
|
stream,
|
|
function,
|
|
kernel_metadata,
|
|
launch_metadata,
|
|
launch_enter_hook,
|
|
launch_exit_hook,
|
|
*args,
|
|
):
|
|
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
|
|
|
|
# args: Sequence[TypedValue] = TypedValue.make_typed_values(self.signature, args)
|
|
|
|
return self.impl.launch(
|
|
gridX,
|
|
gridY,
|
|
gridZ,
|
|
stream,
|
|
function,
|
|
kernel_metadata,
|
|
launch_metadata,
|
|
launch_enter_hook,
|
|
launch_exit_hook,
|
|
global_scratch,
|
|
profile_scratch,
|
|
args,
|
|
)
|
|
|
|
|
|
class TVMFFIDriver(CudaDriver):
|
|
def __init__(self, *args, **kwargs) -> TVMFFIDriver:
|
|
super().__init__(*args, **kwargs)
|
|
self.utils = utils
|
|
self.launcher_cls: Type[TVMLauncher] = TVMLauncher
|
|
|
|
|
|
del CudaDriver
|