mirror of
https://github.com/sgjzfzzf/triton-tvm-ffi.git
synced 2026-07-01 08:51:56 +08:00
89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, List, Optional, Type
|
|
from triton.backends.nvidia.driver import CudaDriver
|
|
from triton.runtime import _allocation
|
|
from . import TypedValue, 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: 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,
|
|
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
|
|
assert len(self.signature) == len(args)
|
|
|
|
def canonicalize(arg: Any, sig: str) -> TypedValue:
|
|
ty: Optional[int] = string_to_type(sig)
|
|
assert ty is not None, sig
|
|
return TypedValue(ty, arg)
|
|
|
|
args = [canonicalize(arg, sig) for arg, sig in zip(args, self.signature)]
|
|
|
|
return self.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 = utils
|
|
self.launcher_cls: Type[TVMLauncher] = TVMLauncher
|
|
|
|
|
|
del CudaDriver
|