This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 123543a62 [Fix][CUDA] Align device target architecture detection (#193)
123543a62 is described below
commit 123543a6274f29685bdb09e813f8dc9ef41729ee
Author: Shushi Hong <[email protected]>
AuthorDate: Wed Nov 4 16:07:43 2026 -0500
[Fix][CUDA] Align device target architecture detection (#193)
This pr appends the architecture-specific suffix for detected CUDA
devices with compute capability 9.0 or later, matching the existing C++
defaults. Document the compatibility policy and cover both input forms,
architecture boundaries, and nonzero device indices.
---
python/tvm_tirx/backend/cuda/__init__.py | 8 +++++++-
python/tvm_tirx/target/target.py | 4 ++++
tests/python/target/test_target.py | 2 +-
3 files changed, 12 insertions(+), 2 deletions(-)
diff --git a/python/tvm_tirx/backend/cuda/__init__.py
b/python/tvm_tirx/backend/cuda/__init__.py
index a26155860..497426972 100644
--- a/python/tvm_tirx/backend/cuda/__init__.py
+++ b/python/tvm_tirx/backend/cuda/__init__.py
@@ -47,13 +47,19 @@ _LAZY_SUBMODULES = {
def _detect_target_from_device(dev):
from tvm_tirx.target import Target # pylint:
disable=import-outside-toplevel
+ compute_version = dev.compute_version.replace(".", "")
+ arch = "sm_" + compute_version
+ # Match UpdateCUDAAttrs when inferring an architecture for a local device.
+ if int(compute_version) >= 90:
+ arch += "a"
+
return Target(
{
"kind": "cuda",
"max_shared_memory_per_block": dev.max_shared_memory_per_block,
"max_threads_per_block": dev.max_threads_per_block,
"thread_warp_size": dev.warp_size,
- "arch": "sm_" + dev.compute_version.replace(".", ""),
+ "arch": arch,
}
)
diff --git a/python/tvm_tirx/target/target.py b/python/tvm_tirx/target/target.py
index b6598a3b9..b9c31313f 100644
--- a/python/tvm_tirx/target/target.py
+++ b/python/tvm_tirx/target/target.py
@@ -174,6 +174,10 @@ class Target(tvm_ffi.Object):
"""Detects Target associated with the given device. If the device does
not exist,
there will be an Error.
+ CUDA devices with compute capability 9.0 or later use the
architecture-specific
+ ``sm_XXa`` variant, matching ``Target("cuda")``. Code compiled for
this variant
+ is specific to that compute capability.
+
Parameters
----------
dev : Union[str, Device]
diff --git a/tests/python/target/test_target.py
b/tests/python/target/test_target.py
index 3e94a16e7..fe633d1cb 100644
--- a/tests/python/target/test_target.py
+++ b/tests/python/target/test_target.py
@@ -375,7 +375,7 @@ def test_target_from_device_cuda(input_form):
assert target.attrs["max_threads_per_block"] ==
dev.max_threads_per_block
assert int(target.attrs["max_shared_memory_per_block"]) ==
dev.max_shared_memory_per_block
assert int(target.attrs["thread_warp_size"]) == dev.warp_size
- assert str(target.attrs.get("arch", "")) == "sm_" +
dev.compute_version.replace(".", "")
+ assert target.arch == Target("cuda").arch
tvm_tirx.testing.run_with_gpu_lock(run_and_check)
