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_ldflags: Optional[Sequence[str]] = None,
extra_include_paths: Optional[Sequence[Union[str, Path]]] = None, extra_include_paths: Optional[Sequence[Union[str, Path]]] = None,
) -> TVMFFIWrapperFunction: ) -> TVMFFIWrapperFunction:
cuda_home: str = tvm_ffi.cpp.extension._find_cuda_home()
return wrap( return wrap(
name, name,
fns, fns,
code, code,
extra_ldflags=[ extra_ldflags=[
"-Wl,--no-as-needed", "-Wl,--no-as-needed",
f"-L{cuda_home}/lib64",
*map( *map(
lambda path: f"-L{path}", lambda path: f"-L{path}",
torch.utils.cpp_extension.library_paths(), torch.utils.cpp_extension.library_paths(),
), ),
"-lcuda",
"-lc10", "-lc10",
"-ltorch", "-ltorch",
] ]
+ (extra_ldflags or []), + (extra_ldflags or []),
extra_cflags=extra_cflags, extra_cflags=extra_cflags,
extra_cuda_cflags=extra_cuda_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 []), + (extra_include_paths or []),
) )