Skip to content

Commit 445107f

Browse files
committed
env and test fix
1 parent 13bdcd6 commit 445107f

File tree

2 files changed

+4
-5
lines changed

2 files changed

+4
-5
lines changed

testing/python/language/test_tilelang_language_let.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ def main(A_ptr: T.handle):
1616

1717
mod = tvm.IRModule({"main": main})
1818
mod = tvm.compile(mod, target="cuda")
19-
assert "float4 b" in mod.mod.imported_modules[0].get_source()
19+
assert "float4 b" in mod.mod.imports[0].inspect_source()
2020

2121

2222
if __name__ == "__main__":

tilelang/env.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -297,12 +297,11 @@ def prepend_pythonpath(path):
297297
if env.TVM_IMPORT_PYTHON_PATH is not None:
298298
prepend_pythonpath(env.TVM_IMPORT_PYTHON_PATH)
299299
else:
300-
tvm_path = os.path.join(THIRD_PARTY_ROOT, "tvm")
300+
tvm_path = os.path.join(THIRD_PARTY_ROOT, 'tvm', 'python')
301301
assert os.path.exists(tvm_path), tvm_path
302302
if tvm_path not in sys.path:
303-
tvm_python_binding = os.path.join(tvm_path, 'python')
304-
prepend_pythonpath(tvm_python_binding)
305-
env.TVM_IMPORT_PYTHON_PATH = tvm_python_binding
303+
prepend_pythonpath(tvm_path)
304+
env.TVM_IMPORT_PYTHON_PATH = tvm_path
306305

307306
if os.environ.get("TVM_LIBRARY_PATH") is None:
308307
os.environ['TVM_LIBRARY_PATH'] = env.TVM_LIBRARY_PATH = os.pathsep.join(TL_LIBS)

0 commit comments

Comments
 (0)