JinjieLiu/triton-tvm-ffi
1
0
Fork 0
mirror of https://github.com/sgjzfzzf/triton-tvm-ffi.git synced 2026-05-02 03:52:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix bugs on lacking cudatoolkit

Browse Source 
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
This commit is contained in:
Jinjie Liu
2026-02-04 10:27:44 +08:00
parent dc8c2c17e0
commit 6e4c2d4a43
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/triton_tvm_ffi/wrap.py
View File

@@ -155,22 +155,28 @@ def torch_wrap(
extra_ldflags: Optional[Sequence[str]] = None,
extra_include_paths: Optional[Sequence[Union[str, Path]]] = None,
) -> TVMFFIWrapperFunction:
cuda_home: str = tvm_ffi.cpp.extension._find_cuda_home()
return wrap(
name,
fns,
code,
extra_ldflags=[
"-Wl,--no-as-needed",
f"-L{cuda_home}/lib64",
*map(
lambda path: f"-L{path}",
torch.utils.cpp_extension.library_paths(),
),
"-lcuda",
"-lc10",
"-ltorch",
]
+ (extra_ldflags or []),
extra_cflags=extra_cflags,
extra_cuda_cflags=extra_cuda_cflags,
extra_include_paths=[*torch.utils.cpp_extension.include_paths()]
extra_include_paths=[
f"{cuda_home}/include",
*torch.utils.cpp_extension.include_paths(),
]
+ (extra_include_paths or []),
)