File size: 3,457 Bytes
fa1aa1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | from __future__ import annotations
import pathlib
from functools import lru_cache
from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union
if TYPE_CHECKING:
from tvm_ffi import Module
def _make_wrapper(tup: Tuple[str, str]) -> str:
export_name, kernel_name = tup
return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));"
@lru_cache()
def _resolve_kernel_path() -> pathlib.Path:
cur_dir = pathlib.Path(__file__).parent.resolve()
# first, try this directory structure
def _environment_install():
candidate = cur_dir.resolve()
if (candidate / "include").exists() and (candidate / "csrc").exists():
return candidate
return None
def _package_install():
# TODO: support find path by package
return None
path = _environment_install() or _package_install()
if path is None:
raise RuntimeError("Cannot find sgl-kernel/jit path")
return path
KERNEL_PATH = _resolve_kernel_path()
DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")]
DEFAULT_CFLAGS = ["-std=c++20", "-O3"]
DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"]
DEFAULT_LDFLAGS = []
CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool]
class CPPArgList(list[str]):
def __str__(self) -> str:
return ", ".join(self)
def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList:
def _convert(arg: CPP_TEMPLATE_TYPE) -> str:
if isinstance(arg, bool):
return "true" if arg else "false"
if isinstance(arg, (int, float)):
return str(arg)
raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}")
return CPPArgList(_convert(arg) for arg in args)
def load_jit(
*args: str,
cpp_files: List[str] | None = None,
cuda_files: List[str] | None = None,
cpp_wrappers: List[Tuple[str, str]] | None = None,
cuda_wrappers: List[Tuple[str, str]] | None = None,
extra_cflags: List[str] | None = None,
extra_cuda_cflags: List[str] | None = None,
extra_ldflags: List[str] | None = None,
extra_include_paths: List[str] | None = None,
build_directory: str | None = None,
) -> Module:
from tvm_ffi.cpp import load_inline
cpp_files = cpp_files or []
cuda_files = cuda_files or []
cpp_wrappers = cpp_wrappers or []
cuda_wrappers = cuda_wrappers or []
extra_cflags = extra_cflags or []
extra_cuda_cflags = extra_cuda_cflags or []
extra_ldflags = extra_ldflags or []
extra_include_paths = extra_include_paths or []
# include cpp files
cpp_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cpp_files]
cpp_sources = [f'#include "{path}"' for path in cpp_paths]
cpp_sources += [_make_wrapper(tup) for tup in cpp_wrappers]
# include cuda files
cuda_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cuda_files]
cuda_sources = [f'#include "{path}"' for path in cuda_paths]
cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers]
return load_inline(
"sgl_kernel_jit_" + "_".join(str(arg) for arg in args),
cpp_sources=cpp_sources,
cuda_sources=cuda_sources,
extra_cflags=DEFAULT_CFLAGS + extra_cflags,
extra_cuda_cflags=DEFAULT_CUDA_CFLAGS + extra_cuda_cflags,
extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
build_directory=build_directory,
)
|