40 Commits

Author SHA1 Message Date
Tianqi Chen 9bfefb7e4b [TIRx] Introduce first-class Return statement (#20018)
Return is control flow, but TIRx currently represents it as an
Evaluate-wrapped intrinsic call. This prevents return values from
participating naturally in statement traversal and requires special-case
handling across the pipeline.

This change introduces a reflected tirx.Return statement carrying an
Expr, wires it through TVMScript, statement visitors and mutators,
lowering, storage planning, and C/LLVM code generation, and removes the
legacy tirx.ret and T.ret surfaces.
2026-07-16 17:34:50 -04:00
Shushi Hong 1a764d7993 [Tests] Reduce runtime of slow Python tests (#20006)
This PR reduces the runtime of several slow Python test groups:

- Parameterize LLVM division and CUDA vectorized-cast cases so
pytest-xdist can schedule them independently.
- Replace exhaustive ONNX execution with structural importer checks plus
representative numerical cases, and avoid registering unsupported
backend cases.
- Reuse compiled paged-attention kernels across compatible test cases.

Targeted measurements showed:

- LLVM division: 25.34s → 19.87s
- CUDA vectorized casts: 142.85s → 108.40s
- Paged-attention CPU: 316.32s → 192.50s
- ONNX Conv: 25.84s → 1.81s
- ONNX Reduce: 20.66s → 4.04s

This PR also fixes a latent CUDA Graph cleanup bug that could leave
`cudaErrorStreamCaptureInvalidated` in the worker thread and cause
unrelated subsequent GPU tests to fail.
2026-07-16 06:07:20 +08:00
Tianqi Chen e479a5dbe7 [RUNTIME][PYTHON] Add explicit Target device conversion (#20005)
## Summary

Compiler Targets can carry device-type semantics that runtime
device-name parsing does not preserve.

- add `tvm.device_from_target` for canonical Target-to-Device
translation
- use explicit runtime constructors where the device kind is fixed
- update target-derived utilities, tests, and documentation to use the
explicit boundary
2026-07-15 05:34:21 +08:00
Tianqi Chen 3452fd4ffa [TEST] Serialize local GPU execution under pytest-xdist (#19942)
Add tvm.testing.run_with_gpu_lock backed by the existing
tvm_ffi.utils.FileLock. Migrate live local GPU tests to acquire the
machine-local lock around device execution, synchronization, host
transfer, and checks while leaving target construction and compilation
outside the critical section.

Replace the custom xdist scheduler with standard xdist_group placement
for the order-dependent test family. RPC tests retain dynamic port
allocation and per-test process isolation rather than gaining a broad
category lock.
2026-07-04 17:49:45 -04:00
Shushi Hong a8635b04cc [Tests] Migrate off tvm.testing.parametrize_targets to native pytest (#19826)
This pr moves target selection and per-target device gating off the TVM
pytest plugin onto plain pytest, and remove the now-dead machinery.
2026-06-18 08:48:53 -04:00
Shushi Hong e4da848e57 [Tests] Modernize test gating (#19777)
This pr modernizes test gating. It replaces the heavy
`tvm.testing.Feature` machinery with a thin `tvm.testing.env` module of
`has_*()` capability probes, used via standard pytest.mark + skipif. And
markers move to `pyproject.toml`
2026-06-15 18:50:57 -04:00
Tianqi Chen 96cba60464 [PYTHON] Autoload backends; simplify library loading; remove TVMError for native errors (#19727)
This PR adds an autoload mechanism for out-of-tree backends, simplifies
TVM's Python library loading, and removes `TVMError` in favor of native
Python errors.

## Autoload out-of-tree backends

Out-of-tree packages can register an autoload callable under the
`tvm.backends` entry-point group (mirroring torch's device-backend
autoload). At `import tvm` startup each entry point is discovered and
its callable invoked once, after the core runtime and the `tvm`
namespace are fully initialized, so an extension can register
ops/targets/funcs or load extra libraries.

```toml
[project.entry-points."tvm.backends"]
tvm_foo = "tvm_foo:_autoload"
```

A failing extension is caught and surfaced via `warnings.warn` so it
cannot break `import tvm`. Autoload can be disabled with
`TVM_DEVICE_BACKEND_AUTOLOAD=0`.

## Simplify library loading

The library-loading path in `base.py` is consolidated around a single
`_LOADED_LIBS` dict (basename to ctypes handle) so downstream and
autoloaded extensions can skip already-loaded libraries; the per-backend
runtime DSO list is folded into `load_backend_libs`. Accumulated cruft
is removed: the Python-3.9 check, the readline shim, the `_FFI_MODE`
ctypes check, the `base.__version__` re-export, and `py_str` (call sites
inline `.decode("utf-8")`).

## Remove TVMError in favor of native Python errors

`TVMError` added a layer atop `RuntimeError` that downstream code had to
import and learn. It is removed; the registered FFI error kinds
(`InternalError`, `RPCError`, `OpError`, `DiagnosticError`,
`ScheduleError`) now subclass `RuntimeError` directly while staying
registered, so the FFI keeps throwing the right kinds. All `TVMError`
imports, `except`/`raise`/`isinstance` uses, and
`pytest.raises(tvm.TVMError)` sites move to the `RuntimeError` builtin.
2026-06-11 13:50:38 -04:00
Tianqi Chen ffea531107 [REFACTOR][PYTHON] Lift compiler/CLI/process modules from tvm.contrib to tvm.support (#19624)
## Summary

Lifts 10 host-toolchain / CLI / process / utility modules from
`python/tvm/contrib/` to a new `python/tvm/support/` package, and
deletes two dead contrib shims.

`tvm.support` is the home for Python helpers that integrate TVM with
external CLIs and host-side tools — compilers, archivers, subprocess
pools, and build-info queries. These are load-bearing internal pieces
that TVM's compile/link/run paths depend on. `tvm.contrib` is reserved
for optional vendor SDK integrations and experimental features. The
distinction is documented in the `tvm.support` package docstring.

Moved (one commit each):

- `tvm.contrib.cc` → `tvm.support.cc`
- `tvm.contrib.nvcc` → `tvm.support.nvcc`
- `tvm.contrib.rocm` → `tvm.support.rocm`
- `tvm.contrib.ndk` → `tvm.support.ndk`
- `tvm.contrib.xcode` → `tvm.support.xcode`
- `tvm.contrib.clang` → `tvm.support.clang`
- `tvm.contrib.emcc` → `tvm.support.emcc`
- `tvm.contrib.popen_pool` → `tvm.support.popen_pool`
- `tvm.contrib.utils` → `tvm.support.utils`
- `tvm.contrib.tar` → `tvm.support.tar`

Deleted:
- `tvm.contrib.spirv` — single `optimize()` wrapping `spirv-opt`; zero
importers.
- `tvm.contrib.rpc` — self-deprecation shim with "removed in 0.5"
banner; honoring it.

Package conversion:
- `python/tvm/support.py` → `python/tvm/support/__init__.py` with
inclusion-rule docstring.
- `libinfo()` extracted into `python/tvm/support/libinfo.py`.
- `FrontendTestModule` dropped (audit confirmed zero callers outside its
own definition).

## Compatibility

Hard break — no `tvm.contrib.<mod>` re-export shims. All callers updated
in this PR.

C++-side FFI registry keys (`tvm.contrib.nvcc.*`, etc.) are unchanged —
only the Python module path moves. Renaming the FFI keys is a separate
follow-up.
2026-05-27 15:31:12 -04:00
Bohan Hou 859498dc01 [TIRx] Bringup TIRx Infrastructure (#19581)
## Summary

This PR adds the initial TIRx support needed for low-level programming
of Blackwell-class GPU architectures. As part of the ongoing TIRx
refactor, it introduces TVMScript support for directly scripting
advanced hardware features without relying on scheduling as the primary
programming interface.

The change keeps existing `s_tir` script support intact while making
direct scripting a first-class path for TIRx programs.

## Main Changes

- Add TIRx operator dispatch and layout infrastructure.
- Add TVMScript support for new low-level TIRx operations.
- Add analysis, transform, and lowering support for TIRx IR nodes.
- Add CUDA/Blackwell-oriented codegen and intrinsic coverage.
- Add Python and C++ integration points for TIRx scripting and runtime
support.

## Validation

- `pre-commit run --all-files`
- `ninja -C build -j32`
- `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16`
  - `1723 passed, 47 skipped, 32 warnings`
- `CUDA_VISIBLE_DEVICES=2 python -m pytest -v
tests/python/all-platform-minimal-test`
  - `37 passed, 105 skipped`
- `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis
tests/python/tirx-base tests/python/tirx-transform -n 16`
  - `664 passed, 25 skipped, 9 xfailed, 1 xpassed`

## Local CI Notes

Some full CI-equivalent jobs were not locally reproducible because this
machine is missing parts of the Apache TVM CI environment, including
`llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten,
and ARM/QEMU cross-toolchain components. Metal-specific tests were
skipped locally because no Metal runtime is available.
2026-05-18 16:44:43 -07:00
Tianqi Chen 6e8f77d664 [REFACTOR][RUNTIME][CODEGEN] Backend specific target and runtime to enable cross-compile fallback (#19465)
## Why

This refactor reshapes each backend into a self-contained
`src/target/<X>/`
cluster (with optional `src/target/<X>/llvm/` for LLVM-dependent
codegen) and
introduces a per-backend fallback module that absorbs the cross-compile
role
cleanly — without `target/opt/` stubs, without `DeviceSourceModuleNode`,
and
without leaking a synthetic `kind()` to consumers.

## High-level principles

- **One directory per backend.** All codegen-side files for backend
`<X>` live
under `src/target/<X>/`. Optional `src/target/<X>/llvm/` subdir for
files
  that require `USE_LLVM` at build time. Backend grouping wins over
build-dependency grouping (the latter being upstream's `target/llvm/`).
- **Plugin-only runtime modules.** `src/runtime/<X>/<X>_module.h` is
deleted.
The runtime's real `<X>ModuleNode` is reachable only via the FFI
registry
  (`ffi.Module.create.<kind>`, `ffi.Module.load_from_bytes.<kind>`). No
  C++ API surface other than the static registrations.
- **Per-backend fallback module for cross-compile.** Each `<X>` gets a
`<X>FallbackModuleNode` in `src/target/<X>/<X>_fallback_module.{h,cc}`.
Same `kind()` as the real module. Codegen-time only — never reachable
via
load. `GetFunction` errors with a backend-specific "runtime not linked"
  message; `InspectSource` works.
- **Codegen-side wrapper does the fallback selection.** Codegen calls
`<X>ModuleCreateWithFallback(...)`, which tries
`ffi.Module.create.<kind>`
  via the registry; on miss, falls through to `<X>FallbackModuleCreate`
(plain C++; reachable directly from the fallback header). When
`USE_<X>=ON`
  is in effect, the registry hit returns the real module; when
  `USE_<X>=OFF`, the fallback is what codegen gets. No CMake `if/else`
  gating; fallback always compiled.

## Specific changes

### New per-backend directories (codegen + fallback)

- `src/target/cuda/` — `codegen_cuda.cc` + `intrin_rule_cuda.cc` +
fallback module pair + `llvm/codegen_nvptx.cc`
- `src/target/rocm/` — fallback module pair + `llvm/codegen_amdgpu.cc` +
`llvm/intrin_rule_rocm.cc`
- `src/target/hexagon/` — fallback module pair +
`llvm/codegen_hexagon.cc` + `llvm/intrin_rule_hexagon.cc`
- `src/target/metal/` — `codegen_metal.cc` + `intrin_rule_metal.cc` +
fallback module pair
- `src/target/vulkan/` — `build_vulkan.cc` + the rest of `target/spirv/`
absorbed + fallback module pair
- `src/target/opencl/` — `codegen_opencl.cc` + `intrin_rule_opencl.cc` +
fallback module pair
- `src/target/webgpu/` — `codegen_webgpu.cc` + fallback module pair (was
`WebGPUSourceModuleNode`, renamed)

### New fallback classes

`<X>FallbackModuleNode` for X in {`CUDA`, `ROCm`, `Hexagon`, `Metal`,
`Vulkan`, `OpenCL`, `WebGPU`}. Each:
- `kind()` matches the real backend
- Stores `(code or smap, fmt, fmap, source)` — no driver/runtime calls
- `GetFunction` errors with backend-specific "runtime not linked"
message
- `InspectSource` works fully
- `SaveToBytes` byte-identical to real
2026-04-29 14:48:31 -04:00
Tianqi Chen 141c22fd8a [Refactor] Bring up tirx namespace (#18913)
This PR brings up the tirx namespace. We have been spliting out the
original tir namespace to include high-level component s_tir and this PR
updates the remaining low-level part as tirx namespace
2026-03-19 21:27:54 -07:00
Tianqi Chen 0fba1606be [REFACTOR][TIR] Introduce AllocBuffer and phase out Allocate+DeclBuffer (#18865)
## Summary

This PR introduces `AllocBufferNode`/`AllocBuffer` as a single TIR
statement that both allocates memory and declares a buffer into scope.
This replaces the previous pattern of `Allocate(var, dtype, shape, cond,
DeclBuffer(buf, body))` with the simpler `AllocBuffer(buf, body)`.

### Main changes

- **New IR node** `AllocBufferNode` with fields `{buffer, annotations,
body}` — same semantics as `DeclBuffer` but also allocates memory
- **TVMScript**: `T.alloc_buffer(shape, dtype, scope)` now emits
`AllocBuffer` directly (statement-level allocation).
`T.sblock_alloc_buffer(...)` for SBlock-level buffer allocation (full
parameter set)
- **All codegen backends** (C, CUDA, Metal, OpenCL, WebGPU, LLVM, NVPTX,
AMDGPU, SPIR-V) updated to handle `AllocBufferNode`
- **All TIR transforms** (storage_rewrite, flatten_buffer,
vectorize_loop, lower_warp_memory, etc.) updated
- **All S-TIR transforms** (compact_buffer_region, merge_shared_memory,
inject_double_buffer, etc.) updated
- **Removed `AllocateNode`** entirely — `AllocBuffer` is now the sole
allocation primitive
- **Removed `AllocDescriptor`** from merge_shared_memory_allocations —
uses `Buffer` objects directly
- **Added `AllocBuffer::ConstantAllocationSize()`** inline helper method

### Design rationale

The old `Allocate + DeclBuffer` pair was a historical artifact:
`AllocateNode` stored raw fields (`buffer_var`, `dtype`, `extents`,
`condition`) separate from the `Buffer` object, requiring pattern
matching (`IsAllocateDeclBufferPattern`) to reconstruct the buffer
association. `AllocBuffer` unifies this into a single node with a proper
`Buffer` reference, simplifying codegen backends and transform passes.

225 files changed, ~3500 insertions/deletions (net near-zero, mostly
mechanical migration).

## Test plan

- [x] All TIR base tests pass
- [x] All TIR transform tests pass
- [x] TVMScript roundtrip tests pass
- [x] S-TIR transform tests pass
- [x] Codegen tests pass
- [x] All-platform minimal tests pass
- [x] C++ functor tests pass
- [x] Pre-commit clean (clang-format, ruff, etc.)
2026-03-04 11:59:20 -05:00
Tianqi Chen 33dcea1686 [REFACTOR][LINT] Modernize ruff config (#18810)
This PR removes the extra lint violations from the codebase so lint
aligns with the latest style
2026-02-23 07:29:21 -05:00
Tianqi Chen aa2e609136 [LINT] Modernize lint to use pre-commit hooks (#18807)
This PR migrates existing lint to use pre-commit hooks
2026-02-22 11:03:21 -05:00
Tianqi Chen 283fd19683 [REFACTOR][TARGET] Further cleanup target python api (#18793)
This PR cleans up the target python api.

- Removes the indirections of attribute exposure
- Move tag registry to python so it is easily configurable
- Remove legacy constructors in favor of tags
2026-02-17 21:39:27 -05:00
Tianqi Chen 3dce4aedb5 [REFACTOR][S-TIR] Lift transform passes to s_tir namespace (#18722)
This PR migrates the s_tir related transform passes into s_tir namespace
instead. This set of changes can minimize the overall tir namespace to
make it more focused.
2026-02-07 12:01:13 -05:00
Tianqi Chen 198df475fa [REFACTOR][TEST] Migrate all codegen test to tvmscript (#18719)
This PR migrates all the codegen tests to explicitly using tvmscript
instead of indirectly via s_tir.Schedule. They makes the test surface
more unit, contains less dep and more maintainable.
2026-02-07 07:09:59 -05:00
Tianqi Chen d76c729259 [REFACTOR][S-TIR] Initialize the s_tir module (#18712)
This PR initalizes the s_tir for scheduable TensorIR. The change mainly
starts from python side, the we will gradually move towards the c++ side
in followup PRs. The python main change:

tir.Schedule => s_tir.Schedule
2026-02-05 09:40:31 -05:00
Kathryn (Jinqi) Chen fa905d2b69 [Compile] accelerate compilation speed using NVRTC (#18519)
This PR supports NVRTC as an alternative to NVCC for faster, device-side
JIT compilation of CUDA kernels, in favor of the PR
[https://github.com/apache/tvm-ffi/pull/283](https://github.com/apache/tvm-ffi/pull/283).

It enhances the CUDA compilation backend by:
- Adding Python NVRTC support using cuda-python bindings
- Removing legacy C++ NVRTC fallback in favor of a Python-first approach
- Keeping nvcc as the default compiler with fatbin output (no behavior
change for existing users)

Users can choose the compilation backend using an environment variable
`TVM_CUDA_COMPILE_MODE`, choosing from "nvcc" and "nvrtc". For example,

`TVM_CUDA_COMPILE_MODE=nvrtc python3 your_program.py`

Here is a short benchmark of the compilation speed of kernels in
`test_target_codegen_cuda.py`.

### NVCC vs NVRTC Compilation Time Comparison (Python-side Call)

| Test Case | Code Size | NVCC Time (ms) | NVRTC Time (ms) | Speedup |
| :--- | :--- | :--- | :--- | :--- |
| `test_crossthread_reduction1` | 1945 B | 241.27 | 51.23 | **4.7x** |
| `test_cuda_bf16_vectorize_add` | 3760 B | 342.72 | 44.50 | **7.7x** |
| `test_cuda_const_float_to_half` | 12394 B | 272.85 | 31.99 | **8.5x**
|
| `test_cuda_device_func_call` | 975 B | 215.58 | 21.47 | **10.0x** |
| `test_cuda_float_const_hex_format` | 685 B | 217.39 | 20.52 |
**10.6x** |
| `test_cuda_floordiv_with_vectorization` | 1050 B | 213.88 | 23.32 |
**9.2x** |
| `test_cuda_inf_nan` | 673 B | 214.33 | 24.94 | **8.6x** |
| `test_cuda_tensormap` | 755 B | 213.91 | 20.74 | **10.3x** |
| `test_cuda_thread_sync_inside_condition` | 1007 B | 213.43 | 28.29 |
**7.5x** |
| `test_cuda_vectorize_add` | 908 B | 226.81 | 40.39 | **5.6x** |
| `test_cuda_vectorize_load` | 734 B | 217.25 | 24.02 | **9.0x** |
| `test_device_host_call_same_func` | 924 B | 216.03 | 21.21 | **10.2x**
|
| `test_vectorized_intrin1` | 847 B | 226.15 | 26.34 | **8.6x** |

### NVSHMEM Support

Currently, NVSHMEM is **not** supported via NVRTC.
- Fallback Behavior: When NVSHMEM is required, the compilation pipeline
will automatically fall back to NVCC, even if `TVM_CUDA_COMPILE_MODE` is
set to nvrtc.
- Future Roadmap: Support for NVRTC with NVSHMEM is planned for
follow-up PRs.
2026-01-08 11:08:06 -05:00
wrongtest 13ea9dc104 [TIR] Add step attribute to ForNode (Initial codes) (#18421)
An initial change to add `ForNode::step`.

- Add `Optional<PrimExpr>` typed step attribute to ForNode. Then add
minimal codes for
    - Roundtrip support for TIR tvmscript grammar
    - Correctness of TIR lowering pipeline:
        - Canonicalize the loop in default pipeline
- Ensure the original `ForNode::step` is not dropped by mutations on
`ForNode`.
    - CodeGen support for non-zero min and non-trivial step.

- TODOs in the future (hopefully)
- For **all transformations and analysis tools**, make adaptions to
non-consecutive loop iteration indices
    - Correctness of TensorIR schedule and MetaSchedule

---------

Co-authored-by: baoxinqi <bao.xinqi@intellif.com>
2025-11-24 08:30:16 -05:00
Lei Wang 5ee38eae80 [TIR][CUDA] Preserve float precision in codegen with hexfloat output (#18320)
Previously, `float` constants in codegen were always emitted in **scientific decimal format**, e.g.:

```cpp
bfloat16_t(3.487723e-05f);
```

This could introduce slight **rounding differences** compared to the actual binary representation, since the constant is printed and then re-parsed in decimal. we now emit the value in **hexadecimal floating-point format** (`std::hexfloat`) to preserve the exact binary value, and additionally include the decimal form as a comment for readability:

```cpp
bfloat16_t(0x1.2492492492492p-15f /*3.487723e-05*/)
```
2025-09-19 09:49:01 -04:00
Tianqi Chen 3c36ce2ec6 [FFI][REFACTOR][ABI] Rename NDArray to Tensor (#18275)
This PR Updates the NDArray => Tensor.

Both tensor and ndarray are commonly used terms.

Because the term Tensor is getting more common in the context of ML,
we do the rename to stay more aligned with torch.Tensor and DLTensor.
2025-09-06 14:33:59 -07:00
Tianqi Chen b8eb80b968 [FFI] Formalize ffi.Module (#18213)
This PR formalizes original runtime::Module into ffi
as ffi.Module and cleans the APIs around it.

The goal is to stablize the Module API as extra API that can benefit the overall
ffi interactions. We also refactors the c++ code that depends on the Module.
2025-08-17 23:33:05 +08:00
Ruihang Lai 532db3392a [TIR] Fix host/device function check for build (#18199)
This PR fixes a bug of deciding whether a function is host
or device function in TIR build.

Previously the decision is made based on checking whether `"cpu"`
is a substring of the target string. This check fails to work
for ROCm target, which usually comes with an `"mcpu"` attribute
that also contains `"cpu"`.

This PR fixes by checking target kind. Targets with kind `"llvm"`
or `"c"` will be treated as host functions.
2025-08-08 14:54:06 -04:00
Siyuan Feng ea4369c221 [TIR] Add T.thread_return() for early thread exit in CUDA kernels (#18134)
This commit implements T.thread_return() functionality that allows threads
to exit early from CUDA kernels. The feature is useful for cases where
threads need to conditionally return based on thread indices or other
conditions.

Key changes:
- Add thread_return builtin in TIR
- Implement CUDA codegen for thread_return
- Add Python bindings for T.thread_return()
- Update TIR IR builder to support thread_return
- Add tests demonstrating thread_return usage

Example usage:
```python
@T.prim_func
def main(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")):
    for i in T.thread_binding(16, thread="blockIdx.x"):
        for j in T.thread_binding(32, thread="threadIdx.x"):
            if j >= 16:
                T.thread_return()  # Early exit for threads with j >= 16
            B[i, j] = A[i, j]
```

and generate code is:

```cuda
extern "C" __global__ void __launch_bounds__(32) main_kernel(float* __restrict__ A, float* __restrict__ B) {
  if (16 <= ((int)threadIdx.x)) {
    return;
  }
  B[((((int)blockIdx.x) * 16) + ((int)threadIdx.x))] = A[((((int)blockIdx.x) * 16) + ((int)threadIdx.x))];
}
```
2025-07-14 09:13:19 -04:00
Siyuan Feng acdc164d1f [Target] Support CUDA device function calls (#18055)
[TIR][Target] Support device call compilation

This PR introduces support for device call compilation in TVM by enhancing the BindTarget pass to properly handle functions called from both host and device contexts. The key improvement is the ability to automatically create host-specific duplicates of functions that are called from both host and device code, ensuring proper target binding for heterogeneous compilation.

- **Function Classification**: Analyzes call patterns to identify functions called from host vs device contexts
- **Smart Target Binding**: Automatically binds appropriate targets based on calling context:
  - Functions called only from host → host target
  - Functions called only from device → device target
  - Functions called from both → device target + host duplicate
- **Call Site Updates**: Updates call sites in externally exposed functions to use appropriate duplicates

- Improved device function extraction and kernel generation
- Better handling of error propagation for different device types
- Enhanced buffer declaration and parameter management

- Support for `__device__` function calls in CUDA kernels
- Proper function signature generation for device functions
- Enhanced calling convention handling

- Updated build pipeline to handle device call compilation
- Improved target-specific compilation logic

The following example demonstrates how the BindTarget pass handles functions called from both host and device contexts:

```python
@I.ir_module
class Module:
    @T.prim_func(private=True)
    def add(a: T.int32, b: T.int32) -> T.int32:
        return a + b

    @T.prim_func
    def main(
        A: T.Buffer((128, 128), "int32"),
        B: T.Buffer((128, 128), "int32"),
        C: T.Buffer((128, 128), "int32"),
    ):
        T.func_attr({"global_symbol": "main"})
        length: T.int32 = Module.add(64, 64)  # Call from host
        for bx in T.thread_binding(length, "blockIdx.x"):
            for tx in T.thread_binding(length, "threadIdx.x"):
                C[bx, tx] = Module.add(A[bx, tx], B[bx, tx])  # Call from device
```

After applying `BindTarget(cuda, host="llvm")`, the pass automatically:
1. Creates a device version of `add` with CUDA target
2. Creates a host duplicate `add_host` with LLVM target
3. Updates the main function to call `add_host` from host context and `add` from device context

This enables seamless compilation of mixed host/device code while maintaining proper target-specific optimizations and code generation.

- **Automatic Target Binding**: No manual target annotation required for most use cases
- **Heterogeneous Compilation**: Proper support for functions called from multiple contexts
- **Code Reuse**: Shared functions can be called from both host and device without duplication
- **Performance**: Maintains target-specific optimizations for each context
- **Developer Experience**: Simplifies writing mixed host/device code

The implementation is backward compatible and integrates seamlessly with existing TVM compilation pipelines.
2025-07-11 11:53:13 +08:00
Bohan Hou b6db2ec89f [Runtime] CutensorMap support (#18097)
This PR introduces Cutensor map support in the runtime module. It enables calling kernels whose arguments are cuTensorMap, these arguments are passed as handle(address) and associated with arg_extra_tags that indicate indicate it is tensor map. The TensorMap is allocated on stack with a runtime API
2025-06-30 07:59:52 -04:00
Tianqi Chen 95d1268982 [REFACTOR] Introduce and modernize FFI system (#17920)
This PR modernizes the FFI foundation of the project and introduce
a new minimal and lightweight module [tvm ffi](https://github.com/apache/tvm/tree/refactor-s3/ffi)
based on our lessons in the past few years. It implements a modern
version of the [Unified Packed and Object RFC](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0097-unify-packed-and-object.md)
that unifies the packed function call and object systems.

Summary of the change:
- A dedicated clean Any/AnyView that can store strong and weak
references of items
- Function(previously PackedFunc) system built on top of the Any/AnyView
- A minimal C API that backs the overall calls. We are stabilizing the
API with a goal to bring clean, stable FFI conventions for both compiled
and registered code
- A rewrite of core python binding and generated code based on the module
- Update existing code and test cases to the new module
- Latest dlpack support
 
The new module brings many benefits thanks to the cleaner design,
to name a few:
- Any can support both POD types(int) and object types.
- Containers (e.g. Array) can now also contain Any value, e.g. now
`Array<int>` is supported, no need for boxed types
- Error handling now upgrades to object-based, allowing cleaner
traceback across languages
- Map now preserves insertion orders
- Path toward isolated stabilize minimum core ABI/API foundation module
- Type traits based design that cleanly defines how values interact
with Any system
- Automatic conversion of different types based on traits if needed 

Because FFI upgrade is at heart of the project, the change touches every
component of the system. Importantly, this is an upgrade of the ABI so the
change is not backward compatible.  The code compiled under the old
FFI won't work under the new one. We did provide example ABI translation
(e.g. LegacyTVMArgValueToFFIAny) functions for compatibility. 
The PR tries to leave files in their old places while creating redirections.
The goal is to have the first milestone landed and infrastructure in place,
so we can do further refactors to complete features and cleanup legacy code
as trackable PRs. As of now, python binding and compiled code are under the
new convention while RPC and some  other bindings still relies on legacy ABI
translation. We will work on upgrades in the coming PRs, including areas such
as reflection, phasing out legacy redirections etc.
2025-05-06 19:18:33 -04:00
Siyuan Feng be8e43814a [Refactor] Migrate build API to tvm.compile (#17718)
* tvm.build -> tvm.compile

* relax.build -> tvm.compile

* update
2025-03-09 07:23:52 -04:00
Bohan Hou fc6775e770 [REFACTOR] move build flow from C++ to Python (#17665)
This PR moves build flow from C++ to python, enables more developer productivity and readabilities
2025-02-20 14:47:06 -05:00
Tianqi Chen 9f846bda5b [REFACTOR] Phase out te.schedule python components (#17658)
* [REFACTOR] Phase out te.schedule python components

This PR phases out te.schedule python components.
te.compute is kept around for future usages.
tir.Schedule is a more modern version of the scheduling that we can use onwards.

Doing so also helps us to cleanup the testcases that relies on
explicit full build and execution. As we move future unit testcases
towards structural equality based unit tests.

* Simplify CI to focus on UT

The main rationale is that we should only have very few target
dependent UT in tests/python/codegen and possible
a new category in future for op-level integration if needed.

* Re-enable wasm

* fix lint

* remove hybrid,sparse autodoc and remove tests

---------

Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
2025-02-16 17:13:10 +08:00
Tianqi Chen ccaa534b2c [REFACTOR] Phase out relay python components (#17656)
This PR starts the step 0 to phase out relay from the current
development main branch.  This PR focuses on the python
components of relay, autotvm, auto_scheduler. To make the change
manageable, we will also do followup steps on te.Schedule and
c++ components in followup PRs.

To continue support community members who depends on
legacy flows, the [v0.19.0](https://github.com/apache/tvm/tree/v0.19.0)
branch will continue contain these components.


As noted in [discussion on phasing out legacy components](https://discuss.tvm.apache.org/t/phasing-out-legacy-components/17703/30),
this would help us to do two purposes:

- By removing outdated or redundant elements, we can significantly
reduce complexity and improve maintainability.
- Unify our focus: Concentrating our efforts on the new unity flow
will allow for more efficient development and innovation.

It is also a good opportunity for us to revisit and reduce CI time.
The past relay legacy flow contains a lot of end to end tests that
requires hardware resources to run and causing long CI time.
Moving onwards, we can focus more on unit-tests that focuses
on structural equality and runs within seconds, while be mindful
about tests that requires hardware resources (by restricting them
to specific folders and CI nightly in some cases).

---

Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
2025-02-15 13:48:28 -05:00
Eric Lunderberg 02f48828e4 [FFI] Re-introduce the boxed primitive values (#17257)
* Revert "Revert "[FFI][RUNTIME] Introduce runtime boxed types for int/float/bool" (#17252)"

This reverts commit 11be832620.

* [FFI] Re-introduce the boxed primitive values

Initially introduced in https://github.com/apache/tvm/pull/16183,
these changes were reverted in
https://github.com/apache/tvm/pull/17252 due to performance
degredation in some Relax models.  This could occur when a model
contained a large number of calls to `"vm.builtin.tuple_getitem"`,
which may occur when model weights are provided as a tuple.

This PR re-applies the changes from
https://github.com/apache/tvm/pull/16183, but with the performance
degredation resolved.  The root cause was unnecessary type-checking
when converting from an untyped `tvm::ArrayNode*` to the typed
`tvm::Array<T>`, in the case where `T` is `ObjectRef`.

* Correct typo from T to U
2024-08-12 08:36:17 -04:00
Tianqi Chen 11be832620 Revert "[FFI][RUNTIME] Introduce runtime boxed types for int/float/bool" (#17252)
Revert "[FFI][RUNTIME] Introduce runtime boxed types for int/float/bool (#16183)"

This reverts commit 5f22be4d83.
2024-08-07 12:19:13 -04:00
Eric Lunderberg 5f22be4d83 [FFI][RUNTIME] Introduce runtime boxed types for int/float/bool (#16183)
* [Container] Support non-nullable types in Array::Map

Prior to this commit, the `Array::Map` member function could only be
applied to nullable object types.  This was due to the internal use of
`U()` as the default value for initializing the output `ArrayNode`, where
`U` is the return type of the mapping function.  This default
constructor is only available for nullable types, and would result in
a compile-time failure for non-nullable types.

This commit replaces `U()` with `ObjectRef()` in `Array::Map`,
removing this limitation.  Since all items in the output array are
overwritten before returning to the calling scope, initializing the
output array with `ObjectRef()` does not violate type safety.

* [FFI] Separate runtime types from IR types for int/float/bool

Prior to this commit, `int`, `float`, and `bool` arguments from Python
were converted to `IntImm`, `FloatImm`, and `Bool`.  These are
subtypes of `PrimExpr`, and should only be used at compile-time.  By
automatically applying this conversion as part of the FFI, these types
are required to be present whenever a primitive is converted to a
`tvm::ObjectRef`.

This can become especially fragile for an end-user when storing
objects into a TVM container.  Because TVM containers require all
contents to be `ObjectRef` subclasses, an automatic conversion may be
applied on storing into a container, resulting in an unexpected type
being retrieved from the container.  For example, this currently
occurs in Relax when extracting a `R.Prim` from a `R.Tuple`.

This commit introduces a `Box<T>` type for storage of boxed primitives
at runtime, distinct from the IR types.

* Primitive arguments provided to a PackedFunc that requires an
  `ObjectRef` will be converted to the corresponding boxed type.
  (e.g. Passing a Python `int` to a C++ function accepting `ObjectRef`
  produces a `Box<int64_t>`.

* Boxed primitives provided to a PackedFunc that requires an unboxed
  primitive will be converted to the corresponding primitive.

* PackedFunc return values of `ObjectRef` are converted to the
  corresponding primitive, if present.  (e.g. If a `tuple_getitem`
  with static return type `ObjectRef` returns a `Box<int64_t>`, it
  will be unwrapped to a python `int`.)

Together, these three rules provide backwards compatibility for
existing PackedFunc definitions, while avoiding exposing the user to
any container-induced type conversions betweeen primitive types and
`ObjectRef`.

* Fix unit test failure after merge

* Fix breakage in new unit test
2024-08-05 09:19:20 -04:00
Wuwei Lin a64d1f1cc3 [TIR] Make T.reinterpret nop when dtype is the same (#16879)
* [TIR] Make T.reinterpret nop when dtype is the same

* fix scalable vec handling
2024-04-14 11:21:30 -04:00
Wuwei Lin 109804cc6a [Codegen] Add check to disable invalid reinterpret (#16786)
* [Codegen] Add check to disable invalid reinterpret
2024-03-29 13:58:23 -04:00
Siyuan Feng 268d15c987 [CI] Fix CI Script and Broken Tests (#16521)
* [CI] Fix CI Script and Broken Tests

Co-authored-by: Shengjie Liu <Shengjie.Liu@armchina.com>

* Enhance IterMapSimplify to support uncommon predicate

* Fix runtime traced_callpacked

* Fix derived object attribute get

* update debug line info testcase

* fix relay/relax import and debug_info

* fix lint

---------

Co-authored-by: Shengjie Liu <Shengjie.Liu@armchina.com>
Co-authored-by: tqchen <tianqi.tchen@gmail.com>
2024-02-07 15:54:35 -05:00
Bohan Hou 5308739741 [TIR] Allow sync threads inside condition (#16345)
Originally, it is not allowed to sync threads inside a condition `while, if`.

This PR introduces `tvm_thread_invariant` op to annotate the condition to be thread id invariant and get around the check.
2024-01-04 09:53:25 -08:00
Siyuan Feng bd67d2e5eb [CI] Refactor unittest folder (#16110)
The current unittest folder is too large and contains too many files and
too many components. This PR refactors the unittest folder by moving the
files to the corresponding folders.
2023-11-15 08:23:38 -05:00