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.git
The following commit(s) were added to refs/heads/main by this push:
new 3a0b83a7e9 refactor(infra)!: pass backend configuration through Target
(#20651)
3a0b83a7e9 is described below
commit 3a0b83a7e966ddd35fff6bc811e4923f50cbfb1e
Author: Bohan Hou <[email protected]>
AuthorDate: Sun Oct 11 07:27:33 2026 -0400
refactor(infra)!: pass backend configuration through Target (#20651)
CUDA backend configuration already travels through Target attributes.
Remove the parallel `backend_config` keyword from `tvm.compile`,
`tvm.tirx.build`, `tvm.relax.build`, and `IketProfiler.compile`;
compilation now resolves its defaults from the supplied Target.
```python
import tvm
from tvm.backend.cuda import BackendConfig
target = tvm.target.Target({
"kind": "cuda",
"arch": "sm_100a",
"backend_config": {
"cuda": BackendConfig(compiler="nvrtc", nvrtc=["--use_fast_math",
"--ftz=false"])
},
})
executable = tvm.compile(func, target=target)
```
Each `device_entry(backend_config=...)` keeps its independent overrides,
including architecture and toolchain settings. The existing per-entry
resolution and module grouping continue to apply. Update callers, source
helpers, and documentation to use Target defaults followed by entry
overrides.
Follow-up to merged #20645.
Validation:
- Built successfully on the merged #20645 baseline (`0f861ee568`) in 67
seconds.
- Full TIRx suite: 2,689 passed / 290 skipped on seven B200 GPUs, 165
seconds.
- Backend configuration, launch, and IKET regressions: 131 passed / 2
skipped.
- CUDA codegen and Relax VM build regressions: 415 passed / 4 skipped.
- Updated TIRx-kernels callers: 94 focused tests passed; all eight MQA
registry correctness cases passed. Both workload and strict registry
import gates passed (274 workloads / 93 kernels).
- Changed-file and staged pre-commit checks passed.
The latest upstream commit (`27576ce2d1`) fails the local NCCL-enabled
C++ build in `cuda_ipc_memory.cc` after its object-handle changes. The
checks above therefore use the merged #20645 baseline; this follow-up
changes Python APIs, tests, and documentation only.
---
docs/tirx/api/cuda_compile.rst | 11 +++-
docs/tirx/native_basics/cuda/compiling.rst | 23 ++++---
python/tvm/backend/config.py | 32 ++++------
python/tvm/backend/cuda/iket.py | 6 +-
python/tvm/backend/cuda/transforms/__init__.py | 7 +--
python/tvm/driver/build_module.py | 8 +--
python/tvm/relax/vm_build.py | 15 +----
python/tvm/tirx/build.py | 13 ++--
tests/python/codegen/test_target_codegen_cuda.py | 79 ++++++++++++++----------
tests/python/disco/test_nvshmem.py | 3 +-
tests/python/tirx/codegen/test_codegen_cuda.py | 15 ++++-
tests/python/tirx/test_cuda_backend_config.py | 69 +++++++++++----------
tests/python/tirx/test_cuda_launch.py | 2 +-
13 files changed, 150 insertions(+), 133 deletions(-)
diff --git a/docs/tirx/api/cuda_compile.rst b/docs/tirx/api/cuda_compile.rst
index dc89fb615b..2f5c1d8e47 100644
--- a/docs/tirx/api/cuda_compile.rst
+++ b/docs/tirx/api/cuda_compile.rst
@@ -17,7 +17,8 @@
CUDA backend configuration
==========================
-Pass a nested mapping directly to ``tvm.compile(..., backend_config=...)`` or
+Set the nested mapping on ``Target.backend_config`` and pass that Target to
+``tvm.compile``. Individual kernels can override it with
``T.device_entry(..., backend_config=...)``. Each backend owns its
configuration;
CUDA currently provides these fixed keys:
@@ -58,10 +59,14 @@ compiler arguments are not modeled as separate Python
fields.
from tvm.backend.cuda import BackendConfig
cuda_config = BackendConfig(
- arch="sm_100a", compiler="nvrtc",
+ compiler="nvrtc",
nvrtc=["--use_fast_math", "--ftz=false"],
)
- executable = tvm.compile(func, backend_config={"cuda": cuda_config})
+ target = tvm.target.Target({
+ "kind": "cuda", "arch": "sm_100a",
+ "backend_config": {"cuda": cuda_config},
+ })
+ executable = tvm.compile(func, target=target)
.. autoclass:: tvm.backend.cuda.BackendConfig
:members:
diff --git a/docs/tirx/native_basics/cuda/compiling.rst
b/docs/tirx/native_basics/cuda/compiling.rst
index 68b56b3deb..7611b2a2fa 100644
--- a/docs/tirx/native_basics/cuda/compiling.rst
+++ b/docs/tirx/native_basics/cuda/compiling.rst
@@ -19,7 +19,8 @@ Compiling and inspecting
========================
CUDA compilation accepts a lightweight ``backend_config`` mapping. The same
-nested structure is accepted by ``tvm.compile``, ``T.device_entry``, and
+nested structure is accepted by the Target's ``backend_config`` attribute and
+``T.device_entry``. Pass the configured Target to ``tvm.compile`` or
``txl.Kernel.compile``. ``tvm.backend.cuda.BackendConfig`` is an optional
``TypedDict`` helper, also available as ``T.cuda.BackendConfig`` and
``txl.cuda.BackendConfig``:
@@ -46,8 +47,11 @@ nested structure is accepted by ``tvm.compile``,
``T.device_entry``, and
y = T.cuda.thread_idx("x")
A[y] = B[y] + T.float32(1)
- cuda_config = BackendConfig(arch="sm_100a")
- exe = tvm.compile(pipeline, backend_config={"cuda": cuda_config})
+ target = tvm.target.Target({
+ "kind": "cuda", "arch": "sm_100a",
+ "backend_config": {"cuda": BackendConfig(compiler="nvrtc")},
+ })
+ exe = tvm.compile(pipeline, target=target)
Each entry resolves its configuration before architecture-sensitive lowering.
Entries with different resolved targets or options compile as separate CUDA
@@ -58,8 +62,8 @@ Kernels must be compatible with the device on which they
execute.
Defaults and overrides
----------------------
-Precedence, from lowest to highest, is backend defaults, Target/tag defaults,
-compile-call overrides, then device-entry overrides. ``None`` and ``{}`` add no
+Precedence, from lowest to highest, is backend defaults, Target/tag settings,
+then device-entry overrides. ``None`` and ``{}`` add no
overrides. Toolchain lists replace inherited lists in full, and ``[]`` clears a
list. To keep fast math while changing FTZ, include both arguments as above.
Configuration boundaries snapshot dictionaries and lists, so later mutation of
@@ -79,11 +83,10 @@ Target tags use the existing target registry and carry
defaults in Target attrs:
Both ``Target("local/blackwell")`` and ``Target({"tag": "local/blackwell"})``
preserve these defaults in their expanded attributes. There is no separate
-backend tag registry. An explicit ``Target.arch`` conflicting with build-level
-``backend_config["cuda"]["arch"]`` is rejected. An entry can override the arch.
+backend tag registry. An entry can override the architecture supplied by its
Target.
Online builds can detect the GPU architecture. Offline builds require an
-explicit architecture in the Target, build configuration, or every device
entry.
+explicit architecture in the Target or every device entry.
A generic ``Target("cuda")`` may remain without an architecture for backend
discovery and IR construction; compilation validates it before lowering.
Architecture-dependent Python factories resolve and record their architecture
@@ -127,8 +130,8 @@ Migration
Replace the former compilation dataclass with the nested mapping above. Convert
math, debugging, include, and assembler options to native argument strings.
-Replace ``@txl.kernel(arch=...)`` and ``tirx.cuda_arch`` with an entry or
compile
-``backend_config={"cuda": {"arch": ...}}``.
+Replace ``@txl.kernel(arch=...)`` and ``tirx.cuda_arch`` with ``Target.arch``
or
+an entry's ``backend_config={"cuda": {"arch": ...}}``.
Compiler/math/ptxas environment policies remain removed. CPU benchmark
preparation receives an explicit architecture through ``backend_config`` and
diff --git a/python/tvm/backend/config.py b/python/tvm/backend/config.py
index 2d263bc471..8994d62d37 100644
--- a/python/tvm/backend/config.py
+++ b/python/tvm/backend/config.py
@@ -56,13 +56,12 @@ def parse_backend_config(value):
return copy_backend_config(json.loads(value) if value else None)
-def prepare_target(target, config=None, mod=None):
+def prepare_target(target, mod=None):
"""Resolve the build target, including architectures supplied by CUDA
entries."""
import tvm_ffi
from tvm.target import Target
- config = copy_backend_config(config)
active = Target.current() if target is None else target
if isinstance(active, Target) and active.kind.name == "cuda" and "arch"
not in active.attrs:
active = dict(active.export())
@@ -71,8 +70,13 @@ def prepare_target(target, config=None, mod=None):
or active == "cuda"
or (isinstance(active, Mapping) and active.get("kind") == "cuda" and
"arch" not in active)
)
+ if not generic_cuda:
+ return target
+ config = (
+ copy_backend_config(active.get("backend_config", {})) if
isinstance(active, Mapping) else {}
+ )
entries = []
- if mod is not None and generic_cuda:
+ if mod is not None:
from tvm.ir import RegionStmt
def visit(node):
@@ -83,23 +87,13 @@ def prepare_target(target, config=None, mod=None):
if not config and not any(entries):
return target
cuda = config.get("cuda", {})
- if generic_cuda:
- arch = cuda.get("arch")
- if arch is None and entries and all(c.get("cuda", {}).get("arch") for
c in entries):
- arch = entries[0]["cuda"]["arch"]
- attrs = dict(active) if isinstance(active, Mapping) else {"kind":
"cuda"}
- if arch is not None:
- attrs["arch"] = arch
- return Target(attrs)
- active = Target(active)
- if cuda and active.kind.name != "cuda":
- raise ValueError("CUDA backend_config requires a CUDA target")
arch = cuda.get("arch")
- if arch is not None and active.attrs.get("arch") != arch:
- raise ValueError(
- f"Target.arch={active.attrs.get('arch')!r} conflicts with
backend_config arch={arch!r}"
- )
- return active
+ if arch is None and entries and all(c.get("cuda", {}).get("arch") for c in
entries):
+ arch = entries[0]["cuda"]["arch"]
+ attrs = dict(active) if isinstance(active, Mapping) else {"kind": "cuda"}
+ if arch is not None:
+ attrs["arch"] = arch
+ return Target(attrs)
def argparse_backend_config(value):
diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py
index ab148c0e20..3912adcb0d 100644
--- a/python/tvm/backend/cuda/iket.py
+++ b/python/tvm/backend/cuda/iket.py
@@ -309,16 +309,14 @@ class IketProfiler:
def sentinel_token(self, name: str):
return T.cuda.iket.sentinel_token(name)
- def compile(self, mod, target=None, *, tir_pipeline="tirx",
backend_config=None):
+ def compile(self, mod, target=None, *, tir_pipeline="tirx"):
"""Compile official IKET metadata and NativeDump placeholders."""
if isinstance(mod, tvm.tirx.Function):
mod = tvm.IRModule.from_expr(mod)
if not isinstance(mod, tvm.IRModule):
raise TypeError("IketProfiler.compile expects a TIRx Function or
IRModule")
enabled_mod = mod.with_attr("tirx.iket.enabled", True)
- executable = tvm.compile(
- enabled_mod, target=target, tir_pipeline=tir_pipeline,
backend_config=backend_config
- )
+ executable = tvm.compile(enabled_mod, target=target,
tir_pipeline=tir_pipeline)
return _OfficialIketExecutable(executable)
diff --git a/python/tvm/backend/cuda/transforms/__init__.py
b/python/tvm/backend/cuda/transforms/__init__.py
index 3859b03ddb..fd37f294cb 100644
--- a/python/tvm/backend/cuda/transforms/__init__.py
+++ b/python/tvm/backend/cuda/transforms/__init__.py
@@ -38,12 +38,9 @@ def LowerIket():
return get_global_func("tirx.backend.cuda.transforms.LowerIket")()
-def BindBackendConfig(backend_config=None):
+def BindBackendConfig():
"""Resolve per-entry compilation settings and targets before tile
dispatch."""
- from tvm.backend.config import backend_config_json
-
- config = "" if backend_config is None else
backend_config_json(backend_config)
- return
get_global_func("tirx.backend.cuda.transforms.BindBackendConfig")(config)
+ return
get_global_func("tirx.backend.cuda.transforms.BindBackendConfig")("")
def SpecializeDeviceHelpers():
diff --git a/python/tvm/driver/build_module.py
b/python/tvm/driver/build_module.py
index c1296e439f..8dceee2413 100644
--- a/python/tvm/driver/build_module.py
+++ b/python/tvm/driver/build_module.py
@@ -42,7 +42,6 @@ def compile( # pylint: disable=redefined-builtin
*,
relax_pipeline: tvm.transform.Pass | Callable | str | None = "default",
tir_pipeline: tvm.transform.Pass | Callable | str | None = "default",
- backend_config=None,
) -> Executable:
"""
Compile an IRModule to a runtime executable.
@@ -56,14 +55,12 @@ def compile( # pylint: disable=redefined-builtin
The input module to be compiled. Can be a Function or an IRModule
containing
TIR or Relax functions.
target : Optional[Target]
- The target platform to compile for.
+ The target platform and its backend compilation configuration.
relax_pipeline : Optional[Union[tvm.transform.Pass, Callable, str]]
The compilation pipeline to use for Relax functions.
Only used if the module contains Relax functions.
tir_pipeline : Optional[Union[tvm.transform.Pass, Callable, str]]
The compilation pipeline to use for TIR functions.
- backend_config : Optional[dict[str, dict]]
- Per-backend compiler defaults, overridden by each device entry.
Returns
-------
@@ -77,7 +74,6 @@ def compile( # pylint: disable=redefined-builtin
target,
relax_pipeline=relax_pipeline,
tir_pipeline=tir_pipeline,
- backend_config=backend_config,
)
- lib = tvm.tirx.build(mod, target, pipeline=tir_pipeline,
backend_config=backend_config)
+ lib = tvm.tirx.build(mod, target, pipeline=tir_pipeline)
return Executable(lib)
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index bd2bfc60ae..db65ecbf68 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -103,7 +103,6 @@ def _vmlink(
params: dict[str, list] | None = None,
*,
system_lib: bool | None = None,
- backend_config=None,
):
"""
Internal codegen function to make executable.
@@ -147,9 +146,7 @@ def _vmlink(
tir_ext_libs = []
if tir_mod is not None and len(tir_mod.get_global_vars()) > 0:
tir_mod = _auto_attach_system_lib_prefix(tir_mod, target, system_lib)
- lib = tvm.tirx.build(
- tir_mod, target=target, pipeline=tir_pipeline,
backend_config=backend_config
- )
+ lib = tvm.tirx.build(tir_mod, target=target, pipeline=tir_pipeline)
for ext_mod in ext_libs:
if _is_device_module(ext_mod):
tir_ext_libs.append(ext_mod)
@@ -172,7 +169,6 @@ def build(
tir_pipeline: None | str | tvm.transform.Pass = "default",
*,
system_lib: bool | None = None,
- backend_config=None,
) -> Executable:
"""
Build an IRModule to VM executable.
@@ -206,9 +202,6 @@ def build(
auto registers generated functions to the system.
By default auto detects based on the target.
- backend_config : Optional[dict[str, dict]]
- CUDA compiler defaults forwarded to the generated TIRx functions.
-
Returns
-------
ex: tvm.relax.Executable
@@ -236,10 +229,9 @@ def build(
constants = attrs.get("const_name_to_constant", {})
return ext_libs, constants
- from tvm.backend.config import copy_backend_config, prepare_target
+ from tvm.backend.config import prepare_target
- backend_config = copy_backend_config(backend_config)
- target = prepare_target(target, backend_config, mod)
+ target = prepare_target(target, mod)
if isinstance(target, str):
target = tvm.target.Target(target)
if not params:
@@ -280,7 +272,6 @@ def build(
ext_libs=ext_libs,
params=params,
system_lib=system_lib,
- backend_config=backend_config,
)
diff --git a/python/tvm/tirx/build.py b/python/tvm/tirx/build.py
index 31b9d17c36..8e65df8cbf 100644
--- a/python/tvm/tirx/build.py
+++ b/python/tvm/tirx/build.py
@@ -163,8 +163,6 @@ def build(
mod: Function | IRModule,
target: str | Target | None = None,
pipeline: str | tvm.transform.Pass | None = "default",
- *,
- backend_config=None,
):
"""Build a function with a signature, generating code for devices
coupled with target information.
@@ -174,11 +172,9 @@ def build(
mod : Union[Function, IRModule]
The input to be built.
target : Optional[Union[str, Target]]
- The target for compilation.
+ The target for compilation, including backend configuration defaults.
pipeline : Union[None, str, tvm.transform.Pass]
The pipeline to use for compilation.
- backend_config : Optional[dict[str, dict]]
- Per-backend compiler defaults, overridden by each device entry.
Returns
-------
@@ -191,11 +187,10 @@ def build(
else:
assert isinstance(mod, tvm.IRModule)
- from tvm.backend.config import copy_backend_config, prepare_target
+ from tvm.backend.config import prepare_target
from tvm.backend.cuda.transforms import BindBackendConfig,
SpecializeEntryHelpers
- backend_config = copy_backend_config(backend_config)
- target = prepare_target(target, backend_config, mod)
+ target = prepare_target(target, mod)
# Step 0: Determine the target in environment
# It's used to bind the Function without target attr to serve as a default
target
@@ -228,7 +223,7 @@ def build(
# Step 3: Bind the target to the input module
mod = tvm.tirx.transform.BindTarget(target_to_bind)(mod)
- mod = BindBackendConfig(backend_config)(mod)
+ mod = BindBackendConfig()(mod)
mod = SpecializeEntryHelpers()(mod)
# Step 4: Apply the tirx pipeline
diff --git a/tests/python/codegen/test_target_codegen_cuda.py
b/tests/python/codegen/test_target_codegen_cuda.py
index 3a8181e5b9..31b23b0176 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -55,7 +55,9 @@ def test_cuda_host_bundle(tmp_path, *, backend_config):
B[tx] = A[tx] + T.float32(1)
target = tvm.target.Target("cuda", host="cuda_host")
- built = tvm.compile(add_one, target=target,
backend_config=backend_config).mod
+ built = tvm.compile(
+ add_one, target={**tvm.target.Target(target).export(),
"backend_config": backend_config}
+ ).mod
source = export_cuda_host(built)
assert "tvm_cuda_binary_0" in source
assert "__global__" not in source
@@ -116,9 +118,8 @@ def test_cuda_host_bundle_bf16_cluster(tmp_path, *,
backend_config):
with target:
built = tvm.compile(
tvm.IRModule({"main": main}),
- target=target,
+ target={**tvm.target.Target(target).export(), "backend_config":
backend_config},
tir_pipeline="tirx",
- backend_config=backend_config,
).mod
source = export_cuda_host(built)
assert "cudaLaunchAttributeClusterDimension" in source
@@ -169,7 +170,9 @@ def
test_cuda_host_bundle_programmatic_dependent_launch(tmp_path, *, backend_con
target = tvm.target.Target("cuda", host="cuda_host")
source = export_cuda_host(
- tvm.compile(add_one, target=target, backend_config=backend_config).mod
+ tvm.compile(
+ add_one, target={**tvm.target.Target(target).export(),
"backend_config": backend_config}
+ ).mod
)
assert "cudaLaunchAttributeProgrammaticStreamSerialization" in source
library = tvm_ffi.cpp.build_inline(
@@ -217,9 +220,8 @@ def test_cuda_host_bundle_tensor_map_parameter(tmp_path,
monkeypatch, *, backend
with target:
built = tvm.compile(
tvm.IRModule({"main": main}),
- target=target,
+ target={**tvm.target.Target(target).export(), "backend_config":
backend_config},
tir_pipeline="tirx",
- backend_config=backend_config,
).mod
source = export_cuda_host(built)
assert "(CUtensorMap*)" in source
@@ -259,7 +261,7 @@ def test_cuda_vectorize_add(*, backend_config):
if i_0 * num_thread + i_1 < n:
B[i_0 * num_thread + i_1] = A[i_0 * num_thread +
i_1] + one
- fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ fun = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -323,7 +325,7 @@ def test_cuda_bf16_vectorize_add(*, backend_config):
with tvm.transform.PassContext(
disabled_pass=["tirx.BF16Promote", "tirx.BF16CastElimination",
"tirx.BF16TypeLowering"]
):
- fun = tvm.compile(Module, target="cuda",
backend_config=backend_config)
+ fun = tvm.compile(Module, target={"kind": "cuda",
"backend_config": backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -375,7 +377,7 @@ def test_cuda_multiply_add(*, backend_config):
ty="int32",
)
- fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ fun = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
np_a = np.random.randint(low=-128, high=127, size=(n, lanes))
np_b = np.random.randint(low=-128, high=127, size=(n, lanes))
@@ -414,7 +416,7 @@ def test_cuda_vectorize_load(*, backend_config):
for i_1 in T.thread_binding(num_thread,
thread="threadIdx.x"):
B[i_0 * num_thread + i_1] = A[i_0 * num_thread + i_1]
- fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ fun = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -449,7 +451,7 @@ def test_cuda_make_int8(*, backend_config):
for j in T.vectorized(lanes):
A[i, j] = const_value
- fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ fun = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -487,7 +489,7 @@ def test_cuda_inf_nan(*, backend_config):
if i_0 * 8 + i_1 < n:
C[i_0 * 8 + i_1] = inf_value
- fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ fun = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -546,7 +548,10 @@ def test_crossthread_reduction1(target, *, backend_config):
B[i] = reduced[0]
fun = tvm.compile(
- Module, target=target, backend_config=backend_config if target ==
"cuda" else None
+ Module,
+ target={**tvm.target.Target(target).export(), "backend_config":
backend_config}
+ if target == "cuda"
+ else target,
)
return fun
@@ -627,7 +632,10 @@ def test_crossthread_reduction2(target, *, backend_config):
B[i] = reduced[0]
func = tvm.compile(
- Module, target=target, backend_config=backend_config if target ==
"cuda" else None
+ Module,
+ target={**tvm.target.Target(target).export(), "backend_config":
backend_config}
+ if target == "cuda"
+ else target,
)
return func
@@ -670,7 +678,7 @@ def test_cuda_reduction_binding(*, backend_config):
B[m_0 * 32 + m_1] = T.float32(0.0)
B[m_0 * 32 + m_1] = B[m_0 * 32 + m_1] + A[m_0 * 32 +
m_1, k]
- func = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ func = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
@pytest.mark.gpu
@@ -702,7 +710,7 @@ def test_cuda_const_float_to_half(*, backend_config):
]
)
- func = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ func = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
shape = (2, 3, 4)
a_np = np.random.uniform(size=shape).astype("float16")
@@ -738,7 +746,7 @@ def test_cuda_floordiv_with_vectorization(*,
backend_config):
(i_0 * 256 + i_1_0 * 4 + i_1_1) // 37
]
- func = tvm.compile(Module, target="cuda",
backend_config=backend_config)
+ func = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -772,7 +780,7 @@ def test_cuda_floormod_with_vectorization(*,
backend_config):
(i_0 * 256 + i_1_0 * 4 + i_1_1) % 37
]
- func = tvm.compile(Module, target="cuda",
backend_config=backend_config)
+ func = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -855,7 +863,7 @@ def test_vectorized_casts(t0, t1, factor, *,
backend_config):
t0, B[i_0 * factor + i_1]
)
- func = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ func = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
# correctness
def run_and_check():
@@ -894,7 +902,7 @@ def sched(compute_fn, dtype, n=128, *, backend_config=None):
A[i0_1_0 * 4 + i0_1_1_0 * 4 + i0_1_1_1]
)
- return tvm.compile(Module, target="cuda", backend_config=backend_config)
+ return tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
@pytest.mark.gpu
@@ -1063,7 +1071,7 @@ def test_cuda_vectorize_load_permute_pad(*,
backend_config):
j < padding or l + padding <= j, zero, A[i *
lanes + k, j - padding]
)
- fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ fun = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
np_a = np.random.randint(low=-128, high=127, size=(n, l)).astype(dtype)
np_a_reshape = np_a.reshape(n // lanes, lanes, l).transpose(0, 2, 1)
@@ -1102,7 +1110,7 @@ def test_try_unaligned_vector_load(*, backend_config):
for i_1 in T.vectorized(2):
C[i_0 * 2 + i_1] = A[i_0 * 2 + i_1 + offset]
- f = tvm.tirx.build(Module, target="cuda",
backend_config=backend_config)
+ f = tvm.tirx.build(Module, target={"kind": "cuda", "backend_config":
backend_config})
kernel_source = f.imports[0].inspect_source()
a_data = np.arange(0, N).astype("float16")
@@ -1163,7 +1171,7 @@ def test_cuda_thread_sync_inside_condition(*,
backend_config):
for func in [func2, func3]:
mod = tvm.IRModule({"main": func})
- compiled = tvm.compile(mod, target="cuda",
backend_config=backend_config)
+ compiled = tvm.compile(mod, target={"kind": "cuda", "backend_config":
backend_config})
assert "__syncthreads()" in compiled.mod.imports[0].inspect_source()
@@ -1176,7 +1184,7 @@ def test_invalid_reinterpret(*, backend_config):
B[tx] = T.call_intrin("tirx.reinterpret", A[tx], ty="uint8")
with pytest.raises(RuntimeError):
- tvm.compile(func, target="cuda", backend_config=backend_config)
+ tvm.compile(func, target={"kind": "cuda", "backend_config":
backend_config})
@pytest.mark.gpu
@@ -1197,7 +1205,7 @@ def test_cuda_tensormap(*, backend_config):
# fmt: on
mod = tvm.IRModule({"main": main})
- mod = tvm.compile(mod, target="cuda", backend_config=backend_config)
+ mod = tvm.compile(mod, target={"kind": "cuda", "backend_config":
backend_config})
assert (
"""
extern "C" __global__ void __launch_bounds__(128) main_kernel(float*
__restrict__ A_ptr, const __grid_constant__ CUtensorMap A_map) {
@@ -1228,7 +1236,7 @@ def test_cuda_device_func_call(*, backend_config):
for tx in T.thread_binding(1024, "threadIdx.x"):
C[bx, tx] = Module.add(A[bx, tx], B[bx, tx])
- lib = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ lib = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
cuda_code = lib.mod.imports[0].inspect_source()
assert 'extern "C" __device__ float add(float a, float b) {\n return (a +
b);\n}' in cuda_code
@@ -1248,7 +1256,7 @@ def test_cuda_float_const_hex_format(*, backend_config):
for tx in T.thread_binding(1024, "threadIdx.x"):
A[bx, tx] = T.float32(1 / 27)
- lib = tvm.compile(Module, target="cuda", backend_config=backend_config)
+ lib = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
cuda_code = lib.mod.imports[0].inspect_source()
assert "0x1.2f684bda12f68p-5f" in cuda_code
@@ -1279,7 +1287,9 @@ def test_device_host_call_same_func(*, backend_config):
# 2. We set a dummy mcpu value for testing purpose,
# in order to avoid checking a function is host or device based on the
"cpu" substring.
target = tvm.target.Target({"kind": "cuda", "mcpu": "dummy_mcpu"},
host="c")
- lib = tvm.compile(Module, target=target, backend_config=backend_config)
+ lib = tvm.compile(
+ Module, target={**tvm.target.Target(target).export(),
"backend_config": backend_config}
+ )
cuda_code = lib.mod.imports[0].inspect_source()
assert 'extern "C" __device__ int add(int a, int b) {\n return (a +
b);\n}' in cuda_code
@@ -1313,7 +1323,11 @@ def test_thread_return(monkeypatch, *, backend_config):
B[bx, tx] = A[bx, tx]
lib = tvm.compile(
- Module, backend_config=merge_backend_configs(backend_config, {"cuda":
{"arch": "sm_80"}})
+ Module,
+ target={
+ "kind": "cuda",
+ "backend_config": merge_backend_configs(backend_config, {"cuda":
{"arch": "sm_80"}}),
+ },
)
cuda_code = lib.mod.imports[0].inspect_source()
assert "return;" in cuda_code
@@ -1336,7 +1350,10 @@ def test_cuda_loop_step(*, backend_config):
C[i] = A[i] + B[i]
target = tvm.target.Target({"kind": "cuda"})
- lib = tvm.compile(cuda_loop_step, target=target,
backend_config=backend_config)
+ lib = tvm.compile(
+ cuda_loop_step,
+ target={**tvm.target.Target(target).export(), "backend_config":
backend_config},
+ )
cuda_src = lib.mod.imports[0].inspect_source()
assert "i += 96" in cuda_src
@@ -1371,7 +1388,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path,
*, backend_config):
B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
- host_lib = tvm.compile(Module, target="cuda",
backend_config=backend_config)
+ host_lib = tvm.compile(Module, target={"kind": "cuda", "backend_config":
backend_config})
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
lib_path = str(tmp_path / "lib.so")
diff --git a/tests/python/disco/test_nvshmem.py
b/tests/python/disco/test_nvshmem.py
index 771701d4ac..2e6f7b183e 100644
--- a/tests/python/disco/test_nvshmem.py
+++ b/tests/python/disco/test_nvshmem.py
@@ -339,7 +339,8 @@ def _kernel_compile(compile_mode):
target = tvm.target.Target("cuda")
tvm.compile(
- NvshmemQueryModule, target=target,
backend_config=backend_config
+ NvshmemQueryModule,
+ target={**tvm.target.Target(target).export(),
"backend_config": backend_config},
).export_library(path)
mod = sess.load_vm_module(path)
result = mod["main"]()
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index 1bb4f79fab..64f23928e3 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -23,6 +23,7 @@ import pytest
import tvm
import tvm.testing
+from tvm.backend.config import merge_backend_configs
from tvm.backend.cuda.lang.clc import query_cancel_first_ctaid_x
from tvm.script import tirx as T
from tvm.testing import env
@@ -34,9 +35,21 @@ def _get_source(
if target is None:
target = {"kind": "cuda", "arch": "sm_100a"}
target = tvm.target.Target(target)
+ target = tvm.target.Target(
+ {
+ **target.export(),
+ "backend_config": merge_backend_configs(
+ target.attrs.get("backend_config"), backend_config
+ ),
+ }
+ )
mod = tvm.IRModule({"main": func})
with target:
- mod = tvm.compile(mod, target=target, tir_pipeline="tirx",
backend_config=backend_config)
+ mod = tvm.compile(
+ mod,
+ target=target,
+ tir_pipeline="tirx",
+ )
src = mod.mod.imports[0].inspect_source()
return src, mod
diff --git a/tests/python/tirx/test_cuda_backend_config.py
b/tests/python/tirx/test_cuda_backend_config.py
index 5dd2ed765e..a1de86d7a2 100644
--- a/tests/python/tirx/test_cuda_backend_config.py
+++ b/tests/python/tirx/test_cuda_backend_config.py
@@ -192,17 +192,21 @@ def test_entry_ir_roundtrip_and_resolution():
kernel, tvm.script.from_source(kernel.script(), extra_vars={"T": T})
)
tvm.ir.assert_structural_equal(kernel,
tvm.ir.load_json(tvm.ir.save_json(kernel)))
- target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"},
host="llvm")
- mod = tvm.tirx.transform.BindTarget(target)(tvm.IRModule({"main": kernel}))
- mod = BindBackendConfig(
+ target = tvm.target.Target(
{
- "cuda": {
- "arch": "sm_100a",
- "nvcc": ["--use_fast_math", "--generate-line-info"],
- "nvrtc": ["--use_fast_math", "--generate-line-info"],
- }
- }
- )(mod)
+ "kind": "cuda",
+ "arch": "sm_100a",
+ "backend_config": {
+ "cuda": {
+ "nvcc": ["--use_fast_math", "--generate-line-info"],
+ "nvrtc": ["--use_fast_math", "--generate-line-info"],
+ }
+ },
+ },
+ host="llvm",
+ )
+ mod = tvm.tirx.transform.BindTarget(target)(tvm.IRModule({"main": kernel}))
+ mod = BindBackendConfig()(mod)
entries = []
tvm_ffi.structural_walk(
mod,
@@ -220,14 +224,6 @@ def test_entry_ir_roundtrip_and_resolution():
assert "--generate-line-info" not in config["cuda"]["nvrtc"]
-def test_target_arch_conflict():
- from tvm.backend.config import prepare_target
-
- with pytest.raises(ValueError, match="conflicts"):
- prepare_target({"kind": "cuda", "arch": "sm_90"}, {"cuda": {"arch":
"sm_100a"}})
- assert prepare_target(None, {"cuda": {"arch": "sm_100a"}}).arch ==
"sm_100a"
-
-
@pytest.mark.gpu
@pytest.mark.skipif(not tvm.cuda().exist, reason="requires CUDA")
@pytest.mark.parametrize("compiler", ["nvrtc", "nvcc"])
@@ -247,13 +243,16 @@ def test_backend_config_controls_actual_ftz(compiler):
for ftz in (False, True):
module = tvm.compile(
kernel,
- backend_config={
- "cuda": {
- "arch": arch,
- "compiler": compiler,
- "nvcc": ["--use_fast_math", "--ftz" + "=" +
str(ftz).lower()],
- "nvrtc": ["--use_fast_math", "--ftz" + "=" +
str(ftz).lower()],
- }
+ target={
+ "kind": "cuda",
+ "backend_config": {
+ "cuda": {
+ "arch": arch,
+ "compiler": compiler,
+ "nvcc": ["--use_fast_math", "--ftz" + "=" +
str(ftz).lower()],
+ "nvrtc": ["--use_fast_math", "--ftz" + "=" +
str(ftz).lower()],
+ }
+ },
},
)
b = tvm.runtime.empty((32,), "float32", tvm.cuda())
@@ -302,7 +301,9 @@ def test_two_entries_keep_compilers_math_and_order():
C[y] = B[y] * T.float32(0.5)
arch = env.cuda_arch(0)
- module = tvm.compile(kernel, backend_config={"cuda": {"arch": arch}})
+ module = tvm.compile(
+ kernel, target={"kind": "cuda", "backend_config": {"cuda": {"arch":
arch}}}
+ )
assert len(module.mod.imports) == 2
configs = [parse_backend_config(m.inspect_source("backend_config")) for m
in module.mod.imports]
assert {c["cuda"]["compiler"] for c in configs} == {"nvcc", "nvrtc"}
@@ -365,7 +366,7 @@ def
test_serialized_config_replays_in_fresh_process(tmp_path, monkeypatch, fallb
"nvrtc": ["--use_fast_math", "--ftz=false",
"--generate-line-info"],
}
}
- module = tvm.compile(main, backend_config=config)
+ module = tvm.compile(main, target={"kind": "cuda", "backend_config":
config})
library = str(tmp_path / "replay.so")
module.export_library(library)
script = """
@@ -612,7 +613,9 @@ from tvm.script import tirx as T
target = tvm.target.Target("cuda", host="llvm")
assert "arch" not in target.attrs
assert not tvm.testing.device_enabled("cuda")
-resolved = prepare_target(target, {"cuda": BackendConfig(arch="sm_100a")})
+resolved = prepare_target({
+ **target.export(), "backend_config": {"cuda":
BackendConfig(arch="sm_100a")}
+})
assert resolved.arch == "sm_100a"
assert resolved.host.kind.name == "llvm"
@@ -685,7 +688,9 @@ def
test_source_fallback_does_not_invoke_compiler(monkeypatch):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
A[T.cuda.thread_idx("x")] = T.float32(1)
- module = tvm.compile(main, backend_config={"cuda": {"arch": "sm_100a"}})
+ module = tvm.compile(
+ main, target={"kind": "cuda", "backend_config": {"cuda": {"arch":
"sm_100a"}}}
+ )
assert "__global__" in module.mod.imports[0].inspect_source("cuda")
assert (
parse_backend_config(module.mod.imports[0].inspect_source("backend_config"))["cuda"]["arch"]
@@ -705,7 +710,9 @@ def test_identical_entries_share_one_module(monkeypatch):
A[T.cuda.thread_idx("x")] = T.float32(2)
config["cuda"]["nvrtc"].clear()
- module = tvm.compile(main, backend_config={"cuda": {"arch": "sm_100a"}})
+ module = tvm.compile(
+ main, target={"kind": "cuda", "backend_config": {"cuda": {"arch":
"sm_100a"}}}
+ )
assert len(module.mod.imports) == 1
resolved =
parse_backend_config(module.mod.imports[0].inspect_source("backend_config"))
assert resolved["cuda"]["nvrtc"] == ["--ftz=false"]
@@ -726,7 +733,7 @@ def test_build_snapshots_config_before_passes(monkeypatch):
config["cuda"]["nvrtc"] = ["--ftz=true"]
with tvm.transform.PassContext(instruments=[ChangeCallerConfig()]):
- module = tvm.compile(main, backend_config=config)
+ module = tvm.compile(main, target={"kind": "cuda", "backend_config":
config})
resolved =
parse_backend_config(module.mod.imports[0].inspect_source("backend_config"))
assert config["cuda"]["nvrtc"] == ["--ftz=true"]
assert resolved["cuda"]["nvrtc"] == ["--ftz=false"]
diff --git a/tests/python/tirx/test_cuda_launch.py
b/tests/python/tirx/test_cuda_launch.py
index 4eda010109..70bbc25236 100644
--- a/tests/python/tirx/test_cuda_launch.py
+++ b/tests/python/tirx/test_cuda_launch.py
@@ -327,4 +327,4 @@ def
test_static_shared_memory_override_cannot_shrink_allocation():
scratch[0] = T.cuda.thread_idx("x")
with pytest.raises(ValueError, match="smaller than"):
- tvm.compile(kernel, backend_config={"cuda": {"arch": "sm_80"}})
+ tvm.compile(kernel, target={"kind": "cuda", "backend_config": {"cuda":
{"arch": "sm_80"}}})
