You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/11 02:41:53 UTC
(tvm) branch main updated: refactor(infra)!: unify CUDA backend configuration (#20645)
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 0f861ee568 refactor(infra)!: unify CUDA backend configuration (#20645)
0f861ee568 is described below
commit 0f861ee568ba7ecc53d49d8f59061e6110bcca65
Author: Bohan Hou <sp...@gmail.com>
AuthorDate: Sat Oct 10 22:41:45 2026 -0400
refactor(infra)!: unify CUDA backend configuration (#20645)
CUDA compiler settings need to be consistent across build defaults,
target presets, and individual device entries without maintaining a
Python field for every native compiler flag. Introduce a lightweight
per-backend mapping accepted directly by `tvm.compile(...,
backend_config=...)` and `T.device_entry(..., backend_config=...)`.
```python
from tvm.backend.cuda import BackendConfig
cuda_config = BackendConfig(
arch="sm_100a",
compiler="nvrtc",
nvrtc=["--use_fast_math", "--ftz=false"],
)
executable = tvm.compile(func, backend_config={"cuda": cuda_config})
```
- CUDA exposes six optional keys: `arch`, `compiler`, `target_format`,
`nvcc`, `nvrtc`, and `ptxas`. Native argv lists are forwarded to the
selected toolchain. `BackendConfig` is a TypedDict whose constructor
returns a regular dictionary.
- Resolve backend defaults, Target/tag defaults, compile overrides, and
entry overrides before architecture-sensitive lowering. Lists replace
inherited lists; an empty list clears them. Snapshot inputs and use
canonical configuration cache keys.
- Compile entries with different architectures or options in separate
existing CUDA modules. Specialize shared private device helpers before
tile dispatch and dtype legalization so each helper follows its caller's
configuration.
- Keep CUDA runtime metadata to one opaque backend configuration string
appended after the existing serialized module prefix. Remove the
versioned metadata helper and persisted compiler logs. Accept legacy
payloads, avoid invoking a compiler for source-only fallback modules,
and replay saved source independently of an unrelated current Target.
- Export `cuda_host` as ordinary C++ with embedded independently
compiled device binaries. Exported libraries require TVM FFI and CUDA,
with no device compiler required at runtime.
- Replace `CompileConfig`, compiler-policy environment variables, and
decorator-level architecture settings. Obsolete compiler-policy
environment variables are ignored. Update build, Relax, external-kernel,
and IKET consumers and document native compiler options separately from
the runtime schema.
The caller migration is in
[mlc-ai/TIRx-kernels#8](https://github.com/mlc-ai/TIRx-kernels/pull/8).
The earlier CUDA launch configuration PR #20639 is merged.
Validation on the current implementation:
- TVM C++ build and changed-file pre-commit checks passed.
- Full TIRx suite: 2,691 passed / 290 skipped on seven B200 GPUs (160
seconds). Workload import gate (274 workloads / 93 kernels) and strict
registry import gate passed.
- CUDA codegen and Relax VM build regressions: 415 passed / 4 skipped.
- Backend configuration tests: 45 passed, covering native flags and FTZ
behavior, Target/tag defaults, mixed NVCC/NVRTC entries, helper
specialization, concurrent compilation, explicit offline architecture,
source fallback, legacy payloads, fresh-process artifact replay, and
standalone C++ host export.
- TIRx-kernels unit tests: 225 passed, including configuration
forwarding, CPU preparation without CUDA initialization, mutable-input
snapshots, and cache keys across Target defaults.
---
docs/tirx/api/cuda_compile.rst | 77 +++
docs/tirx/api/cuda_compile_options.rst | 55 ++
docs/tirx/api/index.rst | 2 +
docs/tirx/native_basics/cuda/compiling.rst | 133 +++-
docs/tirx/native_basics/cuda/functions.rst | 6 +-
python/tvm/backend/config.py | 121 ++++
python/tvm/backend/cuda/__init__.py | 2 +
python/tvm/backend/cuda/backend_config.py | 160 +++++
python/tvm/backend/cuda/compiler.py | 53 ++
python/tvm/backend/cuda/host.py | 96 ++-
python/tvm/backend/cuda/iket.py | 6 +-
python/tvm/backend/cuda/script.py | 1 +
python/tvm/backend/cuda/transforms/__init__.py | 20 +-
python/tvm/driver/build_module.py | 6 +-
python/tvm/relax/vm_build.py | 14 +-
python/tvm/support/nvcc.py | 261 ++------
python/tvm/tirx/bench.py | 33 +-
python/tvm/tirx/build.py | 53 +-
python/tvm/tirx/compilation_pipeline.py | 5 +-
.../tvm/tirx/script/ir_builder/external_kernel.py | 66 +-
.../tvm/tirx/script/ir_builder/parser_protocol.py | 41 +-
src/backend/cuda/codegen/codegen_cuda.cc | 13 +-
src/backend/cuda/codegen/codegen_cuda_host.cc | 67 +-
src/backend/cuda/codegen/cuda_fallback_module.cc | 28 +-
src/backend/cuda/codegen/target_kind.cc | 9 +-
src/backend/cuda/runtime/cuda_module.cc | 102 ++-
src/backend/cuda/transforms/bind_backend_config.cc | 281 ++++++++
src/target/source/codegen_c_host.h | 4 +-
src/tirx/transform/ir_utils.h | 2 +-
src/tirx/transform/split_host_device.cc | 13 +
src/tirx/transform/tile_dispatch.cc | 12 +-
src/tirx/transform/tvm_ffi_binder.cc | 2 +-
tests/python/codegen/test_target_codegen_cuda.py | 192 +++---
tests/python/disco/test_nvshmem.py | 12 +-
.../relax/test_frontend_from_exported_program.py | 3 +-
.../python/s_tir/base/test_s_tir_texture_scope.py | 2 +-
...s_tir_transform_lower_cross_thread_reduction.py | 2 +-
tests/python/tirx/codegen/conftest.py | 7 +
tests/python/tirx/codegen/test_codegen_cuda.py | 12 +-
tests/python/tirx/codegen/test_ptx_addr.py | 3 +-
tests/python/tirx/codegen/test_ptx_dialect.py | 14 +-
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 11 +-
tests/python/tirx/iket/test_iket_profiler.py | 12 +-
tests/python/tirx/test_cuda_backend_config.py | 732 +++++++++++++++++++++
tests/python/tirx/test_cuda_launch.py | 3 +-
.../tirx/test_tirx_kernels_registry_correctness.py | 34 +-
46 files changed, 2176 insertions(+), 607 deletions(-)
diff --git a/docs/tirx/api/cuda_compile.rst b/docs/tirx/api/cuda_compile.rst
new file mode 100644
index 0000000000..dc89fb615b
--- /dev/null
+++ b/docs/tirx/api/cuda_compile.rst
@@ -0,0 +1,77 @@
+.. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+
+CUDA backend configuration
+==========================
+
+Pass a nested mapping directly to ``tvm.compile(..., backend_config=...)`` or
+``T.device_entry(..., backend_config=...)``. Each backend owns its configuration;
+CUDA currently provides these fixed keys:
+
+.. list-table:: CUDA keys
+ :header-rows: 1
+ :widths: 20 30 50
+
+ * - Key
+ - Type
+ - Meaning
+ * - ``arch``
+ - ``str``
+ - One real architecture, such as ``sm_100a``; used during lowering and compilation.
+ * - ``compiler``
+ - ``"nvcc"`` or ``"nvrtc"``
+ - CUDA frontend; defaults to NVRTC.
+ * - ``target_format``
+ - ``"ptx"``, ``"cubin"``, or ``"fatbin"``
+ - NVRTC supports PTX and cubin; NVCC also supports fatbin.
+ * - ``nvcc``
+ - ``list[str]``
+ - Native NVCC arguments; defaults to ``["--use_fast_math"]``.
+ * - ``nvrtc``
+ - ``list[str]``
+ - Native NVRTC arguments; defaults to ``["--use_fast_math"]``.
+ * - ``ptxas``
+ - ``list[str]``
+ - Native assembler arguments forwarded through the selected frontend.
+ Defaults to ``["-v", "--warn-on-local-memory-usage", "--register-usage-level=10"]``.
+
+``BackendConfig`` is a ``TypedDict`` for completion and static typing. Its
+constructor returns an ordinary dictionary; all keys are optional. Compilation
+and entry construction validate and snapshot the supplied mapping. Native
+compiler arguments are not modeled as separate Python fields.
+
+.. code-block:: python
+
+ from tvm.backend.cuda import BackendConfig
+
+ cuda_config = BackendConfig(
+ arch="sm_100a", compiler="nvrtc",
+ nvrtc=["--use_fast_math", "--ftz=false"],
+ )
+ executable = tvm.compile(func, backend_config={"cuda": cuda_config})
+
+.. autoclass:: tvm.backend.cuda.BackendConfig
+ :members:
+ :undoc-members:
+
+.. autoclass:: tvm.backend.cuda.compiler.CompilationResult
+ :members:
+
+.. autofunction:: tvm.backend.cuda.compiler.compile_source
+
+See :doc:`../native_basics/cuda/compiling` for target-tag defaults, inheritance,
+per-entry compilation, and artifact behavior. See :doc:`cuda_compile_options`
+for native NVCC/NVRTC/ptxas argument usage and official compiler references.
diff --git a/docs/tirx/api/cuda_compile_options.rst b/docs/tirx/api/cuda_compile_options.rst
new file mode 100644
index 0000000000..184fd54340
--- /dev/null
+++ b/docs/tirx/api/cuda_compile_options.rst
@@ -0,0 +1,55 @@
+.. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Native CUDA compiler arguments
+==============================
+
+CUDA uses six fixed configuration keys; see :doc:`cuda_compile`. Pass other
+compiler settings as native arguments in the ``nvcc``, ``nvrtc``, or ``ptxas``
+lists:
+
+.. code-block:: python
+
+ backend_config = {"cuda": {
+ "arch": "sm_100a",
+ "compiler": "nvcc",
+ "nvcc": ["--use_fast_math", "--ftz=false", "--std=c++20", "-I/include"],
+ "nvrtc": ["--use_fast_math", "--ftz=false", "--std=c++20", "-I/include"],
+ "ptxas": ["-O3", "--register-usage-level=10"],
+ }}
+
+Each list contains argv items, without shell parsing. The selected frontend
+receives its own list and the forwarded ``ptxas`` list. Each list replaces the
+inherited list completely: ``nvrtc=[]`` clears the default fast-math flag and
+``ptxas=[]`` clears the default assembler arguments.
+
+TVM validates fixed keys and types. Architecture, output format, and toolchain
+routing are controlled by the corresponding configuration keys and the compiler
+adapter. Native flags that override those controls are reserved. The selected
+compiler validates other arguments, including their values and version support;
+new compiler flags can be used directly through these lists.
+
+TVM adds the integration arguments needed for CUDA headers, architecture, output
+retrieval, and NVSHMEM device linking. Supported outputs are PTX, cubin, and NVCC
+fatbin. Additional artifact or linking workflows require backend support beyond
+passing their compiler flags.
+
+Consult the `NVCC compiler options`_ and `NVRTC compilation options`_ for the
+installed toolkit version. Use ``ptxas --help`` to inspect its assembler options.
+
+.. _NVCC compiler options: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html
+.. _NVRTC compilation options: https://docs.nvidia.com/cuda/nvrtc/index.html
diff --git a/docs/tirx/api/index.rst b/docs/tirx/api/index.rst
index eb77ae7e76..5a47e0c1cc 100644
--- a/docs/tirx/api/index.rst
+++ b/docs/tirx/api/index.rst
@@ -38,6 +38,8 @@ Target APIs
:maxdepth: 1
cuda
+ cuda_compile
+ cuda_compile_options
cuda_launch
ptx
trainium
diff --git a/docs/tirx/native_basics/cuda/compiling.rst b/docs/tirx/native_basics/cuda/compiling.rst
index 2477b8b608..68b56b3deb 100644
--- a/docs/tirx/native_basics/cuda/compiling.rst
+++ b/docs/tirx/native_basics/cuda/compiling.rst
@@ -18,23 +18,126 @@
Compiling and inspecting
========================
-Wrap the ``Function`` in an ``IRModule`` and compile with
-``tvm.compile(mod, target=..., tir_pipeline="tirx")``; it runs the TIRx lowering
-pipeline and returns an ``Executable`` you call directly. With an active CUDA
-device, target ``"cuda"`` auto-detects its architecture (for example
-``sm_100a``). If no device is available during compilation, TVM warns and falls
-back to ``sm_50``; specify ``-arch=...`` when cross-compiling or when the emitted
-instructions require a newer architecture.
+CUDA compilation accepts a lightweight ``backend_config`` mapping. The same
+nested structure is accepted by ``tvm.compile``, ``T.device_entry``, and
+``txl.Kernel.compile``. ``tvm.backend.cuda.BackendConfig`` is an optional
+``TypedDict`` helper, also available as ``T.cuda.BackendConfig`` and
+``txl.cuda.BackendConfig``:
.. code-block:: python
- target = tvm.target.Target("cuda")
- exe = tvm.compile(tvm.IRModule({"main": scale}), target=target, tir_pipeline="tirx")
+ from tvm.backend.cuda import BackendConfig
+
+ @T.function
+ def pipeline(A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float32")):
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={"cuda": {"nvrtc": ["--use_fast_math", "--ftz=false"]}},
+ ):
+ x = T.cuda.thread_idx("x")
+ B[x] = A[x] * T.float32(0.5)
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={"cuda": {
+ "compiler": "nvcc",
+ "nvcc": ["--use_fast_math", "--generate-line-info"],
+ }},
+ ):
+ 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})
+
+Each entry resolves its configuration before architecture-sensitive lowering.
+Entries with different resolved targets or options compile as separate CUDA
+modules; their shared device helpers are copied into each group. The existing
+host module imports these device modules and launches them in program order.
+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
+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
+the input cannot change an already constructed entry or a factory cache key.
+
+Target tags use the existing target registry and carry defaults in Target attrs:
+
+.. code-block:: python
-``tir_pipeline="tirx"`` selects the TIRx lowering pipeline (tile-primitive
-dispatch and cleanup inside ``LowerTIRx`` → host/device split → finalize).
-Compiling inside a ``with target:`` block also works and lets the kernel pick
-up the target context.
+ tvm.target.tag.register_tag("local/blackwell", {
+ "kind": "cuda", "arch": "sm_100a",
+ "backend_config": {"cuda": {
+ "compiler": "nvcc", "nvcc": ["--use_fast_math", "--generate-line-info"],
+ }},
+ })
+ exe = tvm.compile(pipeline, target="local/blackwell")
+
+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.
+
+Online builds can detect the GPU architecture. Offline builds require an
+explicit architecture in the Target, build configuration, 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
+before tracing. Other backend defaults remain overridable at compilation.
+
+The default compiler is NVRTC, with fast math enabled and ptxas register usage
+level 10. NVRTC produces cubin by default, NVCC produces fatbin, and NVSHMEM
+requires cubin. Native options go into ``nvcc``, ``nvrtc``, or ``ptxas`` as
+individual argv strings; the compiler validates their values and availability.
+TVM reserves only routing, architecture, and output controls that it manages.
+Adding a compiler flag does not require adding a Python field or updating TVM.
+
+``LaunchConfig`` supplies runtime grid, block, cluster, and stream settings.
+``KernelAttributes`` describes CUDA kernel declaration attributes.
+``backend_config`` controls device compilation. PassContext still controls TVM
+passes independently; there is no wrapper combining these APIs. Only CUDA
+backend configuration is implemented currently; unsupported backend keys raise
+an explicit error.
+
+Artifacts and standalone hosts
+------------------------------
+
+CUDA modules retain the resolved configuration as one opaque JSON string,
+available through ``module.inspect_source("backend_config")``. Serialization
+keeps the existing format, function map, and code fields, then appends this string.
+It does not persist compiler logs or the complete source map. Source fallback
+artifacts retain source and configuration and compile on loading with a CUDA
+runtime. Replay uses the saved defaults independently of the loader's Target scope. Producing a source fallback does not require NVCC or NVRTC.
+Old binary artifacts remain loadable; old source artifacts without configuration
+use the default compiler settings when compiled.
+
+For ``target=Target({"kind": "cuda", "arch": "sm_100a"}, host="cuda_host")``,
+``tvm.backend.cuda.export_cuda_host(exe.mod)`` returns ordinary C++ host source
+with embedded independently compiled device binaries. Compile it with a C++
+compiler and CUDA headers, and link with tvm-ffi, cudart and the CUDA driver.
+The exported library requires no TVM installation or device compiler at runtime.
+Compiling the host does not recompile kernels.
+
+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": ...}}``.
+
+Compiler/math/ptxas environment policies remain removed. CPU benchmark
+preparation receives an explicit architecture through ``backend_config`` and
+defaults to NVCC. The runner, benchmark, and test CLIs accept
+``--backend-config '{"cuda":{"arch":"sm_100a","compiler":"nvcc"}}'``.
+Factory caches use canonical configuration snapshots; dictionary key order is
+ignored while argument order remains significant.
+
+See :doc:`../../api/cuda_compile` for the complete configuration API.
Inspecting the result
---------------------
@@ -44,8 +147,8 @@ compiled module.
.. code-block:: python
- scale.show() # pretty-print the TIRx (TVMScript)
- print(scale.script()) # ... the same, as a string
+ pipeline.show() # pretty-print the TIRx (TVMScript)
+ print(pipeline.script()) # ... the same, as a string
# the generated CUDA C source, from the compiled Executable:
print(exe.mod.imports[0].inspect_source())
diff --git a/docs/tirx/native_basics/cuda/functions.rst b/docs/tirx/native_basics/cuda/functions.rst
index 8ee5075980..5f136fc88d 100644
--- a/docs/tirx/native_basics/cuda/functions.rst
+++ b/docs/tirx/native_basics/cuda/functions.rst
@@ -288,9 +288,9 @@ both block and cluster dimensions using CUDA 13's ``__block_size__`` declaration
block size. Runtime values belong in ``LaunchConfig``.
Both host backends consume the same validated launch description. The normal
-CUDA module calls the Driver API's ``cuLaunchKernelEx``; exported ``cuda_host``
-code calls the Runtime API's ``cudaLaunchKernelEx`` and needs only CUDA and
-tvm-ffi. They share field decoding, attribute encoding, and resource setup.
+CUDA module and exported ``cuda_host`` code both call the Driver API's
+``cuLaunchKernelEx``. The exported code embeds compiled device binaries and
+needs only CUDA and tvm-ffi. They share field decoding, attribute encoding, and resource setup.
Dynamic grid sizes, stream handles, event handles, and other launch values are
host-side call operands, not device kernel parameters or expression-valued attrs.
diff --git a/python/tvm/backend/config.py b/python/tvm/backend/config.py
new file mode 100644
index 0000000000..2d263bc471
--- /dev/null
+++ b/python/tvm/backend/config.py
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Copy, merge, and serialize explicit per-backend compilation settings."""
+
+import argparse
+import json
+from collections.abc import Mapping
+
+
+def copy_backend_config(config=None):
+ """Validate and snapshot a mapping without filling any defaults."""
+ if config is None:
+ return {}
+ if not isinstance(config, Mapping):
+ raise TypeError("backend_config must be a mapping of backend names to configurations")
+ result = {}
+ for name, value in config.items():
+ if name != "cuda":
+ raise ValueError(f"backend_config does not yet support backend {name!r}")
+ from .cuda.backend_config import validate_backend_config
+
+ result[name] = validate_backend_config(value)
+ return result
+
+
+def merge_backend_configs(*configs):
+ """Overlay explicitly supplied keys; argument lists replace inherited lists."""
+ result = {}
+ for config in configs:
+ for name, value in copy_backend_config(config).items():
+ result.setdefault(name, {}).update(value)
+ return result
+
+
+def backend_config_json(config=None):
+ """Return deterministic JSON, preserving the order of toolchain arguments."""
+ return json.dumps(copy_backend_config(config), sort_keys=True, separators=(",", ":"))
+
+
+def parse_backend_config(value):
+ """Read a configuration snapshot without consulting defaults or process state."""
+ return copy_backend_config(json.loads(value) if value else None)
+
+
+def prepare_target(target, config=None, 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())
+ generic_cuda = (
+ active is None
+ or active == "cuda"
+ or (isinstance(active, Mapping) and active.get("kind") == "cuda" and "arch" not in active)
+ )
+ entries = []
+ if mod is not None and generic_cuda:
+ from tvm.ir import RegionStmt
+
+ def visit(node):
+ if isinstance(node, RegionStmt) and node.op.name == "tirx.device_entry":
+ entries.append(parse_backend_config(node.attrs.get("backend_config", "")))
+
+ tvm_ffi.structural_walk(mod, visit)
+ 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
+
+
+def argparse_backend_config(value):
+ """Parse the same nested mapping accepted by compile and device_entry."""
+ try:
+ return parse_backend_config(value)
+ except (TypeError, ValueError) as error:
+ raise argparse.ArgumentTypeError(str(error)) from error
+
+
+def add_backend_config_argument(parser):
+ """Add the shared JSON configuration argument to a command-line interface."""
+ parser.add_argument(
+ "--backend-config",
+ type=argparse_backend_config,
+ default=None,
+ metavar="JSON",
+ help='Backend settings, e.g. \'{"cuda":{"compiler":"nvcc","nvcc":["--ftz=false"]}}\'',
+ )
diff --git a/python/tvm/backend/cuda/__init__.py b/python/tvm/backend/cuda/__init__.py
index b9145aec00..60d3411e03 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -23,6 +23,7 @@ from tvm_ffi.libinfo import load_lib_ctypes
from tvm.base import _LOADED_LIBS
+from .backend_config import BackendConfig
from .host import export_cuda_host
_LAZY_SUBMODULES = {
@@ -115,6 +116,7 @@ def __getattr__(name: str):
__all__ = [
+ "BackendConfig",
"codegen",
"cpp",
"export_cuda_host",
diff --git a/python/tvm/backend/cuda/backend_config.py b/python/tvm/backend/cuda/backend_config.py
new file mode 100644
index 0000000000..9923d921cc
--- /dev/null
+++ b/python/tvm/backend/cuda/backend_config.py
@@ -0,0 +1,160 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CUDA toolchain selection and native argument lists."""
+
+import re
+from collections.abc import Mapping, Sequence
+from typing import Literal, TypedDict
+
+import tvm_ffi
+
+
+class BackendConfig(TypedDict, total=False):
+ """Optional CUDA overrides. Toolchain arguments are passed as individual argv items."""
+
+ arch: str
+ compiler: Literal["nvcc", "nvrtc"]
+ target_format: Literal["ptx", "cubin", "fatbin"]
+ nvcc: list[str]
+ nvrtc: list[str]
+ ptxas: list[str]
+
+
+# Only arguments that conflict with TVM's architecture, artifact, or tool routing
+# are reserved. General compiler options and their values belong to the compiler.
+_RESERVED = {
+ "-arch",
+ "--gpu-architecture",
+ "-gpu-architecture",
+ "--gpu-name",
+ "-arch-name",
+ "-gencode",
+ "--generate-code",
+ "-code",
+ "--gpu-code",
+ "-o",
+ "--output-file",
+ "-optf",
+ "--options-file",
+ "--cubin",
+ "-cubin",
+ "--ptx",
+ "-ptx",
+ "--fatbin",
+ "-fatbin",
+ "-Xptxas",
+ "--ptxas-options",
+}
+
+
+def validate_backend_config(config):
+ """Check fixed keys and types, and return independently owned containers."""
+ if not isinstance(config, Mapping):
+ raise TypeError("CUDA backend_config must be a mapping")
+ result = {}
+ for name, value in config.items():
+ if name not in BackendConfig.__annotations__:
+ raise ValueError(f"Unknown CUDA backend_config key: {name!r}")
+ if name in ("nvcc", "nvrtc", "ptxas"):
+ if isinstance(value, str | bytes) or not isinstance(value, Sequence):
+ raise TypeError(f"CUDA backend_config.{name} must be a sequence of strings")
+ if not all(isinstance(v, str) for v in value):
+ raise TypeError(f"CUDA backend_config.{name} must be a sequence of strings")
+ for item in value:
+ if item.split("=", 1)[0] in _RESERVED:
+ raise ValueError(
+ f"{name}: {item!r} is managed by backend_config or the compiler adapter"
+ )
+ result[name] = list(value)
+ elif not isinstance(value, str):
+ raise TypeError(f"CUDA backend_config.{name} must be a string")
+ else:
+ result[name] = value
+ if "arch" in result and not re.fullmatch(r"sm_[0-9]+[af]?", result["arch"]):
+ raise ValueError("CUDA backend_config.arch must name a real architecture, e.g. sm_100a")
+ if "compiler" in result and result["compiler"] not in ("nvcc", "nvrtc"):
+ raise ValueError("CUDA backend_config.compiler must be 'nvcc' or 'nvrtc'")
+ if "target_format" in result and result["target_format"] not in ("ptx", "cubin", "fatbin"):
+ raise ValueError("CUDA backend_config.target_format must be 'ptx', 'cubin', or 'fatbin'")
+ if result.get("compiler") == "nvrtc" and result.get("target_format") == "fatbin":
+ raise ValueError("NVRTC cannot produce fatbin; use compiler='nvcc'")
+ return result
+
+
+def default_backend_config():
+ """Return fresh CUDA defaults, preserving the existing compilation policy."""
+ return BackendConfig(
+ compiler="nvrtc",
+ nvcc=["--use_fast_math"],
+ nvrtc=["--use_fast_math"],
+ ptxas=["-v", "--warn-on-local-memory-usage", "--register-usage-level=10"],
+ )
+
+
+def resolve_backend_config(config=None, target=None):
+ """Resolve one CUDA configuration after applying target/tag defaults."""
+ from tvm.target import Target
+
+ from ..config import copy_backend_config
+
+ overrides = validate_backend_config(config or {})
+ if target is None:
+ target = Target.current()
+ if target is None:
+ target = Target(
+ {"kind": "cuda", **({"arch": overrides["arch"]} if "arch" in overrides else {})}
+ )
+ else:
+ target = Target(target)
+ if target.kind.name != "cuda":
+ raise ValueError("CUDA backend_config requires a CUDA target")
+ result = default_backend_config()
+ result.update(copy_backend_config(target.attrs.get("backend_config", {})).get("cuda", {}))
+ result.update(overrides)
+ if "arch" not in result:
+ arch = target.attrs.get("arch")
+ if arch is None:
+ raise ValueError(
+ "CUDA compilation requires an explicit backend_config arch or Target arch "
+ "when no CUDA device is available"
+ )
+ result["arch"] = str(arch)
+ return validate_backend_config(result)
+
+
+def compiler_options(config):
+ """Forward argv to the selected frontend and its ptxas stage."""
+ options = list(config[config["compiler"]])
+ if config["compiler"] == "nvcc":
+ if config["ptxas"]:
+ options.append(f"--ptxas-options={','.join(config['ptxas'])}")
+ else:
+ options.extend(f"--ptxas-options={value}" for value in config["ptxas"])
+ return options
+
+
+@tvm_ffi.register_global_func("cuda.resolve_backend_config")
+def _resolve_entry(local, defaults, target):
+ from tvm.target import Target
+
+ from ..config import backend_config_json, merge_backend_configs, parse_backend_config
+
+ merged = merge_backend_configs(parse_backend_config(defaults), parse_backend_config(local))
+ config = resolve_backend_config(merged.get("cuda"), target)
+ attrs = dict(target.export())
+ attrs["arch"] = config["arch"]
+ return [Target(attrs), backend_config_json({"cuda": config})]
diff --git a/python/tvm/backend/cuda/compiler.py b/python/tvm/backend/cuda/compiler.py
new file mode 100644
index 0000000000..f3ae39329a
--- /dev/null
+++ b/python/tvm/backend/cuda/compiler.py
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Compile one CUDA source group using explicit backend settings."""
+
+from typing import NamedTuple
+
+from ..config import copy_backend_config
+from .backend_config import compiler_options, resolve_backend_config
+
+
+class CompilationResult(NamedTuple):
+ """Device artifact and the effective settings used to produce it."""
+
+ binary: bytes
+ target_format: str
+ config: dict
+
+
+def compile_source(source: str, backend_config=None, *, target=None) -> CompilationResult:
+ """Compile source with per-backend overrides and optional Target/tag defaults.
+
+ When target is omitted, use the current CUDA Target or detect the device.
+ Compiler logs are reported on failure and are not part of the artifact.
+ """
+ from tvm.support import nvcc
+
+ config = resolve_backend_config(copy_backend_config(backend_config).get("cuda"), target)
+ use_nvshmem = "#include <nvshmem.h>" in source or "#include <nvshmemx.h>" in source
+ output_format = config.get("target_format") or (
+ "cubin" if config["compiler"] == "nvrtc" or use_nvshmem else "fatbin"
+ )
+ if use_nvshmem and output_format != "cubin":
+ raise ValueError("NVSHMEM device linking requires target_format='cubin'")
+ config["target_format"] = output_format
+ driver = getattr(nvcc, f"_compile_cuda_{config['compiler']}")
+ binary, _log = driver(
+ source, output_format, config["arch"], compiler_options(config), use_nvshmem=use_nvshmem
+ )
+ return CompilationResult(bytes(binary), output_format, {"cuda": config})
diff --git a/python/tvm/backend/cuda/host.py b/python/tvm/backend/cuda/host.py
index 65e31dc16f..e896ba2013 100644
--- a/python/tvm/backend/cuda/host.py
+++ b/python/tvm/backend/cuda/host.py
@@ -14,40 +14,26 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""CUDA host and device source assembly."""
+"""Export standalone tvm-ffi host C++ with independently compiled CUDA binaries."""
-from tvm_ffi import Module
+import json
+from tvm_ffi import Module
-def export_cuda_host(mod: Module) -> str:
- """Return one CUDA C++ translation unit from a built CUDA-host module.
+from tvm.backend.config import parse_backend_config
- Build ``mod`` with a CUDA target whose host is ``"cuda_host"``. Its CUDA
- device imports must retain their CUDA C++ source. Device definitions are
- emitted before the host wrappers that launch them. This function only
- reads the modules and returns source; compilation, loading and writing
- files are left to the caller. Sources are concatenated without rewriting
- declarations or deduplicating helpers, so the imports must be compatible
- within one translation unit and provide the types needed by NVCC's host pass.
- Generated host wrappers use tvm-ffi and CUDA libraries. TVM runtime workspace
- allocation and other unsupported runtime services are rejected during codegen.
- Tensor-map encoding additionally requires linking the CUDA driver library.
- Parameters
- ----------
- mod : tvm.runtime.Module
- The built CUDA-host source module with its CUDA device imports.
+def export_cuda_host(mod: Module) -> str:
+ """Return C++ host source with embedded device binaries.
- Returns
- -------
- source : str
- CUDA C++ source suitable for compilation with NVCC and tvm-ffi headers.
+ Build with a CUDA Target whose host is ``cuda_host``. Each device import
+ retains its own architecture and BackendConfig. This export preserves those
+ binary boundaries; compiling the returned host source does not invoke a
+ device compiler. Link the result with tvm-ffi, cudart and the CUDA driver.
+ The resulting library needs only those libraries at runtime.
- Raises
- ------
- ValueError
- If the host or its imports are incompatible, or CUDA source is absent
- (for example, after loading a device module saved as a binary).
+ Source-only fallback imports are compiled using their serialized settings.
+ Binary imports work after a save/load roundtrip, without original source.
"""
if not isinstance(mod, Module):
raise TypeError("export_cuda_host expects a runtime Module")
@@ -56,9 +42,8 @@ def export_cuda_host(mod: Module) -> str:
host_source = mod.inspect_source()
if not host_source:
raise ValueError("CUDA-host module has no source")
-
- sources = []
- visited = set()
+ binaries = []
+ visited, symbols = set(), set()
def collect(device_mod):
if device_mod in visited:
@@ -66,13 +51,54 @@ def export_cuda_host(mod: Module) -> str:
visited.add(device_mod)
if device_mod.kind != "cuda":
raise ValueError(f"Expected a CUDA device import, got {device_mod.kind!r}")
- source = device_mod.inspect_source("cuda")
- if not source:
- raise ValueError("CUDA device import has no CUDA C++ source to bundle")
for imported in device_mod.imports:
collect(imported)
- sources.append(source)
+ binary, fmt, names = device_mod["__tvm_cuda_binary"]()
+ if fmt == "cuda":
+ from .compiler import compile_source
+
+ config = device_mod.inspect_source("backend_config")
+ if not config:
+ config = "{}"
+ config = parse_backend_config(config)
+ arch = config.get("cuda", {}).get("arch")
+ target = {"kind": "cuda", "arch": arch} if arch is not None else None
+ result = compile_source(bytes(binary).decode(), config, target=target)
+ binary, fmt = result.binary, result.target_format
+ if fmt not in ("cubin", "fatbin", "ptx"):
+ raise ValueError(f"Cannot embed CUDA format {fmt!r}")
+ names = tuple(str(name) for name in names)
+ duplicate = symbols.intersection(names)
+ if duplicate:
+ raise ValueError(f"Duplicate CUDA kernel symbols: {sorted(duplicate)}")
+ symbols.update(names)
+ binaries.append((bytes(binary), names))
for imported in mod.imports:
collect(imported)
- return "\n\n".join([*sources, host_source])
+ source = [
+ "#define TVM_FFI_CUBIN_LAUNCHER_USE_DRIVER_API 1",
+ "#include <tvm/ffi/extra/cuda/cubin_launcher.h>",
+ "#include <cstring>",
+ ]
+ for index, (binary, _) in enumerate(binaries):
+ # PTX must be NUL terminated; an extra trailing byte is harmless for ELF/fatbin.
+ values = [str(value) for value in binary] + ["0"]
+ lines = [",".join(values[start : start + 32]) for start in range(0, len(values), 32)]
+ source.append(
+ f"alignas(64) static const unsigned char tvm_cuda_binary_{index}[] = {{\n"
+ + ",\n".join(lines)
+ + "\n};"
+ )
+ source.append("static CUkernel tvm_cuda_host_get_kernel(const char* name) {")
+ for index, (_, names) in enumerate(binaries):
+ match = " || ".join(f"std::strcmp(name, {json.dumps(name)}) == 0" for name in names)
+ if match:
+ source.append(
+ f" if ({match}) {{\n"
+ f" static tvm::ffi::CubinModule module(tvm_cuda_binary_{index});\n"
+ " return module.GetKernel(name).GetHandle();\n }"
+ )
+ source.append(' TVM_FFI_THROW(ValueError) << "Unknown CUDA kernel: " << name;\n}')
+ source.append(host_source)
+ return "\n\n".join(source)
diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py
index 3912adcb0d..ab148c0e20 100644
--- a/python/tvm/backend/cuda/iket.py
+++ b/python/tvm/backend/cuda/iket.py
@@ -309,14 +309,16 @@ class IketProfiler:
def sentinel_token(self, name: str):
return T.cuda.iket.sentinel_token(name)
- def compile(self, mod, target=None, *, tir_pipeline="tirx"):
+ def compile(self, mod, target=None, *, tir_pipeline="tirx", backend_config=None):
"""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)
+ executable = tvm.compile(
+ enabled_mod, target=target, tir_pipeline=tir_pipeline, backend_config=backend_config
+ )
return _OfficialIketExecutable(executable)
diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py
index 021cc6d34f..36b7e01225 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -21,6 +21,7 @@ from __future__ import annotations
from tvm import ir as _ir
from tvm.backend.cuda import op as _cuda_op
+from .backend_config import BackendConfig as BackendConfig
from .launch import AccessPolicyWindow as AccessPolicyWindow
from .launch import KernelAttributes as KernelAttributes
from .launch import LaunchCompletionEvent as LaunchCompletionEvent
diff --git a/python/tvm/backend/cuda/transforms/__init__.py b/python/tvm/backend/cuda/transforms/__init__.py
index b43403d337..3859b03ddb 100644
--- a/python/tvm/backend/cuda/transforms/__init__.py
+++ b/python/tvm/backend/cuda/transforms/__init__.py
@@ -38,4 +38,22 @@ def LowerIket():
return get_global_func("tirx.backend.cuda.transforms.LowerIket")()
-__all__ = ["LowerIket"]
+def BindBackendConfig(backend_config=None):
+ """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)
+
+
+def SpecializeDeviceHelpers():
+ """Clone private device helpers into each independently compiled CUDA group."""
+ return get_global_func("tirx.backend.cuda.transforms.SpecializeDeviceHelpers")()
+
+
+def SpecializeEntryHelpers():
+ """Specialize shared helper calls before target-sensitive tile dispatch."""
+ return get_global_func("tirx.backend.cuda.transforms.SpecializeEntryHelpers")()
+
+
+__all__ = ["BindBackendConfig", "LowerIket", "SpecializeDeviceHelpers", "SpecializeEntryHelpers"]
diff --git a/python/tvm/driver/build_module.py b/python/tvm/driver/build_module.py
index bbf84813bc..c1296e439f 100644
--- a/python/tvm/driver/build_module.py
+++ b/python/tvm/driver/build_module.py
@@ -42,6 +42,7 @@ 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.
@@ -61,6 +62,8 @@ def compile( # pylint: disable=redefined-builtin
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
-------
@@ -74,6 +77,7 @@ 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)
+ lib = tvm.tirx.build(mod, target, pipeline=tir_pipeline, backend_config=backend_config)
return Executable(lib)
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index ab1d0c0d34..bd2bfc60ae 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -103,6 +103,7 @@ def _vmlink(
params: dict[str, list] | None = None,
*,
system_lib: bool | None = None,
+ backend_config=None,
):
"""
Internal codegen function to make executable.
@@ -146,7 +147,9 @@ 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)
+ lib = tvm.tirx.build(
+ tir_mod, target=target, pipeline=tir_pipeline, backend_config=backend_config
+ )
for ext_mod in ext_libs:
if _is_device_module(ext_mod):
tir_ext_libs.append(ext_mod)
@@ -169,6 +172,7 @@ 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.
@@ -202,6 +206,9 @@ 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
@@ -229,6 +236,10 @@ def build(
constants = attrs.get("const_name_to_constant", {})
return ext_libs, constants
+ from tvm.backend.config import copy_backend_config, prepare_target
+
+ backend_config = copy_backend_config(backend_config)
+ target = prepare_target(target, backend_config, mod)
if isinstance(target, str):
target = tvm.target.Target(target)
if not params:
@@ -269,6 +280,7 @@ def build(
ext_libs=ext_libs,
params=params,
system_lib=system_lib,
+ backend_config=backend_config,
)
diff --git a/python/tvm/support/nvcc.py b/python/tvm/support/nvcc.py
index cdb602629c..33c57d3fec 100644
--- a/python/tvm/support/nvcc.py
+++ b/python/tvm/support/nvcc.py
@@ -20,87 +20,31 @@
import glob
import os
import platform
-import shlex
import subprocess
import warnings
import tvm_ffi
import tvm
+from tvm.backend.config import parse_backend_config
from tvm.target import Target
from . import utils
-def _ptxas_option_flags():
- """Return ptxas flags forwarded via ``--ptxas-options`` (without the prefix).
+def compile_cuda(code, *, backend_config=None, path_target=None):
+ """Compile source using the nested backend_config mapping and return its bytes.
- Environment Variables
- ---------------------
- TVM_CUDA_PTXAS_REG_LEVEL : str
- ptxas ``--register-usage-level`` (default ``10``).
- TVM_CUDA_PTXAS_EXTRA_OPTS : str
- Extra ptxas flags, shell-tokenized (e.g. ``"-O1"`` or ``"-O2 --def-load-cache=ca"``).
- Each token becomes its own ``--ptxas-options=<token>`` entry for NVRTC, or is
- comma-joined for nvcc.
+ For the binary format and resolved configuration, use
+ :func:`tvm.backend.cuda.compiler.compile_source` directly.
"""
- flags = [
- "-v",
- f"--register-usage-level={os.environ.get('TVM_CUDA_PTXAS_REG_LEVEL', '10')}",
- "--warn-on-local-memory-usage",
- ]
- extra = os.environ.get("TVM_CUDA_PTXAS_EXTRA_OPTS", "").strip()
- if extra:
- flags.extend(shlex.split(extra))
- return flags
-
-
-def compile_cuda(
- code, target_format=None, arch=None, options=None, path_target=None, compiler="nvrtc"
-):
- """Compile CUDA code with NVCC or NVRTC.
-
- Parameters
- ----------
- code : str
- The CUDA code.
-
- target_format : str
- The target format of the compiler ("ptx", "cubin", or "fatbin").
-
- arch : str
- The CUDA architecture.
-
- options : str or list of str
- The additional options.
-
- path_target : str, optional
- Output file.
+ from tvm.backend.cuda.compiler import compile_source
- compiler : str, optional
- Compiler backend: "nvrtc" (default) or "nvcc".
- This can be set by the TVM_CUDA_COMPILE_MODE environment variable.
-
- Returns
- -------
- res_binary : bytearray
- The bytearray of the compiled binary (ptx/cubin/fatbin).
-
- Notes
- -----
- - NVRTC is a "runtime" compilation library and can be faster for JIT compilation.
- - NVRTC requires cuda-bindings: pip install cuda-bindings
- """
- use_nvshmem = "#include <nvshmem.h>" in code or "#include <nvshmemx.h>" in code
-
- if compiler == "nvcc":
- result = _compile_cuda_nvcc(code, target_format, arch, options, path_target, use_nvshmem)
- elif compiler == "nvrtc":
- result = _compile_cuda_nvrtc(code, target_format, arch, options, path_target, use_nvshmem)
- else:
- raise ValueError(f"CUDA compiler must be 'nvcc' or 'nvrtc', got: {compiler}")
-
- return result
+ result = compile_source(code, backend_config)
+ if path_target:
+ with open(path_target, "wb") as output:
+ output.write(result.binary)
+ return bytearray(result.binary)
def _compile_cuda_nvcc(
@@ -142,50 +86,13 @@ def _compile_cuda_nvcc(
target_format = "cubin"
nvshmem_include_path, nvshmem_lib_path = find_nvshmem_paths()
- if arch is None:
- # If None, then it will use `tvm.target.Target.current().arch`.
- # Target arch could be a str like "sm_xx", or a list, such as
- # [
- # "-gencode", "arch=compute_52,code=sm_52",
- # "-gencode", "arch=compute_70,code=sm_70"
- # ]
- target = Target.current(allow_none=True)
- target_arch = getattr(target, "arch", None) if target is not None else None
- if isinstance(target_arch, str) and target_arch.startswith("sm_"):
- suffix = target_arch[3:]
- else:
- suffix = "".join(get_target_compute_version(target).split("."))
- arch = ["-gencode", f"arch=compute_{suffix},code=sm_{suffix}"]
-
temp = utils.tempdir()
file_name = "tvm_kernels"
- if target_format is None and not use_nvshmem:
- target_format = "ptx"
-
- tvm_kernel_dump = os.environ.get("TVM_KERNEL_DUMP", None)
- if tvm_kernel_dump is not None:
- target_format = "fatbin" # use fatbin to get cubin for SASS extraction
-
if target_format not in ["cubin", "ptx", "fatbin"]:
raise ValueError("target_format must be in cubin, ptx, fatbin")
temp_code = temp.relpath(f"{file_name}.cu")
temp_target = temp.relpath(f"{file_name}.{target_format}")
- pass_context = tvm_ffi.get_global_func("transform.GetCurrentPassContext")()
- kernels_output_dir = (
- pass_context.config["cuda.kernels_output_dir"]
- if "cuda.kernels_output_dir" in pass_context.config
- else None
- )
- if tvm_kernel_dump is not None:
- kernels_output_dir = tvm_kernel_dump
-
- if kernels_output_dir is not None:
- if not os.path.isdir(kernels_output_dir):
- os.makedirs(kernels_output_dir)
- temp_code = os.path.join(kernels_output_dir, f"{file_name}.cu")
- temp_target = os.path.join(kernels_output_dir, f"{file_name}.{target_format}")
-
with open(temp_code, "w") as out_file:
out_file.write(code)
@@ -196,12 +103,6 @@ def _compile_cuda_nvcc(
cmd = ["nvcc"]
cmd += [f"--{target_format}", "-O3"]
- if tvm_kernel_dump is not None:
- cmd += ["-lineinfo"]
- cmd += ["--keep", f"--keep-dir={tvm_kernel_dump}"]
- if os.environ.get("TVM_KERNEL_DEBUG", "0") == "1":
- cmd += ["-g"]
- cmd += ["-G"]
if isinstance(arch, list):
cmd += arch
elif isinstance(arch, str):
@@ -216,12 +117,8 @@ def _compile_cuda_nvcc(
"-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
"--expt-relaxed-constexpr",
"--expt-extended-lambda",
- *([] if os.environ.get("TVM_CUDA_NVCC_NO_FAST_MATH") else ["--use_fast_math"]),
- f"--ptxas-options={','.join(_ptxas_option_flags())}",
]
- major, _ = parse_compute_version(get_target_compute_version(Target.current(allow_none=True)))
-
if options:
if isinstance(options, str):
cmd += [options]
@@ -258,12 +155,7 @@ def _compile_cuda_nvcc(
# Second stage for NVSHMEM
if use_nvshmem:
- target = Target.current(allow_none=True)
- target_arch = getattr(target, "arch", None) if target is not None else None
- if isinstance(target_arch, str) and target_arch.startswith("sm_"):
- compute_version = target_arch[3:]
- else:
- compute_version = "".join(get_target_compute_version(target).split("."))
+ compute_version = arch.removeprefix("sm_")
cmd = ["nvlink"]
cmd += [f"-arch=sm_{compute_version}"]
cmd += ["-L", nvshmem_lib_path]
@@ -289,7 +181,7 @@ def _compile_cuda_nvcc(
data = bytearray(f.read())
if not data:
raise RuntimeError("Compilation error: empty result is generated")
- return data
+ return data, out.decode("utf-8", errors="replace")
def _find_cuda_target_include(cuda_path):
@@ -489,6 +381,13 @@ namespace std {
b"--device-int128",
]
+ # CUDA 12.9+ enables a driver-backed cache and calls cuInit implicitly.
+ # Keep compilation driver-independent. Its cubin cache also reused an
+ # artifact across FTZ changes in CUDA 13.2. Keep different compiler
+ # configurations independent of that cache.
+ status, major, minor = nvrtc.nvrtcVersion()
+ if status == nvrtc.nvrtcResult.NVRTC_SUCCESS and (major, minor) >= (12, 9):
+ compile_opts.append(b"--no-cache")
if use_nvshmem:
compile_opts.extend([b"-rdc", b"true"])
@@ -568,63 +467,11 @@ namespace std {
b"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
b"-U__CUDA_NO_BFLOAT162_OPERATORS__",
b"-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
- b"--use_fast_math",
]
)
- # Mirror the nvcc path's ptxas options. register-usage-level drives ptxas
- # register allocation / instruction scheduling and is perf-relevant (FA4 was
- # tuned around it, hence the env-driven default); -v and
- # --warn-on-local-memory-usage are diagnostic. NVRTC rejects -O3 and
- # --register-usage-level as top-level flags but forwards them to its internal
- # ptxas via --ptxas-options (ptxas already defaults to -O3). NB: unlike nvcc,
- # NVRTC does not comma-split --ptxas-options, so each ptxas flag must be its
- # own entry. The nvcc-only --expt-relaxed-constexpr / --expt-extended-lambda
- # have no NVRTC equivalent and are intentionally not mirrored.
- for flag in _ptxas_option_flags():
- compile_opts.append(f"--ptxas-options={flag}".encode())
-
- # Extra NVRTC frontend flags (shell-tokenized), appended after all built-in
- # defaults so they can override them (e.g. TVM_CUDA_NVRTC_EXTRA_OPTS="--ftz=false"
- # to undo the -ftz=true implied by --use_fast_math).
- nvrtc_extra = os.environ.get("TVM_CUDA_NVRTC_EXTRA_OPTS", "").strip()
- if nvrtc_extra:
- compile_opts.extend(t.encode() for t in shlex.split(nvrtc_extra))
-
- # Add user-provided options, filtering out nvcc-specific flags that nvrtc doesn't support
- if options:
- nvcc_only_prefixes = (
- "-c",
- "-O",
- "-std",
- "--std",
- "-Xcompiler",
- "-Xlinker",
- "-Xarchive",
- "-Xcudafe",
- "-Xptxas",
- "--compile",
- "--compiler-options",
- "--linker-options",
- "-fPIC",
- "-shared",
- "-o",
- )
- if isinstance(options, str):
- options = [options]
- for opt in options:
- if isinstance(opt, str):
- opt_str = opt
- elif isinstance(opt, bytes):
- opt_str = opt.decode()
- else:
- opt_str = str(opt)
- skip = any(
- opt_str.startswith(prefix) or opt_str == prefix for prefix in nvcc_only_prefixes
- )
- if skip:
- continue
- compile_opts.append(opt.encode() if isinstance(opt, str) else opt)
+ # Options have already been validated and translated for this backend.
+ compile_opts.extend(option.encode() for option in (options or ()))
# Compile
(result,) = nvrtc.nvrtcCompileProgram(prog, len(compile_opts), compile_opts)
@@ -644,6 +491,12 @@ namespace std {
nvrtc.nvrtcDestroyProgram(prog)
raise RuntimeError(error_msg)
+ result_log, log_size = nvrtc.nvrtcGetProgramLogSize(prog)
+ log_buf = bytearray(log_size)
+ if result_log == nvrtc.nvrtcResult.NVRTC_SUCCESS and log_size:
+ nvrtc.nvrtcGetProgramLog(prog, log_buf)
+ log = log_buf.decode("utf-8", errors="replace").rstrip("\0")
+
# Get compiled binary
if target_format == "cubin":
result, binary_size = nvrtc.nvrtcGetCUBINSize(prog)
@@ -676,7 +529,7 @@ namespace std {
if path_target:
with open(path_target, "wb") as f:
f.write(binary_buf)
- return binary_buf
+ return binary_buf, log
def _link_nvshmem_nvrtc(binary_buf, nvshmem_lib_path):
@@ -907,51 +760,17 @@ def find_nvshmem_paths() -> tuple[str, str]:
@tvm_ffi.register_global_func
-def tvm_callback_cuda_compile(code):
- """
- Compile CUDA code using the configured backend (nvcc or nvrtc).
-
- This callback is invoked by TVM's C++ backend during CUDA module compilation.
- By default, uses nvrtc to generate cubin. The current target is fetched
- inside the callback (via ``tvm.target.Target.current(allow_none=True)``)
- so the caller does not need to push/pop a target scope around the
- invocation.
-
- Environment Variables
- ---------------------
- TVM_CUDA_COMPILE_MODE : str
- Compiler backend: "nvrtc" (default) or "nvcc"
- - "nvrtc": Use NVRTC via cuda-bindings for faster JIT, generates cubin
- - "nvcc": Use nvcc subprocess, generates fatbin
- TVM_KERNEL_DUMP : str
- If set, dump generated CUDA/intermediate files and append "-lineinfo" so profilers can
- correlate SASS back to the dumped source.
- TVM_CUDA_PTXAS_REG_LEVEL : str
- Forwarded to ptxas ``--register-usage-level`` (default ``10``).
- TVM_CUDA_PTXAS_EXTRA_OPTS : str
- Extra ptxas flags (shell-tokenized), e.g. ``"-O1"`` or ``"-O2"``.
-
- Parameters
- ----------
- code : str
- CUDA source code to compile
-
- Returns
- -------
- bytes
- Compiled binary (fatbin for nvcc, cubin for nvrtc)
- """
- # The current Target is fetched inside compile_cuda via
- # tvm.target.Target.current(allow_none=True) when arch is unset; the
- # caller no longer needs to push/pop a target scope.
- compiler = os.environ.get("TVM_CUDA_COMPILE_MODE", "nvrtc").lower()
-
- if compiler == "nvrtc":
- return compile_cuda(code, target_format="cubin", compiler="nvrtc")
- if compiler == "nvcc":
- return compile_cuda(code, target_format="fatbin", compiler="nvcc")
-
- raise ValueError(f"Invalid TVM_CUDA_COMPILE_MODE: {compiler}. Expected 'nvcc' or 'nvrtc'.")
+def tvm_callback_cuda_compile(code, config_json):
+ """Compile one device group using its serialized, resolved configuration."""
+ from tvm.backend.cuda.compiler import compile_source
+
+ config = parse_backend_config(config_json)
+ arch = config.get("cuda", {}).get("arch")
+ # A saved device group already contains its resolved defaults. Do not let
+ # an unrelated Target scope supply new defaults when replaying its source.
+ target = {"kind": "cuda", "arch": arch} if arch is not None else None
+ result = compile_source(code, config, target=target)
+ return [bytearray(result.binary), result.target_format]
@tvm_ffi.register_global_func("tvm_callback_libdevice_path")
diff --git a/python/tvm/tirx/bench.py b/python/tvm/tirx/bench.py
index 9dc425804f..ffc2686ee0 100644
--- a/python/tvm/tirx/bench.py
+++ b/python/tvm/tirx/bench.py
@@ -34,11 +34,10 @@ from enum import Enum
import numpy as np
import torch
import triton.profiler as proton
-import tvm_ffi
import tvm
+from tvm.backend.config import merge_backend_configs
from tvm.script import tirx as T
-from tvm.support import nvcc
_DISTRIBUTED_KINETO_WARMUP_ITERATIONS = 5
_DISTRIBUTED_KINETO_REPEAT_ITERATIONS = 30
@@ -80,20 +79,26 @@ def is_running_under_pytest():
def setup():
+ """Parse benchmark compile settings for explicit forwarding by the caller."""
parser = argparse.ArgumentParser()
- parser.add_argument("--dump-ptx", type=str, help="Dump PTX code to specified file")
- parser.add_argument("--dump-source", action="store_true", help="Dump source code")
- args = parser.parse_args()
-
- if args.dump_ptx:
-
- @tvm_ffi.register_global_func("tvm_callback_cuda_compile", override=True)
- def tvm_callback_cuda_compile(code, target):
- ptx = nvcc.compile_cuda(code, target_format="ptx")
- with open(args.dump_ptx, "w", encoding="utf-8") as f:
- f.write(ptx.decode())
- return ptx
+ parser.add_argument("--arch", help="CUDA architecture, for example sm_100a")
+ parser.add_argument("--compiler", choices=("nvcc", "nvrtc"))
+ parser.add_argument("--target-format", choices=("ptx", "cubin", "fatbin"))
+ parser.add_argument("--dump-source", action="store_true", help="Print generated source")
+ from tvm.backend.config import add_backend_config_argument
+ add_backend_config_argument(parser)
+ args = parser.parse_args()
+ overrides = {
+ key: value
+ for key, value in {
+ "arch": args.arch,
+ "compiler": args.compiler,
+ "target_format": args.target_format,
+ }.items()
+ if value is not None
+ }
+ args.backend_config = merge_backend_configs(args.backend_config, {"cuda": overrides})
return args
diff --git a/python/tvm/tirx/build.py b/python/tvm/tirx/build.py
index 27d79a1937..31b9d17c36 100644
--- a/python/tvm/tirx/build.py
+++ b/python/tvm/tirx/build.py
@@ -27,7 +27,7 @@ from tvm.target import Target
from tvm.tirx import Function
-def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[Target, IRModule]]:
+def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[tuple[Target, str], IRModule]]:
"""Split an IRModule into host and device modules.
This function takes an IRModule containing functions with different target attributes
@@ -46,9 +46,9 @@ def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[Target, IRModu
-------
host_mod : tvm.IRModule
The module containing host functions (CPU-targeted functions)
- device_mod_dict : Dict[Target, tvm.IRModule]
- A dict mapping targets to device modules. Each device module contains
- functions targeting the same device (e.g., CUDA GPU, OpenCL, etc.)
+ device_mod_dict : Dict[Tuple[Target, str], tvm.IRModule]
+ A dict mapping (target, serialized compile configuration) to device modules.
+ Functions in each module share both target and compiler settings.
Examples
--------
@@ -106,18 +106,23 @@ def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[Target, IRModu
host_mod = tvm.tirx.transform.Filter(is_host_func)(mod)
device_mod = tvm.tirx.transform.Filter(lambda f: not is_host_func(f))(mod)
- # TODO(syfeng): Here we use str as key since target hash is not correct
- target_str2target = {}
- device_func_dict = {}
- device_mod_dict: dict[Target, IRModule] = {}
+ # Options participate in grouping even when two kernels share a Target.
+ groups = {}
for gv, func in device_mod.functions.items():
- target = func.attrs.get("target", None)
- target_str = str(target) if target is not None else ""
- target_str2target[target_str] = target # This might be overridden by the last one
- device_func_dict.setdefault(target_str, dict()).update({gv: func})
- for target_str in target_str2target.keys():
- target = target_str2target[target_str]
- device_mod_dict[target] = tvm.IRModule(device_func_dict[target_str], attrs=device_mod.attrs)
+ target = func.attrs["target"]
+ config = func.attrs.get("backend_config", "")
+ key = (str(target), config)
+ if key not in groups:
+ groups[key] = (target, {})
+ groups[key][1][gv] = func
+ device_mod_dict = {}
+ for (_, config), (target, functions) in groups.items():
+ group = tvm.IRModule(
+ functions, attrs=device_mod.attrs, global_infos=device_mod.global_infos
+ )
+ if config:
+ group = group.with_attr("backend_config", config)
+ device_mod_dict[(target, config)] = group
return host_mod, device_mod_dict
@@ -133,9 +138,9 @@ def codegen_build(mod: IRModule, target: Target) -> tvm.runtime.Module:
def tir_to_runtime(
- host_mod: IRModule, device_mod_dict: dict[Target, IRModule], target_host: Target
+ host_mod: IRModule, device_mod_dict: dict[tuple[Target, str], IRModule], target_host: Target
):
- """Convert a collection of TIR IRModules (keyed by Target) into a single runtime Module."""
+ """Compile each device group and attach it to the host runtime module."""
# Get the first module to get the attributes
# necessary for tests/python/codegen/test_target_codegen_blob.py::test_cuda_multi_lib
@@ -143,7 +148,7 @@ def tir_to_runtime(
mhost_all.update(host_mod)
device_modules = []
- for target, device_mod in device_mod_dict.items():
+ for (target, _config), device_mod in device_mod_dict.items():
if len(device_mod.functions) != 0:
device_modules.append(codegen_build(device_mod, target))
@@ -158,6 +163,8 @@ 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.
@@ -170,6 +177,8 @@ def build(
The target for compilation.
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
-------
@@ -182,6 +191,12 @@ def build(
else:
assert isinstance(mod, tvm.IRModule)
+ from tvm.backend.config import copy_backend_config, prepare_target
+ from tvm.backend.cuda.transforms import BindBackendConfig, SpecializeEntryHelpers
+
+ backend_config = copy_backend_config(backend_config)
+ target = prepare_target(target, backend_config, mod)
+
# Step 0: Determine the target in environment
# It's used to bind the Function without target attr to serve as a default target
target_to_bind = Target.current() if target is None else target
@@ -213,6 +228,8 @@ 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 = SpecializeEntryHelpers()(mod)
# Step 4: Apply the tirx pipeline
if pipeline not in (None, "default"):
diff --git a/python/tvm/tirx/compilation_pipeline.py b/python/tvm/tirx/compilation_pipeline.py
index 276755d627..9dd3652d38 100644
--- a/python/tvm/tirx/compilation_pipeline.py
+++ b/python/tvm/tirx/compilation_pipeline.py
@@ -41,7 +41,6 @@ def tirx_pipeline(*, prepare_only=False):
tirx.transform.StmtSimplify(),
tirx.transform.LowerTIRxOpaque(),
tirx.transform.FlattenBuffer(),
- tirx.transform.BF16ComputeLegalize(),
tirx.transform.NarrowDataType(32),
tirx.transform.VectorizeLoop(not bool(config.get("tirx.disable_vectorize", False))),
tirx.transform.UnrollLoop(),
@@ -51,7 +50,6 @@ def tirx_pipeline(*, prepare_only=False):
passes.append(tirx.transform.CommonSubexprElim())
passes.extend(
[
- tirx.transform.FP8ComputeLegalize(),
tirx.transform.VerifyMemory(),
tirx.transform.LowerThreadAllreduce(),
]
@@ -69,6 +67,9 @@ def finalize_tir_pipeline():
tirx.transform.AnnotateEntryFunc(),
tirx.transform.SplitHostDevice(),
cuda_transforms.LowerIket(),
+ cuda_transforms.SpecializeDeviceHelpers(),
+ tirx.transform.BF16ComputeLegalize(),
+ tirx.transform.FP8ComputeLegalize(),
tirx.transform.MakePackedAPI(),
tirx.transform.FP8StorageLegalize(),
tirx.transform.BF16StorageLegalize(),
diff --git a/python/tvm/tirx/script/ir_builder/external_kernel.py b/python/tvm/tirx/script/ir_builder/external_kernel.py
index 3b4abbe7a9..681978b97b 100644
--- a/python/tvm/tirx/script/ir_builder/external_kernel.py
+++ b/python/tvm/tirx/script/ir_builder/external_kernel.py
@@ -19,8 +19,6 @@
import json
import logging
-import os
-import tempfile
from pathlib import Path
from typing import Any
@@ -28,9 +26,9 @@ import tvm_ffi
from tvm import __version__ as tvm_version
from tvm import libinfo, tirx
+from tvm.backend.config import backend_config_json, merge_backend_configs
from tvm.ir import Expr, PointerType, const, is_prim_expr
from tvm.runtime import Module
-from tvm.support import nvcc
class BaseKernel: # pylint: disable=too-few-public-methods
@@ -64,7 +62,13 @@ class BaseKernel: # pylint: disable=too-few-public-methods
return tvm_metadata
def _create_cuda_module(
- self, binary_data, kernel_arg_types, launch_param_tags, kernel_name, fmt="ptx"
+ self,
+ binary_data,
+ kernel_arg_types,
+ launch_param_tags,
+ kernel_name,
+ fmt="ptx",
+ backend_config=None,
):
"""
Create a CUDA module from compiled binary (PTX or cubin) and metadata.
@@ -104,7 +108,8 @@ class BaseKernel: # pylint: disable=too-few-public-methods
load_meta = tvm_ffi.get_global_func("runtime.LoadMetaDataFromJSON")
fmap = load_meta(tvm_metadata)
create_cuda = tvm_ffi.get_global_func("ffi.Module.create.cuda")
- kernel_module = create_cuda(binary_bytes, fmt, fmap, {})
+ source = {"backend_config": backend_config_json(backend_config)} if backend_config else {}
+ kernel_module = create_cuda(binary_bytes, fmt, fmap, source)
return kernel_module
@@ -148,12 +153,9 @@ class SourceKernel(BaseKernel): # pylint: disable=too-few-public-methods
Path(tvm_ffi.libinfo.find_include_path()),
Path(tvm_ffi.libinfo.find_dlpack_include_path()),
]
- compile_options = []
for include_path in dict.fromkeys(include_paths):
- assert include_path.exists(), f"Not found: {include_path!s}"
- assert include_path.is_dir(), f"Not a directory: {include_path!s}"
- compile_options += ["-I", str(include_path)]
- compile_options += ["-c", "-O3", "-std=c++17", "-Xcompiler=-fPIC"]
+ if not include_path.is_dir():
+ raise ValueError(f"Not an include directory: {include_path}")
source_code = self.source_code
try:
source_path = Path(source_code)
@@ -163,33 +165,25 @@ class SourceKernel(BaseKernel): # pylint: disable=too-few-public-methods
except: # pylint: disable=bare-except
pass
- with tempfile.TemporaryDirectory() as temp_dir:
- # Check if NVSHMEM is used - requires cubin output for device library linking
- use_nvshmem = (
- "#include <nvshmem.h>" in source_code or "#include <nvshmemx.h>" in source_code
- )
- target_format = "cubin" if use_nvshmem else "ptx"
- output_path = f"{temp_dir}/{kernel_name}.{target_format}"
-
- compiler = os.environ.get("TVM_CUDA_COMPILE_MODE", "nvrtc")
- nvcc.compile_cuda(
- source_code,
- target_format=target_format,
- options=compile_options,
- path_target=output_path,
- compiler=compiler,
- )
-
- if target_format == "ptx":
- with open(output_path) as f:
- binary_data = f.read()
- else:
- with open(output_path, "rb") as f:
- binary_data = f.read()
+ from tvm.backend.cuda.compiler import compile_source
- kernel_module = self._create_cuda_module(
- binary_data, kernel_arg_types, launch_param_tags, kernel_name, fmt=target_format
- )
+ flags = ["--use_fast_math", "--std=c++17"]
+ config = merge_backend_configs(
+ {"cuda": {"nvcc": flags, "nvrtc": flags}}, kwargs.get("backend_config")
+ )
+ # These headers are part of the external-kernel adapter, independent of
+ # whether the caller replaces the native toolchain argument lists.
+ for compiler in ("nvcc", "nvrtc"):
+ config["cuda"][compiler].extend(f"-I{p}" for p in include_paths)
+ result = compile_source(source_code, config)
+ kernel_module = self._create_cuda_module(
+ result.binary,
+ kernel_arg_types,
+ launch_param_tags,
+ kernel_name,
+ fmt=result.target_format,
+ backend_config=result.config,
+ )
return kernel_name, kernel_module, runtime_args
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index b259e3e7a0..3e7eab240a 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -755,34 +755,39 @@ def thread_binding(
def device_entry(
- *launch_values, launch=None, kernel_attrs=None, attrs=None, body_params=None
+ *launch_values,
+ launch=None,
+ kernel_attrs=None,
+ backend_config=None,
+ attrs=None,
+ body_params=None,
) -> frame.RegionFrame:
- """Enter a device kernel with an independent CUDA launch configuration.
+ """Enter a device kernel with independent CUDA launch and compile settings.
- CUDA entries use ``LaunchConfig(grid=..., block=...)`` and optional
- ``KernelAttributes``. Configuration values are ordinary region operands, so
- host expressions remain visible to substitution and free-variable analysis.
- Canonical operands and attributes are also accepted for IR reconstruction.
- Other backends may use the argument-free device entry.
+ LaunchConfig values are ordinary region operands. BackendConfig supplies
+ static per-entry overrides. Canonical operands and attributes are accepted
+ for IR reconstruction. Other backends may use the argument-free entry.
"""
+ if launch is not None and (launch_values or attrs is not None):
+ raise ValueError("device_entry launch cannot be combined with canonical operands or attrs")
+ attributes = dict(attrs) if attrs is not None else {}
+ if backend_config is not None:
+ from tvm.backend.config import backend_config_json
+
+ if "backend_config" in attributes:
+ raise ValueError("device_entry backend_config cannot duplicate canonical attrs")
+ attributes["backend_config"] = backend_config_json(backend_config)
if launch is None:
if kernel_attrs is not None:
raise ValueError("device_entry kernel_attrs require a launch configuration")
- return region("tirx.device_entry", launch_values, attrs=attrs, body_params=body_params)
- if launch_values or attrs is not None:
- raise ValueError("device_entry launch cannot be combined with canonical operands or attrs")
+ return region("tirx.device_entry", launch_values, attrs=attributes, body_params=body_params)
from tvm.backend.cuda.launch._impl import pack_kernel_attrs, pack_launch
names, values = pack_launch(launch)
- return region(
- "tirx.device_entry",
- values,
- attrs={
- "cuda.launch_fields": names,
- "cuda.kernel_attrs": pack_kernel_attrs(kernel_attrs, launch),
- },
- body_params=body_params,
+ attributes.update(
+ {"cuda.launch_fields": names, "cuda.kernel_attrs": pack_kernel_attrs(kernel_attrs, launch)}
)
+ return region("tirx.device_entry", values, attrs=attributes, body_params=body_params)
def device_context(
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index 7618e3d027..2a95799ec8 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1776,7 +1776,7 @@ void CodeGenCUDA::PrintVecElemLoadExpr(const PrimType& t, int i, const std::stri
// CUDA codegen entry point. Generates CUDA C++ source, optionally lets a
// Python postproc hook rewrite it, and hands the source bytes off to the
// fallback-aware module factory. The factory may JIT to PTX/cubin via
-// `tvm_callback_cuda_compile` (CUDAModuleNode::JitCompileFromSource) when
+// `tvm_callback_cuda_compile` when
// USE_CUDA=ON; on USE_CUDA=OFF builds (or when TVM_COMPILE_FORCE_FALLBACK is
// set), it returns a `CUDAFallbackModuleNode` carrying the raw source for
// later cross-compile.
@@ -1813,10 +1813,19 @@ ffi::Module BuildCUDA(IRModule mod, Target target) {
// Hand off raw CUDA source to the fallback-aware factory. When the real
// CUDA runtime is registered (USE_CUDA=ON and not forced-fallback) the
- // factory invokes JitCompileFromSource via tvm_callback_cuda_compile and
+ // factory invokes tvm_callback_cuda_compile and
// builds a real CUDAModuleNode. Otherwise it stores the source in a
// CUDAFallbackModuleNode for later cross-compile.
ffi::Map<ffi::String, ffi::String> source_map;
+ auto config = mod->GetAttr<ffi::String>("backend_config");
+ if (!config.has_value()) {
+ // Direct codegen callers still resolve an explicit target, without relying
+ // on a current Target scope or compiler environment variables.
+ auto resolve = ffi::Function::GetGlobalRequired("cuda.resolve_backend_config");
+ auto resolved = resolve("", "", target).cast<ffi::Array<ffi::Any>>();
+ config = resolved[1].cast<ffi::String>();
+ }
+ source_map.Set("backend_config", config.value());
return ::tvm::target::CUDAModuleCreateWithFallback(
ffi::Bytes(code.data(), code.size()), ffi::String("cuda"), ExtractFuncInfo(mod), source_map);
}
diff --git a/src/backend/cuda/codegen/codegen_cuda_host.cc b/src/backend/cuda/codegen/codegen_cuda_host.cc
index 407fef9fae..c9ec1005ba 100644
--- a/src/backend/cuda/codegen/codegen_cuda_host.cc
+++ b/src/backend/cuda/codegen/codegen_cuda_host.cc
@@ -56,7 +56,8 @@ class CodeGenCUDAHost : public CodeGenCHost {
<< "#include <tvm/ffi/extra/c_env_api.h>\n"
<< "#include <tvm/ffi/extra/cuda/device_guard.h>\n"
<< "#include <tvm/ffi/function.h>\n"
- << "#include <cuda.h>\n#include <cuda_runtime.h>\n#include <math.h>\n"
+ << "#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cuda_fp16.h>\n#include "
+ "<cuda_bf16.h>\n#include <math.h>\n"
<< "#define TVM_DLL TVM_FFI_DLL_EXPORT\n";
InitGlobalContext();
check_error_ = name_supply_->FreshName("tvm_cuda_host_check");
@@ -70,41 +71,65 @@ class CodeGenCUDAHost : public CodeGenCHost {
<< " return -1;\n}\n";
decl_stream << tvm_cuda_launch::kSupportSource;
decl_stream << "\n#include <map>\n#include <memory>\n";
- // Convert host ABI spellings (e.g. uint16_t for bfloat16) to device types.
+ decl_stream << "static CUkernel tvm_cuda_host_get_kernel(const char* name);\n";
launch_ = name_supply_->FreshName("tvm_cuda_host_launch");
- decl_stream << "template <typename... Params, typename... Args>\n"
- << "static cudaError_t " << launch_ << R"CUDA((
- void (*kernel)(Params...), tvm_cuda_launch::LaunchValues values,
+ decl_stream << "template <typename... Args>\n"
+ << "static void " << launch_ << R"CUDA((
+ const char* name, tvm_cuda_launch::LaunchValues values,
const tvm_cuda_launch::KernelRequirements& requirements, Args... args) {
- static_assert(sizeof...(Params) == sizeof...(Args), "kernel argument count mismatch");
using namespace tvm_cuda_launch;
int device;
Check(cudaGetDevice(&device));
+ CUkernel kernel = tvm_cuda_host_get_kernel(name);
+ CUfunction function;
+ Check(cuKernelGetFunction(&function, kernel));
+ CUcontext context;
+ unsigned long long context_id;
+ Check(cuCtxGetCurrent(&context));
+ Check(cuCtxGetId(context, &context_id));
static std::mutex resource_mutex;
- static std::map<std::pair<const void*, int>, std::unique_ptr<ResourceState>> resource_states;
+ static std::map<std::pair<CUkernel, unsigned long long>, std::unique_ptr<ResourceState>> states;
ResourceState* state;
{
std::lock_guard<std::mutex> lock(resource_mutex);
- auto& slot = resource_states[{reinterpret_cast<const void*>(kernel), device}];
+ auto& slot = states[{kernel, context_id}];
if (!slot) slot = std::make_unique<ResourceState>();
state = slot.get();
}
values = Normalize(values, requirements);
- PrepareResources<RuntimeAPI>(reinterpret_cast<const void*>(kernel), device, &values, state);
- cudaLaunchConfig_t config{};
- std::array<cudaLaunchAttribute, kMaxAttributes> attributes{};
- auto stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, device));
- EncodeRuntime(values, requirements, stream, &config, &attributes);
- return cudaLaunchKernelEx(&config, kernel, ((Params)args)...);
+ PrepareResources<DriverAPI>(function, device, &values, state);
+ CUlaunchConfig config{};
+ std::array<CUlaunchAttribute, kMaxAttributes> attributes{};
+ auto stream = static_cast<CUstream>(TVMFFIEnvGetStream(kDLCUDA, device));
+ EncodeDriver(values, requirements, stream, &config, &attributes);
+ // Host locals have the device ABI types. In particular FP16/BF16 values
+ // retain their 16 bits, pointers remain pointers, and tensor maps are copied
+ // into aligned by-value objects before their addresses are packed.
+ std::array<void*, sizeof...(Args)> parameters{{static_cast<void*>(&args)...}};
+ Check(cuLaunchKernelEx(&config, function, parameters.data(), nullptr));
}
)CUDA";
}
using CodeGenCHost::PrintType;
+ void PrintType(const PrimType& type, std::ostream& os) override {
+ if (type.IsScalar() && type.MatchesElementType(kDLBfloat, 16)) {
+ os << "__nv_bfloat16";
+ } else {
+ CodeGenCHost::PrintType(type, os);
+ }
+ }
+
void PrintType(const Type& type, std::ostream& os) override {
if (type.as<tirx::TensorMapTypeNode>()) {
os << "CUtensorMap";
+ } else if (auto prim = type.as<PrimType>();
+ prim && prim->IsScalar() && prim->MatchesElementType(kDLBfloat, 16)) {
+ os << "__nv_bfloat16";
+ } else if (auto prim = type.as<PrimType>();
+ prim && prim->IsScalar() && prim->MatchesElementType(kDLFloat, 16)) {
+ os << "half";
} else {
CodeGenCHost::PrintType(type, os);
}
@@ -242,7 +267,14 @@ class CodeGenCUDAHost : public CodeGenCHost {
}
std::string name = name_supply_->FreshName("cuda_arg");
PrintIndent();
- stream << "auto " << name << " = " << value << ";\n";
+ Type type = op->args[i]->ty;
+ if (i < launch_begin) {
+ if (auto* ptr = type.as<PointerTypeNode>()) {
+ if (ptr->element_type.as<tirx::TensorMapTypeNode>()) type = ptr->element_type;
+ }
+ }
+ PrintType(type, stream);
+ stream << " " << name << " = " << value << ";\n";
arguments.push_back(name);
}
std::string values = name_supply_->FreshName("cuda_launch_values");
@@ -305,9 +337,10 @@ class CodeGenCUDAHost : public CodeGenCHost {
}
stream << "}};\n";
std::string launch =
- launch_ + "(::" + std::string(symbol) + ", " + values + ", " + requirements;
+ launch_ + "(\"" + std::string(symbol) + "\", " + values + ", " + requirements;
for (size_t i = 0; i + 1 < launch_begin; ++i) launch += ", " + arguments[i];
- CheckError(launch + ")", true);
+ PrintIndent();
+ stream << launch << ");\n";
os << "0";
}
diff --git a/src/backend/cuda/codegen/cuda_fallback_module.cc b/src/backend/cuda/codegen/cuda_fallback_module.cc
index bdf60f362a..a17b220547 100644
--- a/src/backend/cuda/codegen/cuda_fallback_module.cc
+++ b/src/backend/cuda/codegen/cuda_fallback_module.cc
@@ -21,12 +21,14 @@
* \file cuda_fallback_module.cc
* \brief CUDAFallbackModuleNode — codegen-time placeholder used when the CUDA
* runtime is not linked. Mirrors `CUDAModuleNode`'s save/load format
- * byte-for-byte; see one-way comment in `SaveToBytes` below.
+ * byte-for-byte, including the backend configuration string.
* Always compiled (independent of USE_CUDA); never registered as an
* FFI factory or loader.
*/
#include "cuda_fallback_module.h"
+#include <tvm/ffi/cast.h>
+#include <tvm/ffi/container/array.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
@@ -46,7 +48,8 @@ class CUDAFallbackModuleNode : public ffi::ModuleObj {
: code_(std::move(code)),
fmt_(std::move(fmt)),
fmap_(std::move(fmap)),
- source_(std::move(source)) {}
+ source_(std::move(source)),
+ backend_config_(source_.Get("backend_config").value_or("{}")) {}
// Mirror the real module's kind so consumers cannot distinguish at the
// kind/api layer. Saved bytes load back as a real CUDAModuleNode on a
@@ -56,6 +59,14 @@ class CUDAFallbackModuleNode : public ffi::ModuleObj {
int GetPropertyMask() const final { return ffi::Module::kBinarySerializable; }
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) final {
+ if (name == "__tvm_cuda_binary") {
+ auto self = ffi::GetRef<ffi::Module>(this);
+ return ffi::Function::FromTyped([self, this]() {
+ ffi::Array<ffi::String> names;
+ for (auto [name, info] : fmap_) names.push_back(name);
+ return ffi::Array<ffi::Any>{code_, fmt_, names};
+ });
+ }
TVM_FFI_THROW(RuntimeError)
<< "CUDA runtime is not linked into this build; cannot launch kernels. "
<< "Re-link with USE_CUDA=ON or load this module in a CUDA-equipped "
@@ -64,22 +75,18 @@ class CUDAFallbackModuleNode : public ffi::ModuleObj {
}
ffi::Bytes SaveToBytes() const final {
- // NOTE: serialization format MUST remain byte-identical to
- // CUDAModuleNode::SaveToBytes in src/runtime/cuda/cuda_module.cc (the
- // source of truth). Both produce a kind="cuda" artifact that the loader
- // (ffi.Module.load_from_bytes.cuda, registered only when USE_CUDA=ON)
- // deserializes. If the real impl's format changes, mirror the change
- // here. The dependency is one-way: this file follows; cuda_module.cc
- // does not reference this file.
+ // Keep the existing prefix; only compilation settings are appended.
std::string buffer;
support::BytesOutStream stream(&buffer);
stream.Write(fmt_);
stream.Write(fmap_);
stream.Write(code_);
+ stream.Write(backend_config_);
return ffi::Bytes(std::move(buffer));
}
ffi::String InspectSource(const ffi::String& format) const final {
+ if (format == "backend_config") return backend_config_;
if (format == fmt_) {
return ffi::String(code_.data(), code_.size());
}
@@ -104,8 +111,9 @@ class CUDAFallbackModuleNode : public ffi::ModuleObj {
ffi::String fmt_;
// function information table.
ffi::Map<ffi::String, runtime::FunctionInfo> fmap_;
- // In-memory source map for InspectSource — never serialized.
+ // In-memory source for inspection; not serialized.
ffi::Map<ffi::String, ffi::String> source_;
+ ffi::String backend_config_; // Serialized JSON; interpreted only by the compilation backend.
};
ffi::Module CUDAFallbackModuleCreate(ffi::Bytes code, ffi::String fmt,
diff --git a/src/backend/cuda/codegen/target_kind.cc b/src/backend/cuda/codegen/target_kind.cc
index 26b8a506fc..a0c784010f 100644
--- a/src/backend/cuda/codegen/target_kind.cc
+++ b/src/backend/cuda/codegen/target_kind.cc
@@ -77,11 +77,11 @@ ffi::Map<ffi::String, ffi::Any> UpdateCUDAAttrs(ffi::Map<ffi::String, ffi::Any>
int archInt;
ffi::Any version;
if (!DetectDeviceFlag({kDLCUDA, 0}, runtime::kComputeVersion, &version)) {
- LOG(WARNING) << "Unable to detect CUDA version, default to \"-arch=sm_50\" instead";
- archInt = 50;
- } else {
- archInt = std::stod(version.cast<std::string>()) * 10 + 0.1;
+ // Generic targets are also used for backend discovery and IR construction.
+ // Require an explicit architecture when resolving compilation, not here.
+ return target;
}
+ archInt = std::stod(version.cast<std::string>()) * 10 + 0.1;
if (archInt >= 90) {
target.Set("arch", ffi::String("sm_") + std::to_string(archInt) + "a");
} else {
@@ -126,6 +126,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TargetKindDef("cuda")
.set_default_device_type(kDLCUDA)
+ .def_option<ffi::Map<ffi::String, ffi::Any>>("backend_config")
.def_option<ffi::String>("mcpu")
.def_option<ffi::String>("arch")
.def_option<int64_t>("max_shared_memory_per_block")
diff --git a/src/backend/cuda/runtime/cuda_module.cc b/src/backend/cuda/runtime/cuda_module.cc
index 4cd7ebf9c2..4cdf0cdd17 100644
--- a/src/backend/cuda/runtime/cuda_module.cc
+++ b/src/backend/cuda/runtime/cuda_module.cc
@@ -27,6 +27,7 @@
#include <cuda.h>
#include <cuda_runtime.h>
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/container/array.h>
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/extra/cuda/base.h>
#include <tvm/ffi/extra/module.h>
@@ -73,7 +74,11 @@ class CUDAModuleNode : public ffi::ModuleObj {
public:
CUDAModuleNode(ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String, FunctionInfo> fmap,
ffi::Map<ffi::String, ffi::String> source)
- : code_(code), fmt_(fmt), fmap_(fmap), source_(source) {
+ : code_(code),
+ fmt_(fmt),
+ fmap_(fmap),
+ source_(source),
+ backend_config_(source.Get("backend_config").value_or("{}")) {
std::fill(module_.begin(), module_.end(), nullptr);
}
// destructor
@@ -107,21 +112,18 @@ class CUDAModuleNode : public ffi::ModuleObj {
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) final;
ffi::Bytes SaveToBytes() const final {
- // Format: [fmt][fmap][code]. Source map is in-memory inspection only and
- // is NEVER serialized — it is lost on save/load round-trip (matches
- // upstream behavior; the receiver rebuilds source from code bytes if
- // possible). CUDAFallbackModuleNode::SaveToBytes (in
- // src/target/cuda/cuda_fallback_module.cc) MUST mirror this format
- // byte-for-byte; see one-way comment there.
+ // Keep the existing prefix; only compilation settings are appended.
std::string buffer;
support::BytesOutStream stream(&buffer);
stream.Write(fmt_);
stream.Write(fmap_);
stream.Write(code_);
+ stream.Write(backend_config_);
return ffi::Bytes(std::move(buffer));
}
ffi::String InspectSource(const ffi::String& format) const final {
+ if (format == "backend_config") return backend_config_;
// For known compiled formats, return code as string when format matches.
if (format == fmt_) {
return ffi::String(code_.data(), code_.size());
@@ -174,36 +176,6 @@ class CUDAModuleNode : public ffi::ModuleObj {
return state;
}
- /*!
- * \brief JIT-compile raw CUDA C++ source to PTX/cubin/fatbin via the Python
- * compile callback. Called from BOTH the
- * "ffi.Module.create.cuda" lambda (when the codegen hands us
- * fmt=="cuda") AND from LoadFromBytes (when the saved-on-disk fmt is
- * "cuda" — the cross-compile receiver path).
- *
- * \param source Raw CUDA C++ source (text).
- * \return Compiled binary bytes. Determination of the compiled format
- * (ptx vs cubin vs fatbin) is left to the caller (heuristic on
- * first byte: '/' → ptx-text, otherwise binary).
- */
- static ffi::Bytes JitCompileFromSource(const ffi::String& source) {
- // Registry: "tvm_callback_cuda_compile" — Python-side nvcc/nvrtc wrapper.
- // Grep hint: grep -rn 'tvm_callback_cuda_compile' src/ python/
- auto fcompile = ffi::Function::GetGlobal("tvm_callback_cuda_compile");
- TVM_FFI_CHECK(fcompile.has_value(), RuntimeError)
- << "fmt=='cuda' requires tvm_callback_cuda_compile to be registered. "
- << "Import tvm.support.nvcc.";
- return (*fcompile)(source).cast<ffi::Bytes>();
- }
-
- /*! \brief Pick the compiled format from the JIT output's first byte. */
- static ffi::String DetermineCompiledFormat(const ffi::Bytes& compiled) {
- if (compiled.size() > 0 && compiled.data()[0] == '/') {
- return ffi::String("ptx");
- }
- return ffi::String("cubin");
- }
-
private:
// The binary data (compiled PTX/cubin/fatbin, or raw CUDA source if fmt == "cuda").
ffi::Bytes code_;
@@ -211,8 +183,9 @@ class CUDAModuleNode : public ffi::ModuleObj {
ffi::String fmt_;
// function information table.
ffi::Map<ffi::String, FunctionInfo> fmap_;
- // In-memory source map for InspectSource — never serialized.
+ // In-memory source for inspection; not serialized.
ffi::Map<ffi::String, ffi::String> source_;
+ ffi::String backend_config_; // Serialized JSON; interpreted only by the compilation backend.
// the internal modules per GPU, to be lazily initialized.
std::array<CUmodule, kMaxNumGPUs> module_;
// All packed wrappers for one CUfunction share monotonically increasing
@@ -290,6 +263,15 @@ class CUDAWrappedFunc {
};
ffi::Optional<ffi::Function> CUDAModuleNode::GetFunction(const ffi::String& name) {
+ if (name == "__tvm_cuda_binary") {
+ auto self = ffi::GetRef<ffi::Module>(this);
+ return ffi::Function::FromTyped([self, this]() {
+ ffi::Array<ffi::String> names;
+ for (auto [name, info] : fmap_) names.push_back(name);
+ return ffi::Array<ffi::Any>{code_, fmt_, names};
+ });
+ }
+
ffi::ObjectPtr<ffi::Object> sptr_to_self = ffi::GetObjectPtr<ffi::Object>(this);
TVM_FFI_ICHECK_EQ(sptr_to_self.get(), this);
auto opt_info = fmap_.Get(name);
@@ -307,34 +289,42 @@ static ffi::Module CUDAModuleCreateImpl(ffi::Bytes code, ffi::String fmt,
ffi::Map<ffi::String, FunctionInfo> fmap,
ffi::Map<ffi::String, ffi::String> source) {
if (fmt == "cuda") {
- // Stash the CUDA source for InspectSource before we replace `code` with
- // the JIT output.
- if (source.find("cuda") == source.end()) {
- source.Set("cuda", ffi::String(code.data(), code.size()));
- }
- ffi::Bytes compiled =
- CUDAModuleNode::JitCompileFromSource(ffi::String(code.data(), code.size()));
- fmt = CUDAModuleNode::DetermineCompiledFormat(compiled);
- code = std::move(compiled);
+ ffi::String text(code.data(), code.size());
+ auto compile = ffi::Function::GetGlobalRequired("tvm_callback_cuda_compile");
+ auto result =
+ compile(text, source.Get("backend_config").value_or("{}")).cast<ffi::Array<ffi::Any>>();
+ source.Set("cuda", text);
+ code = result[0].cast<ffi::Bytes>();
+ fmt = result[1].cast<ffi::String>();
+ TVM_FFI_CHECK(fmt == "ptx" || fmt == "cubin" || fmt == "fatbin", ValueError)
+ << "Unsupported CUDA compilation output: " << fmt;
}
auto n = ffi::make_object<CUDAModuleNode>(code, fmt, fmap, source);
return ffi::Module(n);
}
static ffi::Module CUDAModuleLoadFromBytes(const ffi::Bytes& bytes) {
- support::BytesInStream stream(bytes);
ffi::String fmt;
ffi::Map<ffi::String, FunctionInfo> fmap;
ffi::Bytes code;
- stream.Read(&fmt);
+ ffi::Map<ffi::String, ffi::String> source;
+ support::BytesInStream stream(bytes);
+ TVM_FFI_ICHECK(stream.Read(&fmt));
TVM_FFI_ICHECK(stream.Read(&fmap));
- stream.Read(&code);
- // Source map is not serialized — it is lost on save/load round-trip.
- // If the receiver wants InspectSource("cuda") to work, the saved bytes must
- // have been written with fmt=="cuda" so the JIT path below re-stuffs the
- // source map with the original C++ source.
- return CUDAModuleCreateImpl(std::move(code), std::move(fmt), std::move(fmap),
- ffi::Map<ffi::String, ffi::String>());
+ TVM_FFI_ICHECK(stream.Read(&code));
+ // Upstream artifacts end after code. Accept that form, but reject a partial
+ // trailing string rather than silently treating corrupt metadata as absent.
+ uint64_t length = 0;
+ size_t count = stream.Read(&length, sizeof(length));
+ if (count != 0) {
+ TVM_FFI_ICHECK(count == sizeof(length) && length <= bytes.size());
+ std::string config(length, '\0');
+ TVM_FFI_ICHECK_EQ(stream.Read(config.data(), length), length);
+ source.Set("backend_config", ffi::String(config));
+ char trailing;
+ TVM_FFI_ICHECK_EQ(stream.Read(&trailing, 1), 0);
+ }
+ return CUDAModuleCreateImpl(std::move(code), std::move(fmt), std::move(fmap), std::move(source));
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/backend/cuda/transforms/bind_backend_config.cc b/src/backend/cuda/transforms/bind_backend_config.cc
new file mode 100644
index 0000000000..95c3d2dcd8
--- /dev/null
+++ b/src/backend/cuda/transforms/bind_backend_config.cc
@@ -0,0 +1,281 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file bind_backend_config.cc
+ * \brief Resolve CUDA configuration before any target-sensitive lowering.
+ */
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/unique_name_supply.h>
+#include <tvm/target/target.h>
+#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/transform.h>
+
+#include <map>
+#include <unordered_set>
+#include <vector>
+
+namespace tvm {
+namespace tirx {
+namespace transform {
+
+class BackendConfigBinder : public StmtExprMutator {
+ public:
+ BackendConfigBinder(Target target, ffi::String defaults)
+ : target_(target),
+ defaults_(defaults),
+ resolve_(ffi::Function::GetGlobalRequired("cuda.resolve_backend_config")) {}
+
+ Function Apply(Function func) {
+ TVM_FFI_CHECK(!func->GetAttr<ffi::String>("tirx.cuda_arch").has_value(), ValueError)
+ << "tirx.cuda_arch was removed; set backend_config['cuda']['arch'] on device_entry or "
+ "tvm.compile";
+ auto resolved =
+ resolve_(func->GetAttr<ffi::String>("backend_config").value_or(""), defaults_, target_)
+ .cast<ffi::Array<ffi::Any>>();
+ target_ = resolved[0].cast<Target>();
+ defaults_ = resolved[1].cast<ffi::String>();
+ if (!func->body.has_value()) return func;
+ auto body =
+ Mutate(func->body.value(), InplaceMode::kDisallow).ValueOrUnchanged(func->body.value());
+ func.CopyOnWrite()->body = body;
+ return WithAttrs(func, {{tvm::attr::kTarget, target_}, {"backend_config", defaults_}});
+ }
+
+ private:
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final {
+ if (op->op->name != "tirx.device_entry" && op->op->name != "tirx.device_scope") {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ auto local = op->attrs->dict.Get("backend_config");
+ auto resolved = resolve_(local.has_value() ? local.value().cast<ffi::String>() : ffi::String(),
+ defaults_, target_)
+ .cast<ffi::Array<ffi::Any>>();
+ auto attrs = op->attrs->dict;
+ attrs.Set(tvm::attr::kTarget, resolved[0]);
+ attrs.Set("backend_config", resolved[1]);
+ auto body = Mutate(op->body, InplaceMode::kDisallow).ValueOrUnchanged(op->body);
+ return RegionStmt(op->op, op->args, op->body_params, DictAttrs(attrs), body, op->result_vars,
+ op->loc);
+ }
+ Target target_;
+ ffi::String defaults_;
+ ffi::Function resolve_;
+};
+
+Pass BindBackendConfig(ffi::String defaults) {
+ auto transform = [defaults](Function func, IRModule mod, PassContext ctx) -> Function {
+ auto target = func->GetAttr<Target>(tvm::attr::kTarget);
+ if (!target.has_value() || target.value()->kind->name != "cuda") return func;
+ return BackendConfigBinder(target.value(), defaults).Apply(func);
+ };
+ return CreateFunctionPass(transform, 0, "cuda.BindBackendConfig");
+}
+
+// Private device helpers belong to a compilation group. Specialize their
+// transitive closure before dtype legalization so every helper sees the same
+// target and options as its callers. No device symbol crosses binary boundaries.
+class DeviceHelperSpecializer : public StmtExprMutator {
+ public:
+ DeviceHelperSpecializer(IRModule input, IRModule output, UniqueNameSupply names, Target target,
+ ffi::String config, bool active = true)
+ : input_(input),
+ output_(output),
+ names_(names),
+ target_(target),
+ config_(config),
+ active_(active) {}
+
+ Function Rewrite(Function func) {
+ if (!func->body.has_value()) return func;
+ auto body =
+ Mutate(func->body.value(), InplaceMode::kDisallow).ValueOrUnchanged(func->body.value());
+ func.CopyOnWrite()->body = body;
+ return func;
+ }
+
+ std::unordered_set<const GlobalVarNode*> originals;
+
+ private:
+ UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final {
+ if (op->op->name != "tirx.device_entry") {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ auto entry_target = op->attrs->dict.Get(tvm::attr::kTarget);
+ if (!entry_target || entry_target.value().cast<Target>()->kind->name != "cuda") {
+ return StmtExprMutator::Mutate_(op, inplace_mode);
+ }
+ auto saved_target = target_;
+ auto saved_config = config_;
+ auto saved_replacements = replacements_;
+ bool saved_active = active_;
+ target_ = entry_target.value().cast<Target>().WithoutHost();
+ config_ = op->attrs->dict.at("backend_config").cast<ffi::String>();
+ replacements_ = {};
+ active_ = true;
+ auto result = StmtExprMutator::Mutate_(op, inplace_mode);
+ target_ = saved_target;
+ config_ = saved_config;
+ replacements_ = saved_replacements;
+ active_ = saved_active;
+ return result;
+ }
+
+ UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) final {
+ auto result =
+ StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Expr>(op));
+ auto call = result.as_or_throw<Call>();
+ auto callee = call->op.as<GlobalVar>();
+ if (!active_ || !callee.has_value() || !input_->functions.count(callee.value())) return result;
+ originals.insert(callee.value().get());
+ auto original = input_->Lookup(callee.value()).as_or_throw<Function>();
+ auto original_target = original->GetAttr<Target>(tvm::attr::kTarget);
+ TVM_FFI_CHECK(original_target.has_value() && original_target.value()->kind->name == "cuda",
+ ValueError)
+ << "CUDA device function cannot call a non-CUDA helper";
+ TVM_FFI_CHECK(
+ original->GetAttr<CallingConv>(tvm::attr::kCallingConv, CallingConv::kDefault).value() ==
+ CallingConv::kDefault,
+ ValueError)
+ << "CUDA device calls must refer to private device helpers";
+ GlobalVar replacement(ffi::UnsafeInit{});
+ if (auto existing = replacements_.Get(callee.value())) {
+ replacement = existing.value();
+ } else {
+ bool same_group = original_target.value()->str() == target_->str() &&
+ original->GetAttr<ffi::String>("backend_config").value_or("") == config_;
+ replacement = same_group
+ ? callee.value()
+ : GlobalVar(names_->FreshName(callee.value()->name_hint + "_cuda", false));
+ replacements_.Set(callee.value(), replacement);
+ auto helper =
+ WithAttrs(original, {{tvm::attr::kTarget, target_}, {"backend_config", config_}});
+ helper = Rewrite(helper);
+ output_->Add(replacement, helper);
+ }
+ call.CopyOnWrite()->op = replacement;
+ return call;
+ }
+ IRModule input_, output_;
+ UniqueNameSupply names_;
+ Target target_;
+ ffi::String config_;
+ bool active_;
+ ffi::Map<GlobalVar, GlobalVar> replacements_;
+};
+
+// Entry targets must also govern helper tile dispatch, which precedes splitting.
+// Clone before lowering, then discard only originals made unreachable by cloning.
+class HelperCallCollector : public StmtExprVisitor {
+ public:
+ std::vector<GlobalVar> calls;
+ ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
+ if (auto gv = op->op.as<GlobalVar>()) calls.push_back(gv.value());
+ return StmtExprVisitor::Visit_(op);
+ }
+};
+
+Pass SpecializeEntryHelpers() {
+ auto transform = [](IRModule mod, PassContext ctx) {
+ IRModule result({}, mod->source_map, mod->attrs, mod->global_infos);
+ UniqueNameSupply names(mod->functions.begin(), mod->functions.end(),
+ [](const auto& kv) { return kv.first->name_hint; });
+ std::unordered_set<const GlobalVarNode*> originals;
+ for (auto [gv, base] : mod->functions) {
+ auto func = base.as<Function>();
+ auto target = func ? func.value()->GetAttr<Target>(tvm::attr::kTarget) : std::nullopt;
+ if (target && target.value()->kind->name == "cuda") {
+ auto rewriter = ffi::make_object<DeviceHelperSpecializer>(
+ mod, result, names, target.value(), ffi::String(), false);
+ result->Add(gv, rewriter->Rewrite(func.value()));
+ originals.insert(rewriter->originals.begin(), rewriter->originals.end());
+ } else {
+ result->Add(gv, base);
+ }
+ }
+ std::unordered_set<const GlobalVarNode*> reachable;
+ std::vector<GlobalVar> pending;
+ for (auto [gv, base] : result->functions) {
+ if (!originals.count(gv.get())) pending.push_back(gv);
+ }
+ while (!pending.empty()) {
+ auto gv = pending.back();
+ pending.pop_back();
+ if (!reachable.insert(gv.get()).second) continue;
+ auto func = result->Lookup(gv).as<Function>();
+ if (!func || !func.value()->body.has_value()) continue;
+ auto collector = ffi::make_object<HelperCallCollector>();
+ collector->Visit(func.value()->body.value());
+ for (auto callee : collector->calls) {
+ if (result->functions.count(callee)) pending.push_back(callee);
+ }
+ }
+ ffi::Array<GlobalVar> remove;
+ for (auto [gv, base] : result->functions) {
+ if (originals.count(gv.get()) && !reachable.count(gv.get())) remove.push_back(gv);
+ }
+ for (auto gv : remove) result->Remove(gv);
+ return result;
+ };
+ return CreateModulePass(transform, 0, "cuda.SpecializeEntryHelpers");
+}
+
+Pass SpecializeDeviceHelpers() {
+ auto transform = [](IRModule mod, PassContext ctx) {
+ IRModule result({}, mod->source_map, mod->attrs, mod->global_infos);
+ UniqueNameSupply names(mod->functions.begin(), mod->functions.end(),
+ [](const auto& kv) { return kv.first->name_hint; });
+ std::map<std::string, ffi::ObjectPtr<DeviceHelperSpecializer>> groups;
+ for (auto [gv, base] : mod->functions) {
+ auto func = base.as<Function>();
+ if (!func.has_value()) {
+ result->Add(gv, base);
+ continue;
+ }
+ auto target = func.value()->GetAttr<Target>(tvm::attr::kTarget);
+ if (!target.has_value() || target.value()->kind->name != "cuda") {
+ result->Add(gv, base);
+ continue;
+ }
+ auto cc = func.value()->GetAttr<CallingConv>(tvm::attr::kCallingConv, CallingConv::kDefault);
+ if (cc.value() != CallingConv::kDeviceKernelLaunch) {
+ if (target.value()->host.has_value()) result->Add(gv, base);
+ continue;
+ }
+ auto config = func.value()->GetAttr<ffi::String>("backend_config").value_or("");
+ std::string key = std::string(target.value()->str()) + "\n" + std::string(config);
+ auto& group = groups[key];
+ if (!group)
+ group =
+ ffi::make_object<DeviceHelperSpecializer>(mod, result, names, target.value(), config);
+ result->Add(gv, group->Rewrite(func.value()));
+ }
+ return result;
+ };
+ return CreateModulePass(transform, 0, "cuda.SpecializeDeviceHelpers");
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef()
+ .def("tirx.backend.cuda.transforms.BindBackendConfig", BindBackendConfig)
+ .def("tirx.backend.cuda.transforms.SpecializeDeviceHelpers", SpecializeDeviceHelpers)
+ .def("tirx.backend.cuda.transforms.SpecializeEntryHelpers", SpecializeEntryHelpers);
+}
+} // namespace transform
+} // namespace tirx
+} // namespace tvm
diff --git a/src/target/source/codegen_c_host.h b/src/target/source/codegen_c_host.h
index ba38338d28..b98ec39de9 100644
--- a/src/target/source/codegen_c_host.h
+++ b/src/target/source/codegen_c_host.h
@@ -58,8 +58,8 @@ class CodeGenCHost : public CodeGenC {
void DefineModuleName();
using CodeGenC::PrintType;
- void PrintType(const PrimType& t, std::ostream& os) final; // NOLINT(*)
- void PrintFuncPrefix(std::ostream& os) final; // NOLINT(*)
+ void PrintType(const PrimType& t, std::ostream& os) override; // NOLINT(*)
+ void PrintFuncPrefix(std::ostream& os) final; // NOLINT(*)
// overload visitor functions
void Dispatch_(const prim::BroadcastNode* op, std::ostream& os) final; // NOLINT(*)
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index 449b3c6d3a..c2ccb4f580 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -172,7 +172,7 @@ inline PrimType APIType(const PrimType& t) {
if (t.MatchesCode(DLDataTypeCode::kDLBool, DLDataTypeCode::kDLUInt, DLDataTypeCode::kDLInt)) {
return PrimType::Int(64);
}
- TVM_FFI_ICHECK_EQ(t.code(), DLDataTypeCode::kDLFloat);
+ TVM_FFI_ICHECK(t.MatchesCode(DLDataTypeCode::kDLFloat, DLDataTypeCode::kDLBfloat));
return PrimType::Float(64);
}
diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc
index aabeb2f4de..009cabf020 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -257,6 +257,11 @@ class HostDeviceSplitter : public StmtExprMutator {
device_func = WithAttrs(std::move(device_func), {{tvm::attr::kTarget, device_target},
{tvm::tirx::attr::kNoAlias, true},
{tvm::tirx::attr::kIsGlobalFunc, true}});
+ if (auto config = region->attrs->dict.Get("backend_config")) {
+ device_func = WithAttr(std::move(device_func), "backend_config", config.value());
+ } else if (auto config = cur_func_->GetAttr<ffi::String>("backend_config")) {
+ device_func = WithAttr(std::move(device_func), "backend_config", config.value());
+ }
bool is_stir = cur_func_->attrs->dict.count(tvm::attr::kSTir);
if (is_stir) {
device_func = WithAttr(std::move(device_func), tvm::attr::kSTir, true);
@@ -777,6 +782,14 @@ class DeviceKernelMutator : public StmtExprMutator {
// same-device-type coincidence.
bool force_kernel_launch = callee_is_kernel && caller_is_host;
+ // CUDA helpers are specialized per device compilation group after splitting.
+ // Keep their GlobalVar edge even when entry architectures differ, so the
+ // helper's transitive closure can be cloned and lowered for each caller.
+ if (!caller_is_host && !callee_is_kernel && caller_target->kind->name == "cuda" &&
+ callee_target->kind->name == "cuda") {
+ return node;
+ }
+
if (!force_kernel_launch) {
bool same_target = caller_target->str() == callee_target->str();
if (same_target) {
diff --git a/src/tirx/transform/tile_dispatch.cc b/src/tirx/transform/tile_dispatch.cc
index dd8435c7e3..3200ca9baa 100644
--- a/src/tirx/transform/tile_dispatch.cc
+++ b/src/tirx/transform/tile_dispatch.cc
@@ -326,6 +326,11 @@ class TileDispatcher : public StmtExprMutator {
Stmt ProcessDeviceEntry(const RegionStmtNode* entry_node) {
Stmt body_to_visit = entry_node->body;
+ Target saved_target = target_;
+ if (auto target = entry_node->attrs->dict.Get(tvm::attr::kTarget)) {
+ target_ = target.value().cast<Target>();
+ }
+ With<Target> target_scope(target_);
bool is_first_block = false;
std::swap(is_first_block, is_first_block_);
@@ -346,6 +351,7 @@ class TileDispatcher : public StmtExprMutator {
auto pop_exec_contexts = [&]() {
if (pushed_base_ctx) ctx_stack_.pop_back();
+ target_ = saved_target;
};
if (!is_first_block) {
@@ -375,8 +381,8 @@ class TileDispatcher : public StmtExprMutator {
alloc_buffers_.clear();
Stmt res = body;
- if (native_launch_) {
- res = CudaIndexLowerer::Lower(res, launch_params_);
+ if (native_launch_) res = CudaIndexLowerer::Lower(res, launch_params_);
+ if (native_launch_ || entry_node->attrs->dict.count("backend_config")) {
res = RegionStmt(Op::Get("tirx.device_scope"), entry_node->args, {}, entry_node->attrs, res);
}
@@ -1456,7 +1462,7 @@ class TileDispatcher : public StmtExprMutator {
std::unordered_map<Var, ScopeIdTarget, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> scope_aliases_;
ffi::Map<Var, Range> var_range_map_;
sym::Analyzer analyzer_;
- const Target& target_;
+ Target target_;
std::vector<ExecContext> ctx_stack_;
std::unordered_map<ffi::String, ffi::Tuple<PrimVar, PrimExpr>> launch_params_;
std::vector<Bind> alloc_buffers_;
diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc
index 825ef396fe..c352e59333 100644
--- a/src/tirx/transform/tvm_ffi_binder.cc
+++ b/src/tirx/transform/tvm_ffi_binder.cc
@@ -588,7 +588,7 @@ void TVMFFIABIBuilder::DecodeParam(int param_index) {
} else if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
arg_value = DecodeParamInt(param_index, type_index.as_or_throw<PrimExpr>(), dtype);
} else {
- TVM_FFI_ICHECK_EQ(dtype.code(), DLDataTypeCode::kDLFloat);
+ TVM_FFI_ICHECK(dtype.MatchesCode(DLDataTypeCode::kDLFloat, DLDataTypeCode::kDLBfloat));
arg_value = DecodeParamFloat(param_index, type_index.as_or_throw<PrimExpr>(), dtype);
}
diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 0e4b0ccf22..3337880afb 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -1,5 +1,5 @@
# Licensed to the Apache Software Foundation (ASF) under one
-# ruff: noqa: E501, E741, F401, F841
+# ruff: noqa: E501, E741, F841
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
@@ -23,40 +23,23 @@ import pytest
import tvm
import tvm.support.nvcc
import tvm.testing
+from tvm.backend.config import merge_backend_configs
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.support.nvcc import have_bf16, have_fp16, have_int8
from tvm.testing import env
-@pytest.fixture(autouse=True, params=["nvcc", "nvrtc"])
-def setup_cuda_compile_mode(request):
- mode = request.param
- if mode == "nvrtc":
- try:
- from cuda.bindings import nvrtc
- except ImportError:
- pytest.skip("cuda-python not available, skipping nvrtc tests")
-
- orig_func = tvm.support.nvcc.tvm_callback_cuda_compile
-
- def compile_mode_wrapper(code):
- if mode == "nvcc":
- return tvm.support.nvcc.compile_cuda(code, target_format="fatbin", compiler="nvcc")
- elif mode == "nvrtc":
- return tvm.support.nvcc.compile_cuda(code, target_format="cubin", compiler="nvrtc")
- else:
- raise ValueError(f"Unknown mode: {mode}")
-
- tvm.register_global_func("tvm_callback_cuda_compile", compile_mode_wrapper, override=True)
- # yield back to the original function so that each test runs twice
- yield
- tvm.register_global_func("tvm_callback_cuda_compile", orig_func, override=True)
+@pytest.fixture(params=["nvcc", "nvrtc"])
+def backend_config(request):
+ if request.param == "nvrtc":
+ pytest.importorskip("cuda.bindings.nvrtc")
+ return {"cuda": {"compiler": request.param}}
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_host_bundle(tmp_path):
+def test_cuda_host_bundle(tmp_path, *, backend_config):
from shutil import which
import tvm_ffi.cpp
@@ -72,12 +55,15 @@ def test_cuda_host_bundle(tmp_path):
B[tx] = A[tx] + T.float32(1)
target = tvm.target.Target("cuda", host="cuda_host")
- built = tvm.compile(add_one, target=target).mod
+ built = tvm.compile(add_one, target=target, backend_config=backend_config).mod
source = export_cuda_host(built)
- assert source.index("__global__") < source.index("cudaLaunchKernelEx")
+ assert "tvm_cuda_binary_0" in source
+ assert "__global__" not in source
+ assert "cuLaunchKernelEx" in source
library = tvm_ffi.cpp.build_inline(
name="cuda_host_add_one",
cuda_sources=source,
+ extra_ldflags=["-lcuda"],
extra_cuda_cflags=[f"-arch={target.arch}"],
build_directory=str(tmp_path),
backend="cuda",
@@ -97,7 +83,7 @@ def test_cuda_host_bundle(tmp_path):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_host_bundle_bf16_cluster(tmp_path):
+def test_cuda_host_bundle_bf16_cluster(tmp_path, *, backend_config):
"""bfloat16 pointers and cluster launch attributes in a CUDA-host bundle."""
from shutil import which
@@ -128,12 +114,18 @@ def test_cuda_host_bundle_bf16_cluster(tmp_path):
arch = env.cuda_arch()
target = tvm.target.Target({"kind": "cuda", "arch": arch}, host="cuda_host")
with target:
- built = tvm.compile(tvm.IRModule({"main": main}), target=target, tir_pipeline="tirx").mod
+ built = tvm.compile(
+ tvm.IRModule({"main": main}),
+ target=target,
+ tir_pipeline="tirx",
+ backend_config=backend_config,
+ ).mod
source = export_cuda_host(built)
assert "cudaLaunchAttributeClusterDimension" in source
library = tvm_ffi.cpp.build_inline(
name="cuda_host_bf16_cluster",
cuda_sources=source,
+ extra_ldflags=["-lcuda"],
extra_cuda_cflags=[f"-arch={arch}"],
build_directory=str(tmp_path),
backend="cuda",
@@ -156,7 +148,7 @@ def test_cuda_host_bundle_bf16_cluster(tmp_path):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_host_bundle_programmatic_dependent_launch(tmp_path):
+def test_cuda_host_bundle_programmatic_dependent_launch(tmp_path, *, backend_config):
"""A programmatic-dependent-launch flag becomes a launch attribute."""
from shutil import which
@@ -176,11 +168,14 @@ def test_cuda_host_bundle_programmatic_dependent_launch(tmp_path):
B[tx] = A[tx] + T.float32(1)
target = tvm.target.Target("cuda", host="cuda_host")
- source = export_cuda_host(tvm.compile(add_one, target=target).mod)
+ source = export_cuda_host(
+ tvm.compile(add_one, target=target, backend_config=backend_config).mod
+ )
assert "cudaLaunchAttributeProgrammaticStreamSerialization" in source
library = tvm_ffi.cpp.build_inline(
name="cuda_host_programmatic_dependent_launch",
cuda_sources=source,
+ extra_ldflags=["-lcuda"],
extra_cuda_cflags=[f"-arch={target.arch}"],
build_directory=str(tmp_path),
backend="cuda",
@@ -200,7 +195,7 @@ def test_cuda_host_bundle_programmatic_dependent_launch(tmp_path):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_host_bundle_tensor_map_parameter(tmp_path, monkeypatch):
+def test_cuda_host_bundle_tensor_map_parameter(tmp_path, monkeypatch, *, backend_config):
"""A tensor-map parameter handle compiles in the C++ host wrapper."""
from shutil import which
@@ -220,13 +215,19 @@ def test_cuda_host_bundle_tensor_map_parameter(tmp_path, monkeypatch):
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"}, host="cuda_host")
with target:
- built = tvm.compile(tvm.IRModule({"main": main}), target=target, tir_pipeline="tirx").mod
+ built = tvm.compile(
+ tvm.IRModule({"main": main}),
+ target=target,
+ tir_pipeline="tirx",
+ backend_config=backend_config,
+ ).mod
source = export_cuda_host(built)
assert "(CUtensorMap*)" in source
monkeypatch.setenv("TVM_FFI_CUDA_ARCH_LIST", "9.0a")
tvm_ffi.cpp.build_inline(
name="cuda_host_tensor_map_parameter",
cuda_sources=source,
+ extra_ldflags=["-lcuda"],
build_directory=str(tmp_path),
backend="cuda",
)
@@ -234,7 +235,7 @@ def test_cuda_host_bundle_tensor_map_parameter(tmp_path, monkeypatch):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_vectorize_add():
+def test_cuda_vectorize_add(*, backend_config):
num_thread = 8
def check_cuda(dtype, n, lanes):
@@ -258,7 +259,7 @@ def test_cuda_vectorize_add():
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")
+ fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -286,7 +287,7 @@ def test_cuda_vectorize_add():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_bf16_vectorize_add():
+def test_cuda_bf16_vectorize_add(*, backend_config):
if not have_bf16(tvm.cuda(0).compute_version):
print("skip because gpu does not support bf16")
return
@@ -322,7 +323,7 @@ def test_cuda_bf16_vectorize_add():
with tvm.transform.PassContext(
disabled_pass=["tirx.BF16Promote", "tirx.BF16CastElimination", "tirx.BF16TypeLowering"]
):
- fun = tvm.compile(Module, target="cuda")
+ fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -344,7 +345,7 @@ def test_cuda_bf16_vectorize_add():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_multiply_add():
+def test_cuda_multiply_add(*, backend_config):
num_thread = 8
def check_cuda(dtype, n, lanes):
@@ -374,7 +375,7 @@ def test_cuda_multiply_add():
ty="int32",
)
- fun = tvm.compile(Module, target="cuda")
+ fun = tvm.compile(Module, target="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))
@@ -397,7 +398,7 @@ def test_cuda_multiply_add():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_vectorize_load():
+def test_cuda_vectorize_load(*, backend_config):
num_thread = 8
def check_cuda(dtype, n, lanes):
@@ -413,7 +414,7 @@ def test_cuda_vectorize_load():
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")
+ fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -434,7 +435,7 @@ def test_cuda_vectorize_load():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_make_int8():
+def test_cuda_make_int8(*, backend_config):
def check_cuda(n, value, lanes):
dtype = "int8"
const_value = tvm.tirx.const(value, dtype=dtype)
@@ -448,7 +449,7 @@ def test_cuda_make_int8():
for j in T.vectorized(lanes):
A[i, j] = const_value
- fun = tvm.compile(Module, target="cuda")
+ fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -472,7 +473,7 @@ def test_cuda_make_int8():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_inf_nan():
+def test_cuda_inf_nan(*, backend_config):
def check_inf_nan(n, value, dtype):
inf_value = tvm.tirx.const(value, dtype=dtype)
@@ -486,7 +487,7 @@ def test_cuda_inf_nan():
if i_0 * 8 + i_1 < n:
C[i_0 * 8 + i_1] = inf_value
- fun = tvm.compile(Module, target="cuda")
+ fun = tvm.compile(Module, target="cuda", backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -511,7 +512,7 @@ def test_cuda_inf_nan():
pytest.param("rocm", marks=pytest.mark.gpu),
],
)
-def test_crossthread_reduction1(target):
+def test_crossthread_reduction1(target, *, backend_config):
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
@@ -544,7 +545,9 @@ def test_crossthread_reduction1(target):
if m_0 == 0:
B[i] = reduced[0]
- fun = tvm.compile(Module, target=target)
+ fun = tvm.compile(
+ Module, target=target, backend_config=backend_config if target == "cuda" else None
+ )
return fun
def verify(nthd):
@@ -576,7 +579,7 @@ def test_crossthread_reduction1(target):
pytest.param("rocm", marks=pytest.mark.gpu),
],
)
-def test_crossthread_reduction2(target):
+def test_crossthread_reduction2(target, *, backend_config):
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
@@ -623,7 +626,9 @@ def test_crossthread_reduction2(target):
if k0_0 == 0 and k1_0 == 0:
B[i] = reduced[0]
- func = tvm.compile(Module, target=target)
+ func = tvm.compile(
+ Module, target=target, backend_config=backend_config if target == "cuda" else None
+ )
return func
def verify(nthdx, nthdy):
@@ -652,7 +657,7 @@ def test_crossthread_reduction2(target):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_reduction_binding():
+def test_cuda_reduction_binding(*, backend_config):
@I.ir_module
class Module:
@T.function
@@ -665,12 +670,12 @@ def test_cuda_reduction_binding():
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")
+ func = tvm.compile(Module, target="cuda", backend_config=backend_config)
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_const_float_to_half():
+def test_cuda_const_float_to_half(*, backend_config):
# This import is required to use nvcc to perform code gen;
# otherwise it is found that the code gen is done by nvrtc.
@@ -697,7 +702,7 @@ def test_cuda_const_float_to_half():
]
)
- func = tvm.compile(Module, target="cuda")
+ func = tvm.compile(Module, target="cuda", backend_config=backend_config)
shape = (2, 3, 4)
a_np = np.random.uniform(size=shape).astype("float16")
@@ -715,7 +720,7 @@ def test_cuda_const_float_to_half():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_floordiv_with_vectorization():
+def test_cuda_floordiv_with_vectorization(*, backend_config):
with tvm.target.Target("cuda"):
# B[i] = A[floordiv(i, k)]
n = 256
@@ -733,7 +738,7 @@ def test_cuda_floordiv_with_vectorization():
(i_0 * 256 + i_1_0 * 4 + i_1_1) // 37
]
- func = tvm.compile(Module, target="cuda")
+ func = tvm.compile(Module, target="cuda", backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -749,7 +754,7 @@ def test_cuda_floordiv_with_vectorization():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_floormod_with_vectorization():
+def test_cuda_floormod_with_vectorization(*, backend_config):
with tvm.target.Target("cuda"):
# B[i] = A[floormod(i, k)]
n = 256
@@ -767,7 +772,7 @@ def test_cuda_floormod_with_vectorization():
(i_0 * 256 + i_1_0 * 4 + i_1_1) % 37
]
- func = tvm.compile(Module, target="cuda")
+ func = tvm.compile(Module, target="cuda", backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -831,7 +836,7 @@ _VECTORIZED_CAST_CASES += [("int8", "uint8", 16), ("uint8", "int8", 16)]
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
@pytest.mark.parametrize("t0,t1,factor", _VECTORIZED_CAST_CASES)
-def test_vectorized_casts(t0, t1, factor):
+def test_vectorized_casts(t0, t1, factor, *, backend_config):
if (t0 == "float16" or t1 == "float16") and not have_fp16(tvm.cuda(0).compute_version):
print("Skip because gpu does not have fp16 support")
return
@@ -850,7 +855,7 @@ def test_vectorized_casts(t0, t1, factor):
t0, B[i_0 * factor + i_1]
)
- func = tvm.compile(Module, target="cuda")
+ func = tvm.compile(Module, target="cuda", backend_config=backend_config)
# correctness
def run_and_check():
@@ -868,7 +873,7 @@ def test_vectorized_casts(t0, t1, factor):
tvm.testing.run_with_gpu_lock(run_and_check)
-def sched(compute_fn, dtype, n=128):
+def sched(compute_fn, dtype, n=128, *, backend_config=None):
"""Create a vectorized CUDA module with the given compute function.
The schedule structure is: split [1, None] -> split [32, None] -> split [None, 4]
@@ -889,12 +894,12 @@ def sched(compute_fn, dtype, n=128):
A[i0_1_0 * 4 + i0_1_1_0 * 4 + i0_1_1_1]
)
- return tvm.compile(Module, target="cuda")
+ return tvm.compile(Module, target="cuda", backend_config=backend_config)
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_vectorized_intrin1():
+def test_vectorized_intrin1(*, backend_config):
test_funcs = [
(tvm.tirx.floor, lambda x: np.floor(x)),
(tvm.tirx.ceil, lambda x: np.ceil(x)),
@@ -936,7 +941,7 @@ def test_vectorized_intrin1():
return
n = 128
- f = sched(tvm_intrin, dtype, n)
+ f = sched(tvm_intrin, dtype, n, backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -954,7 +959,7 @@ def test_vectorized_intrin1():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_round_ties_to_even():
+def test_round_ties_to_even(*, backend_config):
# The CUDA rule lowers tirx.round to nearbyint/nearbyintf (ties-to-even,
# matching constant-folding semantics), so exact midpoints must round to
# even. np.random-based tests never produce midpoints; this pins the tie
@@ -965,7 +970,7 @@ def test_round_ties_to_even():
test_values = np.tile(midpoints, 16) # 128 lanes, matches sched()
expected = np.tile(ties_to_even, 16)
- f = sched(tvm.tirx.round, "float32")
+ f = sched(tvm.tirx.round, "float32", backend_config=backend_config)
dev = tvm.cuda(0)
a = tvm.runtime.tensor(test_values, dev)
b = tvm.runtime.tensor(np.zeros(len(test_values), dtype="float32"), dev)
@@ -982,7 +987,7 @@ def test_round_ties_to_even():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_vectorized_intrin2(dtype="float32"):
+def test_vectorized_intrin2(dtype="float32", *, backend_config):
c2 = tvm.tirx.const(2, dtype=dtype)
test_funcs = [
(tvm.tirx.power, lambda x: np.power(x, 2.0)),
@@ -991,7 +996,7 @@ def test_vectorized_intrin2(dtype="float32"):
def run_test(tvm_intrin, np_func):
n = 128
- f = sched(lambda x: tvm_intrin(x, c2), dtype, n)
+ f = sched(lambda x: tvm_intrin(x, c2), dtype, n, backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -1008,7 +1013,7 @@ def test_vectorized_intrin2(dtype="float32"):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_vectorized_popcount():
+def test_vectorized_popcount(*, backend_config):
def ref_popcount(x):
cnt = 0
while x:
@@ -1018,7 +1023,7 @@ def test_vectorized_popcount():
def run_test(dtype):
n = 128
- f = sched(lambda x: tvm.tirx.popcount(x), dtype, n)
+ f = sched(lambda x: tvm.tirx.popcount(x), dtype, n, backend_config=backend_config)
def run_and_check():
dev = tvm.cuda(0)
@@ -1036,7 +1041,7 @@ def test_vectorized_popcount():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_vectorize_load_permute_pad():
+def test_cuda_vectorize_load_permute_pad(*, backend_config):
def check_cuda(dtype, n, l, padding, lanes):
if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
print("Skip because gpu does not have fp16 support")
@@ -1058,7 +1063,7 @@ def test_cuda_vectorize_load_permute_pad():
j < padding or l + padding <= j, zero, A[i * lanes + k, j - padding]
)
- fun = tvm.compile(Module, target="cuda")
+ fun = tvm.compile(Module, target="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)
@@ -1086,7 +1091,7 @@ def test_cuda_vectorize_load_permute_pad():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_try_unaligned_vector_load():
+def test_try_unaligned_vector_load(*, backend_config):
def build(N, C_N, offset):
@I.ir_module
class Module:
@@ -1097,7 +1102,7 @@ def test_try_unaligned_vector_load():
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")
+ f = tvm.tirx.build(Module, target="cuda", backend_config=backend_config)
kernel_source = f.imports[0].inspect_source()
a_data = np.arange(0, N).astype("float16")
@@ -1131,7 +1136,7 @@ def test_try_unaligned_vector_load():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_thread_sync_inside_condition():
+def test_cuda_thread_sync_inside_condition(*, backend_config):
@T.function
def func2(A: T.Tensor((4, 4), "float32")) -> None:
A_shared = T.alloc_tensor((4, 4), "float32", scope="shared")
@@ -1158,25 +1163,25 @@ def test_cuda_thread_sync_inside_condition():
for func in [func2, func3]:
mod = tvm.IRModule({"main": func})
- compiled = tvm.compile(mod, target="cuda")
+ compiled = tvm.compile(mod, target="cuda", backend_config=backend_config)
assert "__syncthreads()" in compiled.mod.imports[0].inspect_source()
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_invalid_reinterpret():
+def test_invalid_reinterpret(*, backend_config):
@T.function
def func(A: T.Tensor((4,), "uint32"), B: T.Tensor((4,), "uint8")) -> None:
for tx in T.thread_binding(4, "threadIdx.x"):
B[tx] = T.call_intrin("tirx.reinterpret", A[tx], ty="uint8")
with pytest.raises(RuntimeError):
- tvm.compile(func, target="cuda")
+ tvm.compile(func, target="cuda", backend_config=backend_config)
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
-def test_cuda_tensormap():
+def test_cuda_tensormap(*, backend_config):
# fmt: off
@T.function
def main(A: T.Tensor((16, 16), dtype='float32', align=16)):
@@ -1192,7 +1197,7 @@ def test_cuda_tensormap():
# fmt: on
mod = tvm.IRModule({"main": main})
- mod = tvm.compile(mod, target="cuda")
+ mod = tvm.compile(mod, target="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) {
@@ -1206,7 +1211,7 @@ extern "C" __global__ void __launch_bounds__(128) main_kernel(float* __restrict_
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_device_func_call():
+def test_cuda_device_func_call(*, backend_config):
@I.ir_module
class Module:
@T.function(private=True)
@@ -1223,14 +1228,14 @@ def test_cuda_device_func_call():
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")
+ lib = tvm.compile(Module, target="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
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_float_const_hex_format():
+def test_cuda_float_const_hex_format(*, backend_config):
"""Test that float constants are emitted in hexadecimal format for precision"""
@I.ir_module
@@ -1243,14 +1248,14 @@ def test_cuda_float_const_hex_format():
for tx in T.thread_binding(1024, "threadIdx.x"):
A[bx, tx] = T.float32(1 / 27)
- lib = tvm.compile(Module, target="cuda")
+ lib = tvm.compile(Module, target="cuda", backend_config=backend_config)
cuda_code = lib.mod.imports[0].inspect_source()
assert "0x1.2f684bda12f68p-5f" in cuda_code
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_device_host_call_same_func():
+def test_device_host_call_same_func(*, backend_config):
@I.ir_module
class Module:
@T.function(private=True)
@@ -1274,7 +1279,7 @@ def test_device_host_call_same_func():
# 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)
+ lib = tvm.compile(Module, target=target, 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
@@ -1293,7 +1298,10 @@ def test_device_host_call_same_func():
tvm.testing.run_with_gpu_lock(run_and_check)
-def test_thread_return():
+def test_thread_return(monkeypatch, *, backend_config):
+ if not env.has_cuda():
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+
@I.ir_module
class Module:
@T.function
@@ -1304,7 +1312,9 @@ def test_thread_return():
T.gpu_thread_return()
B[bx, tx] = A[bx, tx]
- lib = tvm.compile(Module, target="cuda")
+ lib = tvm.compile(
+ Module, backend_config=merge_backend_configs(backend_config, {"cuda": {"arch": "sm_80"}})
+ )
cuda_code = lib.mod.imports[0].inspect_source()
assert "return;" in cuda_code
assert "return 0;" not in cuda_code
@@ -1312,7 +1322,7 @@ def test_thread_return():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_cuda_loop_step():
+def test_cuda_loop_step(*, backend_config):
@T.function
def cuda_loop_step(
A: T.Tensor((1024,), "float32"),
@@ -1326,7 +1336,7 @@ def test_cuda_loop_step():
C[i] = A[i] + B[i]
target = tvm.target.Target({"kind": "cuda"})
- lib = tvm.compile(cuda_loop_step, target=target)
+ lib = tvm.compile(cuda_loop_step, target=target, backend_config=backend_config)
cuda_src = lib.mod.imports[0].inspect_source()
assert "i += 96" in cuda_src
@@ -1347,7 +1357,7 @@ def test_cuda_loop_step():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_export_load_with_fallback(monkeypatch, tmp_path):
+def test_export_load_with_fallback(monkeypatch, tmp_path, *, backend_config):
"""Force the codegen wrapper into the fallback branch, then export+load+run."""
n = 1024
@@ -1361,7 +1371,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path):
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")
+ host_lib = tvm.compile(Module, target="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 db2ecceedd..771701d4ac 100644
--- a/tests/python/disco/test_nvshmem.py
+++ b/tests/python/disco/test_nvshmem.py
@@ -18,7 +18,6 @@
# pylint: disable=missing-docstring
import multiprocessing
-import os
import shutil
import socket
import subprocess
@@ -288,11 +287,9 @@ extern "C" __global__ void nvshmem_query_kernel(int* my_pe_out, int* n_pes_out)
def _kernel_compile(compile_mode):
- """Compile and run a kernel that calls NVSHMEM functions.
+ """Compile and run a kernel with an explicit NVSHMEM compiler configuration."""
- Runs in a fresh process, so setting the env var is safe.
- """
- os.environ["TVM_CUDA_COMPILE_MODE"] = compile_mode
+ backend_config = {"cuda": {"compiler": compile_mode}}
num_workers = 2
sess = di.ProcessSession(num_workers=num_workers)
@@ -321,6 +318,7 @@ def _kernel_compile(compile_mode):
my_pe_out.data,
n_pes_out.data,
kernel_name="nvshmem_query_kernel",
+ backend_config=backend_config,
)
@R.function
@@ -340,7 +338,9 @@ def _kernel_compile(compile_mode):
path = tmpdir + "/test_nvshmem_kernel.so"
target = tvm.target.Target("cuda")
- tvm.compile(NvshmemQueryModule, target=target).export_library(path)
+ tvm.compile(
+ NvshmemQueryModule, target=target, backend_config=backend_config
+ ).export_library(path)
mod = sess.load_vm_module(path)
result = mod["main"]()
diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py
index e4fbfacca8..ad012a2995 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -3316,7 +3316,8 @@ def test_einsum():
# in either direction check that negative diagonal lengths are clamped to zero.
for offset in [1, 4, 6, -3, -6]:
verify_model_numerically(DirectDiagonal(offset), (torch.randn(3, 4),))
- verify_model_numerically(DirectTrace(), (torch.randn(4, 4),))
+ # Float32 trace reductions may accumulate diagonal entries in a different order.
+ verify_model_numerically(DirectTrace(), (torch.randn(4, 4),), rtol=1e-6, atol=1e-6)
# A square diagonal must read the input directly without O(N^2) copies,
# both before and after legalization.
diff --git a/tests/python/s_tir/base/test_s_tir_texture_scope.py b/tests/python/s_tir/base/test_s_tir_texture_scope.py
index f655c44018..dcb07a92e5 100644
--- a/tests/python/s_tir/base/test_s_tir_texture_scope.py
+++ b/tests/python/s_tir/base/test_s_tir_texture_scope.py
@@ -69,7 +69,7 @@ def test_texture_scope():
lowered = pipeline(lowered)
_, device_mods = split_host_device_mods(lowered)
assert len(device_mods) == 1
- device_target, device_mod = next(iter(device_mods.items()))
+ (device_target, _backend_config), device_mod = next(iter(device_mods.items()))
device_mod = finalize_device_passes()(device_mod)
source = tvm.get_global_func("target.build.opencl")(device_mod, device_target).inspect_source()
assert "__read_only image2d_array_t" in source
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
index 7edd2deb48..f9937fbce6 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
@@ -1845,7 +1845,7 @@ def test_thread_broadcast_rewrite_2():
def test_thread_broadcast_rewrite_2_full_pipeline():
target_host = "llvm" if tvm.runtime.enabled("llvm") else "c"
- target = tvm.target.Target("cuda").with_host(target_host)
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"}).with_host(target_host)
mod = tvm.IRModule.from_expr(thread_broadcast_2.with_attr("global_symbol", "main"))
mod = tvm.tirx.transform.BindTarget(target)(mod)
pipeline, _, _ = tvm.tirx.get_tir_pipeline("s_tir")
diff --git a/tests/python/tirx/codegen/conftest.py b/tests/python/tirx/codegen/conftest.py
index 283dca5706..39bba77146 100644
--- a/tests/python/tirx/codegen/conftest.py
+++ b/tests/python/tirx/codegen/conftest.py
@@ -24,6 +24,13 @@ import pytest
from tvm.testing import env
+@pytest.fixture(autouse=True)
+def _source_codegen_without_cuda(monkeypatch):
+ # CPU-only source checks must not depend on a CUDA compiler being installed.
+ if not env.has_cuda():
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+
+
@pytest.fixture(autouse=True)
def _release_cuda_cache_between_tests():
yield
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py
index c56ede0ab1..ac2cba9aca 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -28,13 +28,15 @@ from tvm.script import tirx as T
from tvm.testing import env
-def _get_source(func: tvm.tirx.Function, target=None) -> tuple[str, tvm.IRModule]:
+def _get_source(
+ func: tvm.tirx.Function, target=None, backend_config=None
+) -> tuple[str, tvm.IRModule]:
if target is None:
target = {"kind": "cuda", "arch": "sm_100a"}
target = tvm.target.Target(target)
mod = tvm.IRModule({"main": func})
with target:
- mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+ mod = tvm.compile(mod, target=target, tir_pipeline="tirx", backend_config=backend_config)
src = mod.mod.imports[0].inspect_source()
return src, mod
@@ -502,9 +504,7 @@ def test_ptx_sub_f16x2_codegen():
not (env.has_cuda_compute(10, 0) and env.has_nvcc_version(13, 4)),
reason="packed bf16 conversion requires sm_100; the dialect certifies on CUDA 13.4",
)
-def test_sparse_decode_conversion_intrinsics_codegen(monkeypatch):
- monkeypatch.setenv("TVM_CUDA_COMPILE_MODE", "nvcc")
-
+def test_sparse_decode_conversion_intrinsics_codegen():
@T.function
def main(
U16: T.Tensor((1,), "uint16"),
@@ -521,7 +521,7 @@ def test_sparse_decode_conversion_intrinsics_codegen(monkeypatch):
T.ptx.cvt.rn.bf16x2.e4m3x2(U32[1], U16[0])
T.ptx.add.f32x2(U64[0], pair, pair)
- src, _ = _get_source(main)
+ src, _ = _get_source(main, backend_config={"cuda": {"compiler": "nvcc"}})
assert "cvt.rz.ue8m0x2.f32 %0, %1, %2;" in src
assert "cvt.rn.bf16x2.ue8m0x2 %0, %1;" in src
assert "cvt.rn.bf16x2.e4m3x2 %0, %1;" in src
diff --git a/tests/python/tirx/codegen/test_ptx_addr.py b/tests/python/tirx/codegen/test_ptx_addr.py
index 03cf0b3dd8..a268bcdf58 100644
--- a/tests/python/tirx/codegen/test_ptx_addr.py
+++ b/tests/python/tirx/codegen/test_ptx_addr.py
@@ -22,9 +22,10 @@ import tvm_ffi
import tvm
from tvm.ir import Call, Op, const
from tvm.script import tirx as T
+from tvm.testing import env
from tvm.tirx.expr import Broadcast, CallEffectKind
-TARGET = tvm.target.Target("cuda")
+TARGET = tvm.target.Target({"kind": "cuda", "arch": env.cuda_arch() or "sm_90"})
def _cuda_source(func) -> str:
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tests/python/tirx/codegen/test_ptx_dialect.py
index 0fc79f6842..fc80ac6be7 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -29,7 +29,8 @@ from tvm.ir import Op
from tvm.script import tirx as T
from tvm.testing import env
-TARGET = tvm.target.Target("cuda")
+TARGET = tvm.target.Target({"kind": "cuda", "arch": env.cuda_arch() or "sm_90"})
+
requires_nvcc = pytest.mark.skipif(shutil.which("nvcc") is None, reason="nvcc not available")
@@ -51,8 +52,15 @@ def _assert_ptxas_ok(src: str, rdc: bool = False, arch: str = PTX_ARCH) -> None:
"""Assemble through ptxas (cubin) — `-ptx` alone never validates inline asm."""
from tvm.support import nvcc
- options = ["-rdc=true"] if rdc else None
- nvcc.compile_cuda(src, target_format="cubin", arch=arch, options=options, compiler="nvcc")
+ options = ["--use_fast_math"] + (["-rdc=true"] if rdc else [])
+ nvcc.compile_cuda(
+ src,
+ backend_config={
+ "cuda": T.cuda.BackendConfig(
+ arch=arch, target_format="cubin", compiler="nvcc", nvcc=options
+ )
+ },
+ )
def test_ptx_registration():
diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
index c115737fbd..9ec4b24032 100644
--- a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
+++ b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
@@ -24,7 +24,8 @@ from tvm.script import tirx as T
from tvm.testing import env
from tvm.tirx.cuda.tile_primitive.copy._common import copy_ptx_form
-TARGET = tvm.target.Target("cuda")
+TARGET = tvm.target.Target({"kind": "cuda", "arch": env.cuda_arch() or "sm_90"})
+
# num_bytes → kernel layout. ``fill_offset`` fills lane i with ``i + fill_offset``.
_SHARED_COPY_CASES = {
@@ -121,7 +122,7 @@ def test_ptx_ld_st_codegen_emits_shared_asm():
T.cuda.tile.st(D[0:4], reg[:])
# fmt: on
- target = tvm.target.Target("cuda")
+ target = TARGET
with target:
mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx")
src = mod.mod.imports[0].inspect_source("cuda")
@@ -193,7 +194,7 @@ def test_ptx_ld_global_nc_v8_codegen():
for i in T.unroll(8):
out[i] = tmp[i]
- target = tvm.target.Target("cuda")
+ target = TARGET
with target:
mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx")
src = mod.mod.imports[0].inspect_source("cuda")
@@ -220,7 +221,7 @@ def test_ptx_ld_global_nc_v4_u64_256b_codegen():
for i in T.unroll(4):
out[i] = tmp[i]
- target = tvm.target.Target("cuda")
+ target = TARGET
with target:
mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx")
src = mod.mod.imports[0].inspect_source("cuda")
@@ -246,7 +247,7 @@ def test_ptx_ld_vector_scatter_dst_codegen():
out[2] = tmp2[0]
out[3] = tmp3[0]
- target = tvm.target.Target("cuda")
+ target = TARGET
with target:
mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx")
src = mod.mod.imports[0].inspect_source("cuda")
diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py
index ab667df6ea..0cde9aac47 100644
--- a/tests/python/tirx/iket/test_iket_profiler.py
+++ b/tests/python/tirx/iket/test_iket_profiler.py
@@ -41,6 +41,13 @@ TARGET = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
ORACLE_PATH = Path(__file__).parent / "oracle" / "iket_official_cutlass_4_6_0_oracle.json"
+@pytest.fixture(autouse=True)
+def _source_codegen_without_cuda(monkeypatch):
+ # Metadata and source checks do not require an installed CUDA compiler.
+ if not env.has_cuda():
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+
+
@T.function
def serial_a(out: T.Tensor((32,), "int32")):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
@@ -350,7 +357,10 @@ def _nvrtc_disassemble(source, tmp_path):
from tvm.support.nvcc import compile_cuda
tmp_path.mkdir(parents=True, exist_ok=True)
- cubin = compile_cuda(source, target_format="cubin", arch="sm_100a", compiler="nvrtc")
+ cubin = compile_cuda(
+ source,
+ backend_config={"cuda": {"arch": "sm_100a", "compiler": "nvrtc", "target_format": "cubin"}},
+ )
cubin_path = tmp_path / "kernel.cubin"
cubin_path.write_bytes(cubin)
return subprocess.run(
diff --git a/tests/python/tirx/test_cuda_backend_config.py b/tests/python/tirx/test_cuda_backend_config.py
new file mode 100644
index 0000000000..5dd2ed765e
--- /dev/null
+++ b/tests/python/tirx/test_cuda_backend_config.py
@@ -0,0 +1,732 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Explicit CUDA compile settings, independent device groups and artifact replay."""
+
+import os
+import shutil
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.backend.config import (
+ backend_config_json,
+ copy_backend_config,
+ merge_backend_configs,
+ parse_backend_config,
+)
+from tvm.backend.cuda import BackendConfig
+from tvm.backend.cuda.backend_config import compiler_options, resolve_backend_config
+from tvm.backend.cuda.transforms import BindBackendConfig
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def _require_compiler(compiler):
+ if compiler == "nvcc":
+ if shutil.which("nvcc") is None:
+ pytest.skip("NVCC is not installed")
+ from tvm.support.nvcc import get_cuda_version
+
+ version = get_cuda_version()
+ else:
+ nvrtc = pytest.importorskip("cuda.bindings.nvrtc")
+ status, major, minor = nvrtc.nvrtcVersion()
+ assert status == nvrtc.nvrtcResult.NVRTC_SUCCESS
+ version = (major, minor)
+ if version < (12, 8):
+ pytest.skip("sm_100a compile tests require CUDA 12.8 or newer")
+
+
+def test_defaults_overrides_snapshot_and_roundtrip():
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+ defaults = resolve_backend_config(target=target)
+ assert defaults["compiler"] == "nvrtc"
+ assert defaults["nvcc"] == defaults["nvrtc"] == ["--use_fast_math"]
+ assert defaults["ptxas"][-1] == "--register-usage-level=10"
+ assert resolve_backend_config({}, target) == defaults
+ local = {"cuda": BackendConfig(nvrtc=[], ptxas=["-O0"])}
+ result = merge_backend_configs({"cuda": defaults}, local)
+ assert result["cuda"]["nvrtc"] == []
+ assert result["cuda"]["ptxas"] == ["-O0"]
+ assert result["cuda"]["arch"] == "sm_100a"
+ local["cuda"]["nvrtc"].append("--ftz=true")
+ defaults["ptxas"].clear()
+ assert result["cuda"]["nvrtc"] == []
+ assert result["cuda"]["ptxas"] == ["-O0"]
+ assert parse_backend_config(backend_config_json(result)) == result
+ assert copy_backend_config(None) == copy_backend_config({}) == {}
+ assert backend_config_json({"cuda": {"arch": "sm_100a", "nvrtc": []}}) == (
+ backend_config_json({"cuda": {"nvrtc": [], "arch": "sm_100a"}})
+ )
+ # JSON sorts mapping keys, but never changes the order of compiler arguments.
+ options = {"cuda": {"nvrtc": ["--ftz=true", "--ftz=false"]}}
+ assert parse_backend_config(backend_config_json(options)) == options
+
+
+@pytest.mark.parametrize(
+ "config,exception",
+ [
+ ({"cuda": {"unknown": []}}, ValueError),
+ ({"cuda": {"compiler": "clang"}}, ValueError),
+ ({"cuda": {"arch": "compute_100"}}, ValueError),
+ ({"cuda": {"arch": T.int32(100)}}, TypeError),
+ ({"cuda": {"nvrtc": "--ftz=false"}}, TypeError),
+ ({"cuda": {"nvcc": [1]}}, TypeError),
+ ({"cuda": {"compiler": "nvrtc", "target_format": "fatbin"}}, ValueError),
+ ({"rocm": {}}, ValueError),
+ ({"compiler": "nvcc"}, ValueError),
+ ],
+)
+def test_fixed_keys_and_types(config, exception):
+ with pytest.raises(exception):
+ copy_backend_config(config)
+
+
+@pytest.mark.parametrize(
+ "name,option",
+ [
+ ("nvrtc", "--gpu-architecture=sm_90"),
+ ("nvcc", "-arch"),
+ ("nvcc", "--gpu-code=sm_90"),
+ ("nvcc", "--options-file=flags.txt"),
+ ("nvcc", "--cubin"),
+ ("nvrtc", "--ptxas-options=-O1"),
+ ("ptxas", "--output-file=kernel.cubin"),
+ ],
+)
+def test_managed_routing_flags(name, option):
+ with pytest.raises(ValueError, match="managed"):
+ copy_backend_config({"cuda": {name: [option]}})
+
+
+@pytest.mark.parametrize("compiler,format", [("nvrtc", "cubin"), ("nvcc", "fatbin")])
+def test_native_flags_and_default_format(monkeypatch, compiler, format):
+ from tvm.backend.cuda.compiler import compile_source
+ from tvm.support import nvcc
+
+ calls = []
+
+ def driver(source, output_format, arch, options, **kwargs):
+ calls.append((output_format, arch, options))
+ return b"artifact", "compiler output"
+
+ monkeypatch.setattr(nvcc, f"_compile_cuda_{compiler}", driver)
+ config = {
+ "cuda": {
+ "arch": "sm_100a",
+ "compiler": compiler,
+ compiler: ["--ftz=false", "-I/include", "--future-toolchain-option=1"],
+ "ptxas": ["-O1", "--register-usage-level=2"],
+ }
+ }
+ result = compile_source("// source", config)
+ assert result.target_format == format
+ assert calls[0][:2] == (format, "sm_100a")
+ assert calls[0][2][:3] == config["cuda"][compiler]
+ assert any("--register-usage-level=2" in flag for flag in calls[0][2])
+ disabled = resolve_backend_config({"arch": "sm_100a", "nvrtc": [], "ptxas": []})
+ assert compiler_options(disabled) == []
+ linked = compile_source("#include <nvshmem.h>", config)
+ assert linked.target_format == "cubin"
+
+
+@pytest.mark.parametrize("tag_target", ["name", "dict", "export"])
+def test_target_tag_defaults_and_precedence(tag_target):
+ name = "test/backend-config"
+ tvm.target.tag.register_tag(
+ name,
+ {
+ "kind": "cuda",
+ "arch": "sm_100a",
+ "backend_config": {"cuda": {"compiler": "nvcc", "nvcc": ["-O1"]}},
+ },
+ override=True,
+ )
+ target = tvm.target.Target(name if tag_target == "name" else {"tag": name})
+ if tag_target == "export":
+ target = tvm.target.Target(target.export())
+ resolved = resolve_backend_config(target=target)
+ assert resolved["compiler"] == "nvcc" and resolved["nvcc"] == ["-O1"]
+ assert resolved["ptxas"][-1] == "--register-usage-level=10"
+ override = resolve_backend_config({"compiler": "nvrtc", "nvcc": []}, target)
+ assert override["compiler"] == "nvrtc" and override["nvcc"] == []
+ assert resolve_backend_config(target=target) == resolved
+
+
+def test_entry_ir_roundtrip_and_resolution():
+ @T.function
+ def kernel(A: T.Tensor((32,), "float32")):
+ T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={
+ "cuda": {
+ "arch": "sm_90a",
+ "nvcc": ["--use_fast_math", "--ftz=false"],
+ "nvrtc": ["--use_fast_math", "--ftz=false"],
+ }
+ },
+ )
+ A[T.cuda.thread_idx("x")] = T.float32(1)
+
+ tvm.ir.assert_structural_equal(
+ 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(
+ {
+ "cuda": {
+ "arch": "sm_100a",
+ "nvcc": ["--use_fast_math", "--generate-line-info"],
+ "nvrtc": ["--use_fast_math", "--generate-line-info"],
+ }
+ }
+ )(mod)
+ entries = []
+ tvm_ffi.structural_walk(
+ mod,
+ lambda node: (
+ entries.append(node)
+ if isinstance(node, tvm.ir.RegionStmt) and node.op.name == "tirx.device_entry"
+ else None
+ ),
+ )
+ assert len(entries) == 1
+ entry = entries[0]
+ assert entry.attrs["target"].arch == "sm_90a"
+ config = parse_backend_config(entry.attrs["backend_config"])
+ assert "--ftz=false" in config["cuda"]["nvrtc"]
+ 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"])
+def test_backend_config_controls_actual_ftz(compiler):
+ _require_compiler(compiler)
+
+ @T.function
+ def kernel(A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ x = T.cuda.thread_idx("x")
+ B[x] = A[x] * T.float32(0.5)
+
+ arch = env.cuda_arch(0)
+ inp = np.full(32, np.finfo(np.float32).tiny, dtype="float32")
+ a = tvm.runtime.tensor(inp, tvm.cuda())
+ outputs = []
+ 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()],
+ }
+ },
+ )
+ b = tvm.runtime.empty((32,), "float32", tvm.cuda())
+ module.mod["kernel"](a, b)
+ outputs.append(b.numpy())
+ metadata = parse_backend_config(module.mod.imports[0].inspect_source("backend_config"))
+ assert f"--ftz={str(ftz).lower()}" in metadata["cuda"][compiler]
+ assert metadata["cuda"]["compiler"] == compiler
+ np.testing.assert_array_equal(outputs[0], inp * np.float32(0.5))
+ np.testing.assert_array_equal(outputs[1], np.zeros_like(inp))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not tvm.cuda().exist, reason="requires CUDA")
+def test_two_entries_keep_compilers_math_and_order():
+ _require_compiler("nvrtc")
+ _require_compiler("nvcc")
+
+ @T.function
+ def kernel(
+ A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float32"), C: T.Tensor((32,), "float32")
+ ):
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={
+ "cuda": {
+ "compiler": "nvrtc",
+ "nvcc": ["--use_fast_math", "--ftz=false"],
+ "nvrtc": ["--use_fast_math", "--ftz=false"],
+ }
+ },
+ ):
+ x = T.cuda.thread_idx("x")
+ B[x] = A[x] * T.float32(0.5)
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={
+ "cuda": {
+ "compiler": "nvcc",
+ "nvcc": ["--use_fast_math", "--ftz=true"],
+ "nvrtc": ["--use_fast_math", "--ftz=true"],
+ }
+ },
+ ):
+ y = T.cuda.thread_idx("x")
+ C[y] = B[y] * T.float32(0.5)
+
+ arch = env.cuda_arch(0)
+ module = tvm.compile(kernel, 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"}
+ inp = np.full(32, np.finfo(np.float32).tiny, dtype="float32")
+ a = tvm.runtime.tensor(inp, tvm.cuda())
+ b = tvm.runtime.empty((32,), "float32", tvm.cuda())
+ c = tvm.runtime.empty((32,), "float32", tvm.cuda())
+ module.mod["kernel"](a, b, c)
+ np.testing.assert_array_equal(b.numpy(), inp * np.float32(0.5))
+ np.testing.assert_array_equal(c.numpy(), np.zeros_like(inp))
+
+
+@pytest.mark.parametrize("compiler", ["nvrtc", "nvcc"])
+def test_concurrent_source_compilation(compiler):
+ _require_compiler(compiler)
+ from tvm.backend.cuda.compiler import compile_source
+
+ source = 'extern "C" __global__ void half_value(float* x) { x[0] = x[0] * 0.5f; }'
+ base = {"cuda": {"arch": "sm_100a", "compiler": compiler, "target_format": "ptx"}}
+ configs = [
+ merge_backend_configs(
+ base,
+ {
+ "cuda": {
+ "nvcc": ["--use_fast_math", "--ftz" + "=" + str(ftz).lower()],
+ "nvrtc": ["--use_fast_math", "--ftz" + "=" + str(ftz).lower()],
+ }
+ },
+ )
+ for ftz in (False, True, False, True)
+ ]
+ with ThreadPoolExecutor(max_workers=4) as pool:
+ results = list(pool.map(lambda cfg: compile_source(source, cfg), configs))
+ for ftz, cfg, result in zip((False, True, False, True), configs, results):
+ assert result.config["cuda"][compiler] == cfg["cuda"][compiler]
+ assert (b"mul.ftz.f32" in result.binary) == ftz
+ assert results[0].config == results[2].config
+ assert results[0].binary == results[2].binary
+ assert results[0].binary != results[1].binary
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not tvm.cuda().exist, reason="requires CUDA")
+@pytest.mark.parametrize("fallback", [False, True])
+@pytest.mark.parametrize("load_target", ["llvm", "cuda"])
+def test_serialized_config_replays_in_fresh_process(tmp_path, monkeypatch, fallback, load_target):
+ _require_compiler("nvrtc")
+
+ @T.function
+ def main(A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ x = T.cuda.thread_idx("x")
+ B[x] = A[x] * T.float32(0.5)
+
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1" if fallback else "0")
+ config = {
+ "cuda": {
+ "arch": env.cuda_arch(0),
+ "nvcc": ["--use_fast_math", "--ftz=false", "--generate-line-info"],
+ "nvrtc": ["--use_fast_math", "--ftz=false", "--generate-line-info"],
+ }
+ }
+ module = tvm.compile(main, backend_config=config)
+ library = str(tmp_path / "replay.so")
+ module.export_library(library)
+ script = """
+import sys
+import numpy as np
+import tvm
+# The loader's Target may describe unrelated work or provide incompatible
+# defaults. Source replay must use the configuration saved in the artifact.
+with tvm.target.Target(LOAD_TARGET):
+ module = tvm.runtime.load_module(sys.argv[1])
+from tvm.backend.config import parse_backend_config
+config = parse_backend_config(module.imports[0].inspect_source("backend_config"))["cuda"]
+assert "--ftz=false" in config["nvrtc"] and "--generate-line-info" in config["nvrtc"]
+x = np.full(32, np.finfo(np.float32).tiny, dtype="float32")
+a = tvm.runtime.tensor(x, tvm.cuda())
+b = tvm.runtime.empty((32,), "float32", tvm.cuda())
+module["main"](a, b)
+np.testing.assert_array_equal(b.numpy(), x * np.float32(0.5))
+"""
+ load_context = (
+ "llvm"
+ if load_target == "llvm"
+ else {
+ "kind": "cuda",
+ "arch": "sm_90a",
+ "backend_config": {"cuda": {"target_format": "fatbin"}},
+ }
+ )
+ script = script.replace("LOAD_TARGET", repr(load_context))
+ environment = dict(os.environ, TVM_COMPILE_FORCE_FALLBACK="0")
+ result = subprocess.run(
+ [sys.executable, "-c", script, library], env=environment, text=True, capture_output=True
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_device_helpers_are_cloned_per_configuration():
+ _require_compiler("nvrtc")
+ from tvm.backend.cuda.transforms import SpecializeDeviceHelpers
+ from tvm.script import ir as I
+
+ @I.ir_module
+ class Module:
+ @T.function(private=True)
+ def leaf(x: T.float32) -> T.float32:
+ return x * T.float32(0.5)
+
+ @T.function(private=True)
+ def helper(x: T.float32) -> T.float32:
+ return Module.leaf(x)
+
+ @T.function
+ def kernel(A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float32")):
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={
+ "cuda": {
+ "arch": "sm_90a",
+ "nvcc": ["--use_fast_math", "--ftz=false"],
+ "nvrtc": ["--use_fast_math", "--ftz=false"],
+ }
+ },
+ ):
+ x = T.cuda.thread_idx("x")
+ A[x] = Module.helper(A[x])
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={
+ "cuda": {
+ "arch": "sm_100a",
+ "nvcc": ["--use_fast_math", "--ftz=true"],
+ "nvrtc": ["--use_fast_math", "--ftz=true"],
+ }
+ },
+ ):
+ y = T.cuda.thread_idx("x")
+ B[y] = Module.helper(B[y])
+
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}, host="llvm")
+ mod = tvm.tirx.transform.BindTarget(target)(Module)
+ mod = BindBackendConfig()(mod)
+ mod = tvm.tirx.transform.LowerTIRx()(mod)
+ mod = tvm.tirx.transform.SplitHostDevice()(mod)
+ mod = SpecializeDeviceHelpers()(mod)
+ groups = {}
+ for gv, fn in mod.functions.items():
+ if fn.attrs["target"].host is not None:
+ continue
+ config = parse_backend_config(fn.attrs["backend_config"])
+ groups.setdefault(
+ (config["cuda"]["arch"], "--ftz=true" in config["cuda"]["nvrtc"]), []
+ ).append(gv.name_hint)
+ assert set(groups) == {("sm_90a", False), ("sm_100a", True)}
+ assert all(len(names) == 3 for names in groups.values()), mod.script()
+ assert set(groups[("sm_90a", False)]).isdisjoint(groups[("sm_100a", True)])
+ dispatched = set()
+
+ @tvm.transform.pass_instrument
+ class CheckHelperTargets:
+ def run_before_pass(self, lowered, info):
+ if info.name == "tirx.TileDispatch":
+ for gv, fn in lowered.functions.items():
+ if gv.name_hint.startswith(("helper", "leaf")):
+ dispatched.add(fn.attrs["target"].arch)
+ assert gv.name_hint not in ("helper", "leaf")
+
+ with tvm.transform.PassContext(instruments=[CheckHelperTargets()]):
+ built = tvm.compile(Module)
+ assert dispatched == {"sm_90a", "sm_100a"}
+ assert len(built.mod.imports) == 2
+ assert {
+ parse_backend_config(m.inspect_source("backend_config"))["cuda"]["arch"]
+ for m in built.mod.imports
+ } == {"sm_90a", "sm_100a"}
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not tvm.cuda().exist, reason="requires CUDA")
+@pytest.mark.parametrize("fallback", [False, True])
+def test_cuda_host_embeds_independent_binaries_in_plain_cpp(tmp_path, monkeypatch, fallback):
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1" if fallback else "0")
+ _require_compiler("nvrtc")
+ _require_compiler("nvcc")
+ import tvm_ffi.cpp
+
+ from tvm.backend.cuda.host import export_cuda_host
+ from tvm.support.nvcc import find_cuda_path
+
+ @T.function
+ def main(
+ A: T.Tensor((32,), "float32"),
+ B: T.Tensor((32,), "float32"),
+ C: T.Tensor((32,), "float32"),
+ factor: T.float16,
+ other: T.bfloat16,
+ ):
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={
+ "cuda": {
+ "compiler": "nvrtc",
+ "nvcc": ["--use_fast_math", "--ftz=false"],
+ "nvrtc": ["--use_fast_math", "--ftz=false"],
+ }
+ },
+ ):
+ x = T.cuda.thread_idx("x")
+ B[x] = A[x] * T.Cast("float32", factor)
+ with T.device_entry(
+ launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={
+ "cuda": {
+ "compiler": "nvcc",
+ "nvcc": ["--use_fast_math", "--ftz=true"],
+ "nvrtc": ["--use_fast_math", "--ftz=true"],
+ }
+ },
+ ):
+ y = T.cuda.thread_idx("x")
+ C[y] = B[y] * T.Cast("float32", other)
+
+ target = tvm.target.Target({"kind": "cuda", "arch": env.cuda_arch(0)}, host="cuda_host")
+ module = tvm.compile(main, target=target)
+ with tvm.target.Target("llvm"):
+ source = export_cuda_host(module.mod)
+ assert "__global__" not in source
+ cuda = find_cuda_path()
+ library = tvm_ffi.cpp.build_inline(
+ name="independent_cuda_entries",
+ cpp_sources=source,
+ backend="cpu",
+ extra_include_paths=[cuda + "/include"],
+ extra_ldflags=["-L" + cuda + "/lib64", "-lcuda", "-lcudart"],
+ build_directory=str(tmp_path),
+ )
+ script = """
+import sys
+import torch
+import tvm_ffi
+assert 'tvm' not in sys.modules
+module = tvm_ffi.load_module(sys.argv[1])
+a = torch.full((32,), torch.finfo(torch.float32).tiny, device='cuda')
+a[1::2] = 2.0
+b, c = torch.empty_like(a), torch.empty_like(a)
+module['main'](a, b, c, 0.5, 0.5)
+assert torch.equal(b, a * 0.5)
+expected = torch.zeros_like(c)
+expected[1::2] = 0.5
+assert torch.equal(c, expected)
+assert 'tvm' not in sys.modules
+"""
+ result = subprocess.run([sys.executable, "-c", script, library], text=True, capture_output=True)
+ assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_legacy_payloads():
+ import struct
+
+ load = tvm.get_global_func("ffi.Module.load_from_bytes.cuda", allow_missing=True)
+ if load is None:
+ pytest.skip("CUDA runtime module loader is not linked")
+
+ def string(value):
+ return struct.pack("<Q", len(value)) + value
+
+ code = b".version 8.0\n.target sm_80\n.address_size 64\n"
+ old_binary = string(b"ptx") + struct.pack("<Q", 0) + string(code)
+ module = load(old_binary)
+ assert module["__tvm_cuda_binary"]()[0] == code
+ old_source = string(b"cuda") + struct.pack("<Q", 0) + string(b"// legacy source")
+ from tvm.support.nvcc import tvm_callback_cuda_compile
+
+ seen = []
+ try:
+
+ @tvm_ffi.register_global_func("tvm_callback_cuda_compile", override=True)
+ def mock_compile(source, config):
+ seen.append(parse_backend_config(config))
+ return [bytearray(code), "ptx"]
+
+ assert load(old_source)["__tvm_cuda_binary"]()[0] == code
+ assert seen == [{}]
+ finally:
+ tvm_ffi.register_global_func(
+ "tvm_callback_cuda_compile", tvm_callback_cuda_compile, override=True
+ )
+ snapshot = backend_config_json({"cuda": {"arch": "sm_100a", "nvrtc": []}})
+ current = old_binary + string(snapshot.encode())
+ assert load(current).inspect_source("backend_config") == snapshot
+ for malformed in (old_binary + b"x", current[:-1], current + b"x"):
+ with pytest.raises(tvm.error.InternalError):
+ load(malformed)
+
+
+def test_generic_cuda_target_without_gpu(tmp_path):
+ script = """
+import tvm
+import tvm.testing
+from tvm.backend.cuda import BackendConfig
+from tvm.backend.config import prepare_target
+from tvm.backend.cuda.transforms import BindBackendConfig
+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")})
+assert resolved.arch == "sm_100a"
+assert resolved.host.kind.name == "llvm"
+
+@T.function
+def main(A: T.Tensor((32,), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ A[T.cuda.thread_idx("x")] = T.float32(1)
+
+mod = tvm.tirx.transform.BindTarget(target)(tvm.IRModule({"main": main}))
+try:
+ BindBackendConfig()(mod)
+except ValueError as error:
+ assert "explicit" in str(error)
+else:
+ raise AssertionError("offline compilation must not guess an architecture")
+"""
+ path = tmp_path / "generic_cuda.py"
+ path.write_text(script)
+ environment = dict(os.environ, CUDA_VISIBLE_DEVICES="", TVM_TEST_TARGETS="llvm;cuda")
+ result = subprocess.run(
+ [sys.executable, str(path)], env=environment, text=True, capture_output=True
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+
+
+@pytest.mark.parametrize("target", [None, "cuda", {"kind": "cuda"}, "target_object"])
+def test_entry_architectures_compile_without_gpu_detection(tmp_path, target):
+ _require_compiler("nvcc")
+ script = """
+import tvm
+from tvm.script import tirx as T
+from tvm.backend.cuda import BackendConfig
+@T.function
+def main(A: T.Tensor((32,), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32),
+ backend_config={"cuda": T.cuda.BackendConfig(arch="sm_100a", compiler="nvcc")})
+ A[T.cuda.thread_idx("x")] = T.float32(1)
+target = TARGET
+if target == "target_object":
+ target = tvm.target.Target("cuda", host="llvm")
+module = tvm.compile(main, target=target)
+metadata = module.mod.imports[0].inspect_source("backend_config")
+from tvm.backend.config import parse_backend_config
+assert parse_backend_config(metadata)["cuda"]["arch"] == "sm_100a"
+assert "arch" not in tvm.target.Target("cuda").attrs
+"""
+ # Use a source file because the TIRx construction decorator reads its definition.
+ path = tmp_path / "offline.py"
+ path.write_text(script.replace("TARGET", repr(target)))
+ environment = dict(os.environ, CUDA_VISIBLE_DEVICES="")
+ result = subprocess.run(
+ [sys.executable, str(path)], env=environment, text=True, capture_output=True
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_source_fallback_does_not_invoke_compiler(monkeypatch):
+ from tvm.support import nvcc
+
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+
+ def unavailable(*args, **kwargs):
+ raise AssertionError("source fallback must not require a CUDA compiler")
+
+ monkeypatch.setattr(nvcc, "_compile_cuda_nvcc", unavailable)
+ monkeypatch.setattr(nvcc, "_compile_cuda_nvrtc", unavailable)
+
+ @T.function
+ def main(A: T.Tensor((32,), "float32")):
+ 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"}})
+ assert "__global__" in module.mod.imports[0].inspect_source("cuda")
+ assert (
+ parse_backend_config(module.mod.imports[0].inspect_source("backend_config"))["cuda"]["arch"]
+ == "sm_100a"
+ )
+
+
+def test_identical_entries_share_one_module(monkeypatch):
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+ config = {"cuda": {"arch": "sm_100a", "nvrtc": ["--ftz=false"]}}
+
+ @T.function
+ def main(A: T.Tensor((32,), "float32")):
+ with T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32), backend_config=config):
+ A[T.cuda.thread_idx("x")] = T.float32(1)
+ with T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32), backend_config=config):
+ A[T.cuda.thread_idx("x")] = T.float32(2)
+
+ config["cuda"]["nvrtc"].clear()
+ module = tvm.compile(main, 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"]
+
+
+def test_build_snapshots_config_before_passes(monkeypatch):
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+ config = {"cuda": {"arch": "sm_100a", "nvrtc": ["--ftz=false"]}}
+
+ @T.function
+ def main(A: T.Tensor((32,), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
+ A[T.cuda.thread_idx("x")] = T.float32(1)
+
+ @tvm.transform.pass_instrument
+ class ChangeCallerConfig:
+ def run_before_pass(self, mod, info):
+ config["cuda"]["nvrtc"] = ["--ftz=true"]
+
+ with tvm.transform.PassContext(instruments=[ChangeCallerConfig()]):
+ module = tvm.compile(main, 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 a58144e58c..5352b09e08 100644
--- a/tests/python/tirx/test_cuda_launch.py
+++ b/tests/python/tirx/test_cuda_launch.py
@@ -165,6 +165,7 @@ def _compile(kernel, host, tmp_path):
name="launch_" + kernel.attrs["global_symbol"],
cuda_sources=source,
extra_cuda_cflags=[f"-arch={arch}"],
+ extra_ldflags=["-lcuda"],
build_directory=str(tmp_path),
backend="cuda",
)
@@ -325,4 +326,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, target="cuda")
+ tvm.compile(kernel, backend_config={"cuda": {"arch": "sm_80"}})
diff --git a/tests/python/tirx/test_tirx_kernels_registry_correctness.py b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
index c198f89404..21db5e2a7e 100644
--- a/tests/python/tirx/test_tirx_kernels_registry_correctness.py
+++ b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
@@ -127,31 +127,6 @@ def _required_cuda_device_count(config):
return int(config.get("num_processes", config.get("world_size", 1)))
-@contextmanager
-def _current_cuda_prepare_arch():
- try:
- import torch
- except ImportError:
- yield
- return
-
- arch = env.cuda_arch(torch.cuda.current_device()) if torch.cuda.is_available() else None
- if arch is None:
- yield
- return
-
- variable = kernel_runner.PREPARE_CUDA_ARCH_ENV
- previous = os.environ.get(variable)
- os.environ[variable] = arch
- try:
- yield
- finally:
- if previous is None:
- os.environ.pop(variable, None)
- else:
- os.environ[variable] = previous
-
-
@contextmanager
def _registry_gpu_lock(kernel_name, config):
try:
@@ -211,5 +186,10 @@ def test_manifest_tirx_kernel_correctness(kernel_name, config):
"MegaMoE requires its dedicated multi-process scheduler; this suite's "
"processes own CUDA contexts that its physical-device assignment rejects"
)
- with _registry_gpu_lock(kernel_name, config), _current_cuda_prepare_arch():
- kernel_runner.run_kernel_test(kernel_name, config, registry=_KERNELS)
+ with _registry_gpu_lock(kernel_name, config):
+ import torch
+
+ backend_config = {"cuda": {"arch": env.cuda_arch(torch.cuda.current_device())}}
+ kernel_runner.run_kernel_test(
+ kernel_name, config, registry=_KERNELS, backend_config=backend_config
+ )