This is an automated email from the ASF dual-hosted git repository.
wuwei pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 5648a8e114 [Runtime] Add property Module.is_device_module (#17407)
5648a8e114 is described below
commit 5648a8e1149294ca0b84151564ac46505fd18279
Author: Wuwei Lin <[email protected]>
AuthorDate: Tue Sep 24 21:09:32 2024 -0700
[Runtime] Add property Module.is_device_module (#17407)
---
python/tvm/relax/vm_build.py | 2 +-
python/tvm/runtime/module.py | 4 ++++
2 files changed, 5 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index 9fd7a74285..cfa4143b66 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -252,7 +252,7 @@ def _vmlink(
runtime=_autodetect_system_lib_req(target, system_lib),
)
for ext_mod in ext_libs:
- if ext_mod.type_key == "cuda":
+ if ext_mod.is_device_module:
tir_ext_libs.append(ext_mod)
else:
relax_ext_libs.append(ext_mod)
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 2c3eff7000..ca151293bb 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -274,6 +274,10 @@ class Module(object):
"""
return (self.get_property_mask() & ModulePropertyMask.RUNNABLE) != 0
+ @property
+ def is_device_module(self):
+ return self.type_key in ["cuda", "opencl", "metal", "hip", "vulkan",
"webgpu"]
+
@property
def is_dso_exportable(self):
"""Returns true if module is 'DSO exportable', ie can be included in
result of
