Skip to content

Commit d321035

Browse files
authored
[Relax][Frontend][NN] Fix SourceModule include resolution for installed wheels (#19822)
`SourceModule.get_includes()` hardcoded `tvm_home()/3rdparty/tvm-ffi/...` paths that only exist in a source checkout. In a wheel the tvm-ffi/DLPack headers ship in the separate `tvm_ffi` package, so the existence assert failed and `nn.SourceModule` could not compile extern sources. This pr resolves includes via libinfo (`libinfo.find_include_path()`, `tvm_ffi.libinfo.find_include_path()`, `find_dlpack_include_path()`), matching `runtime/module.py` and `rpc/minrpc.py`.
1 parent 244499a commit d321035

1 file changed

Lines changed: 8 additions & 5 deletions

File tree

python/tvm/relax/frontend/nn/extern.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@
2424
from collections.abc import Callable
2525
from pathlib import Path
2626

27-
from tvm import tirx
27+
import tvm_ffi
28+
29+
from tvm import libinfo, tirx
2830
from tvm.runtime import Module, load_static_library
2931
from tvm.support import cc as _cc
3032

@@ -306,15 +308,16 @@ def get_includes(tvm_pkg: list[str] | None = None) -> list[Path]:
306308
includes : List[pathlib.Path]
307309
The list of include paths.
308310
"""
309-
tvm_home = SourceModule.tvm_home()
310311
results = [
311-
tvm_home / "include",
312-
tvm_home / "3rdparty/tvm-ffi/include",
313-
tvm_home / "3rdparty/tvm-ffi/3rdparty/dlpack/include",
312+
Path(libinfo.find_include_path()),
313+
Path(tvm_ffi.libinfo.find_include_path()),
314+
Path(tvm_ffi.libinfo.find_dlpack_include_path()),
314315
]
315316
if tvm_pkg:
317+
tvm_home = SourceModule.tvm_home()
316318
for relative in tvm_pkg:
317319
results.append(tvm_home / "3rdparty" / relative)
320+
results = list(dict.fromkeys(results))
318321
for path in results:
319322
assert path.exists(), f"Not found: {path!s}"
320323
assert path.is_dir(), f"Not a directory: {path!s}"

0 commit comments

Comments
 (0)