## Motivation and context
Copy regions can retain constant-valued expressions after variable
substitution. Converting those expressions directly to Python integers
rejects valid slices even though their extents are statically known.
## Changes
- Simplify swizzled-region extents before converting them to Python
integers.
- Continue rejecting genuinely symbolic extents.
- Add a hardware-independent regression to the existing layout test
suite.
## Testing
- `python -m pytest tests/python/tirx/test_layout.py`
- Changed-files pre-commit checks
## Motivation and context
`tcgen05.ld/st` with a `.16x*b` atom accesses one 16-row half-slab from
each 32-row TMEM partition owned by a warp. Across a four-warp
warpgroup, the two physical half-slabs are:
| View | Physical TMEM lanes | PTX row immediate |
|---|---|---|
| lower half | `0..15, 32..47, 64..79, 96..111` | `0` |
| upper half | `16..31, 48..63, 80..95, 112..127` | `16` |
Layout D for an M=128 accumulator occupies both halves. Layout F exposes
an M=64 logical tile over one half, which is useful both for a native
M=64 accumulator and for reading either 64-row half of an existing
Layout D accumulator.
Before this PR, `tmem_datapath_layout("F", 64, cols)` could describe
only the lower half. The copy dispatcher classified a TMEM buffer only
as `"D"` or `"F"` and every M=64 `.16x*b` operation started at `row=0`.
As a result, there was no layout-preserving way to create a recognized
64-row view of the upper half of a Layout D accumulator: a normal Layout
F view still addressed the lower half, while a hand-written `+16@TLane`
layout was not recognized by the dispatcher.
The half-slab selection belongs in the buffer layout because layout is
the source of truth for physical placement in TIRx. It should not be an
out-of-band `copy_async` option. This PR therefore records the selection
in Layout F, carries it through datapath classification, and derives the
PTX row immediate from it.
A related invariant is that a default M=128 TMEM allocation must be
structurally identical to named Layout D. The dispatcher recognizes
datapaths structurally, so keeping a separate hand-written default
layout creates an unnecessary drift risk. This PR makes the default call
the public Layout D factory directly.
Finally, the register-side `tcgen05_atom_layout` must agree with the PTX
mapping from a logical `(row, col)` to `(laneid, wid_in_wg, register)`.
A self-consistent load/store round trip is not enough to prove that
mapping: raw PTX can move the same bits back even when the logical
layout label is wrong. Elementwise dispatch does consume that label, so
this PR adds direct mapping and elementwise compilation fences for all
`.16x*b` atom families. These are coverage additions; the production
atom-layout construction itself is unchanged here.
## Changes
- Add `sub_slab={0,1}` to `tmem_datapath_layout("F", ...)`.
- Encode the upper view as a `+16@TLane` offset and reject invalid
selectors, including nonzero selectors for Layout D.
- Classify TMEM layouts as `(datapath, sub_slab)` and emit `.16x*b` with
`row=(sub_slab + slab) * 16`.
- Preserve the existing M=128 behavior: Layout D with a 128-row `.16x*b`
fragment still emits two operations at rows `0` and `16`.
- Build the default M=128 TMEM layout through `tmem_datapath_layout("D",
...)`.
- Document the physical lane mapping and supported datapath/atom
combinations.
- Add direct atom-layout mapping coverage and a warpgroup elementwise
regression.
## Testing
- Static Layout F checks cover every logical row for both `sub_slab=0`
and `sub_slab=1`.
- B200 readback tests populate one Layout D accumulator, then verify
that lower and upper Layout F views reproduce the two corresponding
register halves for `.16x64b`, `.16x128b`, and `.16x256b`.
- Negative tests cover invalid sub-slab values and incompatible
datapath/atom pairings.
- Direct `(row, col) -> (laneid, wid_in_wg, register)` sweeps cover
supported `.16x*b` shapes and repetitions.
- Warpgroup elementwise codegen verifies that an atom-layout fragment
canonicalizes and slices correctly.
- Changed-files pre-commit checks.
## Motivation and context
ConvertSSA caches remapped buffers while an SSA-renamed variable is in
scope. Cleanup previously removed a cached buffer only when the renamed
variable was its data pointer, so remaps through other buffer fields
could survive after scope exit.
## Changes
- Track dependencies in every buffer field rewritten by
`GetRemappedBuffer`, including shape, strides, element offset, and tile
layout fields.
- Invalidate cached remaps consistently when an SSA-renamed variable
leaves scope.
- Add a regression using reused sibling loop variables and a
variable-dependent element offset.
## Testing
- `python -m pytest
tests/python/tirx-transform/test_tir_transform_convert_ssa.py`
- Changed-files pre-commit checks
## Motivation and context
A TIRx pointer carries two pieces of information that later lowering
needs: the pointee element type and the storage scope. Both must survive
when a pointer-producing expression is assigned to a Python name and
then used as the backing storage of a buffer.
A concrete example is accessing an mbarrier in another CTA through
distributed shared memory:
```python
ptr_ty = PointerType(PrimType("uint64"), "shared")
remote_ptr = T.reinterpret(
ptr_ty,
T.ptx.map_shared_rank(mbar.ptr_to([0]), T.int32(0)),
)
remote_mbar = T.decl_buffer(
[1], "uint64", data=remote_ptr, scope="shared"
)
```
`map_shared_rank` returns the raw `uint64` address produced by PTX
`mapa`, and `reinterpret` gives that address the intended
`PointerType(uint64, shared)`. Because `decl_buffer(data=...)` requires
a pointer `Var`, assigning the expression to `remote_ptr` should create
an immutable typed pointer binding.
Before this PR, an unannotated assignment such as `remote_ptr = <pointer
expression>` followed the same parser path as a numeric assignment. That
path allocates a mutable local scalar and therefore cannot represent a
`PointerType`. The pointer expression could not be carried as a
correctly typed `Var` into `decl_buffer` and CUDA lowering.
This PR makes an unannotated pointer-valued assignment emit a TIRx
`Bind`. The bound `Var` has exactly the type of the right-hand side,
including its element type and storage scope. Pointer bindings are
immutable, so reassignment in the same scope is diagnosed; shadowing a
name supplied through `extra_vars` remains valid. Numeric assignments
keep their existing mutable-local behavior.
## Type propagation fixes
The parser fix exposed several other boundaries where pointer type
information must remain consistent:
| Boundary | Previous behavior | Behavior after this PR |
| --- | --- | --- |
| Unannotated pointer assignment | Tried to materialize the value as a
local scalar | Emits an immutable `Bind` with the RHS `PointerType` |
| `address_of(buffer)` / `buffer.ptr_to(...)` | Reused the raw backing
pointer type | Returns a pointer to `buffer.dtype` while preserving the
backing pointer storage scope |
| `tvm_access_ptr` / `ptr_byte_offset` | Accepted strings or annotation
expressions, but not a `PrimType` object directly | Accepts `PrimType`
and produces the corresponding typed pointer |
| Printed `T.ptx.mapa` call | The printer emits all intrinsic attributes
positionally, but the Python helper required keyword-only arguments |
Accepts the canonical printed form so pointer code round-trips through
TVMScript |
The `address_of` distinction matters for typed views over byte-addressed
storage. For example, if a `float32` buffer is backed by a `uint8*`
allocation in `shared.dyn`, the address of a buffer element must be
`PointerType(float32, shared.dyn)`, not `PointerType(uint8,
shared.dyn)`.
With these changes, the DSMEM example above round-trips through
TVMScript and CUDA codegen declares the remote buffer pointer as
`uint64_t*`.
## TMA dtype normalization
This PR also contains a small, separate type-representation fix in TMA
descriptor construction. `TmaPlan.elem_dtype` is a string consumed by
the host-side `runtime.cuTensorMapEncodeTiled` packed call, but
`_assemble_plan` stored `g_buf.dtype`, which is a `PrimType`. Converting
it with `str(g_buf.dtype)` ensures that the generated packed-call
argument is `StringImm("float16")` rather than an IR type object. This
does not change the pointer-binding semantics described above.
## Testing
- Verify that an unannotated pointer expression creates a `Bind` whose
`Var` type matches the RHS type.
- Verify that pointer reassignment is rejected while shadowing an
`extra_vars` name is allowed.
- Verify parser/printer structural round-tripping for the pointer
binding and canonical `T.ptx.mapa` call.
- Verify that `address_of` uses the logical buffer element type and
preserves the storage scope for byte-backed buffer views.
- Verify that `tvm_access_ptr` and `ptr_byte_offset` accept `PrimType`
inputs.
- Compile the DSMEM `map_shared_rank` example through the CUDA TIRx
pipeline and check for a typed `uint64_t*` remote buffer pointer.
- Verify that the TMA host initialization passes the descriptor dtype as
a `StringImm`.
## Motivation and context
`TilePrimitiveCallNode` participates in FFI reflection but did not
provide the unsafe-init constructor required to reconstruct an object
during deserialization.
## Changes
- Add the reflection unsafe-init constructor to `TilePrimitiveCallNode`.
- Add a pickle round-trip regression that preserves arguments,
workspace, config, dispatch, and execution scope.
## Testing
- `python -m pytest tests/python/tirx/test_op.py`
- Changed-files pre-commit checks
## Motivation and context
The TVMScript parser already tracks Python AST locations for
diagnostics, but TIRx statements and expression results emitted through
`IRBuilder` did not retain those locations. After parsing, a direct
intrinsic call, an inlined helper body, or a `TilePrimitiveCall`
therefore could not be traced back to the source range that produced it.
Inline expansion also needs more than a single flat location. The
generated IR should retain both the caller location and the
helper-definition location, while ordinary nested AST evaluation within
one source should not accumulate redundant enclosing spans.
## Changes
- Add an active source-span stack to `IRBuilder`, with scoped push/pop
support.
- Make the parser activate the current AST source range while visiting
statements and evaluating expressions.
- Attach the active span to emitted TIRx statements and to expression
results that do not already carry an explicit span.
- Normalize nested spans from the same source to the innermost relevant
range.
- Preserve cross-source inline expansion history as a `SequentialSpan`,
ordered from the call site to the expanded definition.
- Reuse the same source-coordinate calculation for diagnostics and IR
spans so their line and column conventions remain consistent.
Source spans remain diagnostic metadata: functions parsed from different
source locations keep the same structural hash and remain structurally
equal.
## Testing
- Verify exact parser source coordinates against diagnostic coordinates.
- Verify spans on direct intrinsic calls and `TilePrimitiveCall` nodes.
- Verify that inline expansion produces a `SequentialSpan` containing
caller and callee ranges.
- Verify direct `IRBuilder.with_source_span` behavior.
- Verify that source spans do not affect structural identity.
- Run changed-files pre-commit checks, including clang-format.
Focused result: 9 tests passed.
## Motivation and context
`MBarrier.remote_view(rank)` represents an mbarrier owned by another CTA
in the same cluster. The existing view kept only a buffer whose pointer
had been mapped to the remote CTA with PTX `mapa`. Calling
`remote_bar.arrive(...)` then followed the inherited local-arrive path
and emitted the local `mbarrier.arrive.shared.b64` form against that
mapped remote address.
A remote arrival must instead use the `shared::cluster` instruction
form. TIRx models that form with the owner-local barrier pointer plus
the destination CTA rank and predicate. Using the local instruction with
a remote address is not equivalent and is reported by synccheck as a
local arrival on a remote mbarrier address.
## Changes
- Keep the owner-local buffer and target CTA rank when constructing a
remote mbarrier view.
- Route `MBarrier` remote arrivals through the cluster helper using the
local barrier pointer and stored CTA rank.
- Apply the same routing to `TMABar`, including
`mbarrier.arrive.expect_tx`.
- Keep a typed `PointerType(uint64, shared)` mapped buffer on the view
so `ptr_to` remains available to operations that explicitly consume a
remote shared-memory pointer.
- Reject operations with ambiguous or invalid ownership semantics:
- initializing or waiting on a remote view,
- supplying another `cta_id` to a view that already fixes its target,
- creating a remote view from another remote view.
- Preserve the existing local-CTA behavior for ordinary barriers.
## Testing
- Verify the typed `mapa` binding and remote buffer in TIRx IR.
- Verify CUDA codegen for plain and counted
`mbarrier.arrive.shared::cluster.b64` forms.
- Verify CUDA codegen for remote
`mbarrier.arrive.expect_tx.shared::cluster.b64`.
- Verify that the corresponding local instruction forms are not emitted
for remote views.
- Verify diagnostics for remote init, wait, nested views, and
conflicting `cta_id` arguments.
- Run changed-files pre-commit checks.
Focused result: 3 tests passed.
This PR fixes invalid pointer arithmetic emitted by C-family codegen for
vector-typed `tvm_access_ptr`.
A vector access pointer is lowered to `address_of(BufferLoad(...))` with
a `Ramp` index describing its lane indices. For example, `Ramp(4, 1, 2)`
represents scalar elements `[4, 5]`, so its address should be the
address of the first lane, `&A[4]`.
LLVM codegen already extracts `Ramp::base` in this case. However,
`CodeGenC` previously passed the complete ramp to pointer arithmetic,
which could generate invalid CUDA code such as:
```cpp
(float*)A + make_int2(4, 5)
```
This PR makes `CodeGenC` use `Ramp::base` when generating the address of
a vector `BufferLoad`. The normalized index is applied to both the
direct pointer-offset path and the general `GetBufferRef` path.
The existing scalar-buffer plus `Ramp` lowering is preserved. This
avoids regressions for padded vector types such as `float32x3` and
packed vector types such as `int4x4`, while making C-family codegen
consistent with LLVM codegen.
Regression tests cover:
- `float32x2` C codegen.
- Padded `float32x3` LLVM codegen.
- Packed `int4x4` CUDA codegen.
Following the recent PrimType refactor, `nn.Tensor.dtype` returns a
`PrimType` object instead of the documented string dtype value. This
breaks consumers such as NumPy's astype. This PR returns the underlying
dtype value and adds an assertion verifying that `Tensor.dtype` remains
string-compatible.
## Summary
This adds support for the `start` and `end` attributes introduced for
ONNX `Shape` in opset 15.
The Relax ONNX frontend previously reused the opset 13 implementation,
which always returned the full input shape. As a result, models using
sliced shape values could construct an incorrect target shape and fail
in downstream operators such as `Reshape`:
```text
ValueError: Reshape expects the new shape to be convertible from the old shape. However, the old shape is R.shape([12]), with product T.int64(12), while the new shape is R.shape([2, 3, 4]), with product T.int64(24)
```
### Minimal reproduce
```python
import onnx
from tvm.relax.frontend.onnx import from_onnx
input_shape = [2, 3, 4]
data_shape = [12]
expected_shape = [3, 4]
start = 1
end = None
opset = 15
shape_attrs = {"start": start}
if end is not None:
shape_attrs["end"] = end
model = onnx.helper.make_model(
onnx.helper.make_graph(
[
onnx.helper.make_node("Shape", ["x"], ["shape"], **shape_attrs),
onnx.helper.make_node("Reshape", ["data", "shape"], ["y"]),
],
"shape_start_end_repro",
[
onnx.helper.make_tensor_value_info(
"x", onnx.TensorProto.FLOAT, input_shape
),
onnx.helper.make_tensor_value_info(
"data", onnx.TensorProto.FLOAT, data_shape
),
],
[
onnx.helper.make_tensor_value_info(
"y", onnx.TensorProto.FLOAT, expected_shape
)
],
),
opset_imports=[onnx.helper.make_opsetid("", opset)],
)
print(f"Shape attributes: start={start}, end={end}")
print(f"Expected Shape output: {input_shape[start:end]}")
print(from_onnx(model, opset=opset).script())
```
The new implementation applies `start` and `end` slicing to static and
symbolic shape expressions. It also handles runtime-produced shape
values by converting them to a tensor, applying `strided_slice`, and
converting the result back to a shape.
This PR excludes `tirx::StmtNode::span` from structural equality and
structural hash calculations.
Source locations are diagnostic metadata and should not affect the
structural identity of a TIRx statement. Other TIRx nodes with spans
already follow this behavior, but `StmtNode::span` was missing the
`SEqHashIgnore` field flag.
A regression test is added to verify that two otherwise identical
statements with different spans are structurally equal and produce the
same structural hash.
- `LegalizeOps` skips grouped `conv1d/2d/3d` when channel size is
symbolic
- The only blocker is `topi.nn.conv`'s divisibility `assert`s, which
fail on symbolic `PrimExpr`; the grouped compute already handles
symbolic dims
## How
- Skip the divisibility check when the channel size is not a constant
int
- Remove the symbolic-channel guards from the conv legalize functions
## Summary
- bump tvm-ffi and include the device definition where its `DLDevice`
traits are instantiated
- keep only the required Tensor wrapper layout fix and register
`ir.Type` before reflected `Expr` fields can materialize a fallback
wrapper
- preserve `BaseFunc.with_attr` callers by moving only method-private
results, never the canonical `self` wrapper
## Rationale
The tvm-ffi lifetime update requires a replacement wrapper to fit the
layout already registered for the same type index. `runtime.Tensor`
replaces the core `ffi.Tensor` wrapper, so it must use empty slots. The
ordinary TVM mixins are first-registered with their concrete descendants
and may safely retain normal Python dictionaries; the additional mixin
and explicit-dictionary slot changes are not required.
Object tying also means `BaseFuncCopy(self)` may return `self`. Passing
that wrapper through `_move()` invalidates the caller. The first update
now passes the alias as an lvalue, forcing native copy-on-write to
create a private result. Only later dictionary updates move a result
that is not `self` and has not escaped the method.
## Validation
- built an exact CPython 3.12 wheel from tvm-ffi `21e30c3b1d` and
rebuilt TVM against it
- direct Type/function/detach regressions: 3 passed
- complete IR plus focused Relax coverage: 111 passed
- prior Relax failure set: 157 passed, 9 skipped
- runtime probe for `relax.Function`, `relax.ExternFunc`, and
`tirx.PrimFunc`: original wrappers preserved; single- and
multi-attribute results distinct and valid
- all touched-file pre-commit hooks passed
---------
Co-authored-by: Yaxing Cai <caiyaxing666@gmail.com>
This PR reduces repeated Relax frontend and integration test work while
preserving distinct coverage. It reworks ONNX ConvTranspose tests into
direct importer checks plus 11 numerical cases covering all ranks,
asymmetric padding, grouping, bias, dilation, and output padding.
Targeted runtime improves from 17.97s to 3.55s.
- removes duplicate ONNX Pow, unused dynamic Squeeze parameterizations,
and irrelevant Resize ROI value permutations.
- consolidates overlapping PyTorch integration tests while preserving
symbolic shapes, TIR, I.pyfunc, and packed-function coverage.
- removes the redundant BasePyModule aggregate suite, moves its unique
output-only call_tir case into the DLPack test, and removes a DLPack
test that swallowed all exceptions.
The ONNX importer's Gather converter asserted that indices must be a
constant whenever the data operand is a ShapeExpr, raising "Only
constant indices supported for shape gather." for any runtime-computed
index. Detection post-processing graphs such as FasterRCNN feed a
dynamic index into a Gather whose data comes from a Shape node, so
import failed before compilation could start.
Keep the fast path for a single constant index, which resolves one
dimension to a PrimValue and preserves shape-specialized handling
downstream. Any other index (dynamic, or a constant selecting multiple
dimensions) materializes the shape as an int64 tensor via
shape_to_tensor and gathers from it at runtime, reusing the existing
negative-index normalization.
Adds a regression test that gathers a dimension out of a Shape result
using a non-constant index, covering positive and negative indices, and
checks it against onnxruntime.
Fixes part of #19965.
## Summary
Before opset 13, ONNX `Squeeze` specifies `axes` as a node attribute
rather than a tensor input. The Relax ONNX importer only implemented
`_impl_v13`, which reads axes from the second input, so for opset < 13
models, the attribute was silently ignored (`axis` defaulted to `None`)
and the importer squeezed every size-1 dimension instead of only the
requested one. This produced tensors with the wrong rank, breaking
downstream ops like `Transpose` whose `perm` no longer matched the
input's actual rank.
Added `_impl_v1` to read `axes` from the node attribute for opset < 13,
and factored the existing squeeze logic into a shared `_squeeze` helper
used by both `_impl_v1` and `_impl_v13`.
## Test plan
- Added `test_squeeze_axes_attribute` to
`tests/python/relax/test_frontend_onnx.py`, covering an opset-11
`Squeeze` node with `axes` as an attribute.
- Ran `pytest tests/python/relax/test_frontend_onnx.py -k squeeze`. All
21 tests pass.
- Verified against the real-world model that triggers this bug,
[PaddlePaddle/PP-OCRv6_tiny_rec_onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_tiny_rec_onnx)
(opset 11, uses attribute-based `Squeeze`): import fails on `main` with
`Transpose: number of axes in perm attribute (3) must equal the number
of input tensor dimensions (-1)`, and succeeds with this fix.
## Real-world reproduction
```python
import urllib.request
import onnx
from tvm.relax.frontend.onnx import from_onnx
# PaddlePaddle/PP-OCRv6_tiny_rec_onnx (opset 11, uses attribute-based Squeeze)
url = "https://huggingface.co/PaddlePaddle/PP-OCRv6_tiny_rec_onnx/resolve/main/inference.onnx"
path = "pp_ocrv6_tiny_rec.onnx"
urllib.request.urlretrieve(url, path)
model = onnx.load(path)
print("opset:", [(o.domain, o.version) for o in model.opset_import])
for node in model.graph.node:
if node.op_type == "Squeeze":
axes_attr = [a for a in node.attribute if a.name == "axes"]
print(node.name, "inputs=", list(node.input), "axes_attr=", axes_attr)
# Fails on main with:
# ValueError: Transpose: number of axes in perm attribute (3) must equal the number of input tensor dimensions (-1)
# Succeeds with this fix.
mod = from_onnx(model)
print("Import succeeded")
```
Fixes (partially) #19965. The shape-Gather and dynamic-TopK issues
reported in that issue are separate and not addressed here.
Followup to #19978: the const-int-bound modular-set fix correctly
prevents the simplifier from over-folding the adaptive pool window
extent. The previous expected IR used the simplified closed form `(v_ax2
% 3 * 4 + 16) // 12 + 1`, which was only reachable because the buggy
bound let `CanProve` prove an invalid predicate. After the fix the
generated IR retains the correct `T.Select` form, so update the expected
IR to match and remove the `xfail` marker that was added in #19978.
This branch is based on current main so the `xfail` removal is explicit
(addressing feedback from @tlopex on the previous attempt in #19995).
This PR:
- Removes duplicate module/functional, alias, positional-argument, and
no-op cases from the PyTorch ExportedProgram tests.
- Consolidates the four GRU configurations into a table-driven loop
without removing any configurations.
- Removes duplicated FX cases already covered through the same shared
converters.
- Restores FX scalar tensor constant coverage and TFLite
constant-parameter Gather and static broadcast/MUL coverage.
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.
Rename the reflected local `Var` field from `name_hint` to `name` and
update its typed C++ consumers. Preserve distinct named-node APIs and
the Python constructor keyword compatibility path, while making `.name`
the sole stored Var property. Upgrade legacy compact JSON records for
current and pre-unification Var schemas.
Validation: full runtime/compiler build, focused C++ Var copy-helper
test, focused Python IR/Relax/TIRx/script tests, Vulkan codegen syntax
build, touched-file pre-commit checks, and `git diff --check`.
Fixes#19887.
This PR fixes several Relax TensorRT BYOC issues exposed by YOLO-style
models:
- adds TensorRT support for SiLU and resize2d
- preserves operand and TupleGetItem ordering during codegen
- fixes cyclic and unsafe Tuple/TGI region merging
- handles static Shape bindings and nested packed-function outputs
- normalizes PrimType dtype arguments passed to relax.arange
With these changes, yolo11n-seg can be merged into a single TensorRT
region, while yolo11n can be imported and partitioned successfully.
This PR removes redundant and misleading Relax ONNX and TFLite frontend
tests.
For ONNX, it removes numerical tests already covered more systematically
by the official ONNX backend suite, duplicate/subset IR checks,
redundant NMS cases, and unused test parameters/helpers.
For TFLite, it removes tests that TensorFlow 2.19 rewrites into
already-covered operators, exact duplicates, no-op models, and checks
superseded by stronger retained tests.
Cleanup pass that relies on IR invariants instead of re-checking
already-guaranteed conditions. No new features; this is a
consolidation/cleanup pass only.
## Changes
- **docsifier (`python_doc_printer.cc`)**: the `ExprStringDoc` escape
scope always wraps the printer's fixed in-memory `ostringstream` sink,
which never short-writes and never enters a fail state. Drop the
streambuf-general short-write reporting in `xsputn`, the ctor `good()`
ICHECK, the dtor `rdstate`/`setstate` dance, and the redundant
post-render `good()` ICHECK; keep the one-line `saw_newline()` contract.
- **relax diagnostics (`well_formed.cc`, `block_builder.cc`)**: the ty
diagnostics test `ty.IsMissing()` on a now non-nullable `Type`, so word
them as "is missing" rather than "is nullptr".
- **relax numeric-gradient tests**: derive the device from the build
target via `tvm.device_from_target` inside the helpers instead of
threading a redundant `dev` argument that duplicates `target` at every
call site; annotate the numpy inputs as `np.ndarray`.
- **target/printer tests**: drop assertions that re-check a condition an
earlier assertion in the same test already guarantees.
### Summary
Onnx subgraph imports should also normalize and generate ty_info for its
ops, this is broken since #19853 refactor.
### Issue
```
tests/python/relax/test_frontend_onnx.py:11581: in test_if_subgraph
tvm_model = from_onnx(model, keep_params_in_input=True)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/usr/lib/python3.15/site-packages/tvm/relax/frontend/onnx/onnx_frontend.py:6283: in from_onnx
return g.from_onnx(graph, opset)
^^^^^^^^^^^^^^^^^^^^^^^^^
/usr/lib/python3.15/site-packages/tvm/relax/frontend/onnx/onnx_frontend.py:5823: in from_onnx
self._construct_nodes(graph)
/usr/lib/python3.15/site-packages/tvm/relax/frontend/onnx/onnx_frontend.py:5969: in _construct_nodes
then_expr = self._convert_subgraph(self.bb, attr["then_branch"])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/usr/lib/python3.15/site-packages/tvm/relax/frontend/onnx/onnx_frontend.py:6166: in _convert_subgraph
op = self._convert_operator(op_name, inputs, attr, self.opset)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/usr/lib/python3.15/site-packages/tvm/relax/frontend/onnx/onnx_frontend.py:6117: in _convert_operator
sym = op_function(self.bb, inputs, attrs, [self._nodes, self._params])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
/usr/lib/python3.15/site-packages/tvm/relax/frontend/onnx/onnx_frontend.py:1882: in _impl_v11
ndim = len(inputs[0].ty.shape)
^^^^^^^^^^^^^^^^^^
E AttributeError: 'Type' object has no attribute 'shape'
```
### Fix
Add conversion check, normalize and populate the final relax op with
ty_info regardless of the graph context.
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.
ONNX integer Div uses truncating division, rounding toward zero. The
Relax ONNX frontend already special-cased integer Div to detect zero
divisors, but its PrimExpr folding path could still use NumPy
floating-point division when one of the inputs was a shape-derived
PrimExpr.
That behavior can produce floating-point TIR values for integer
shape/index computations. For example, a `Shape -> Gather -> Div ->
Slice` subgraph can produce `T.float64(128.666...)` as a Slice bound,
which Relax rejects because strided_slice expects integer PrimExpr
bounds.
This patch handles scalar integer Div inputs that contain a PrimExpr
using TIR `truncdiv`, preserving ONNX semantics while keeping shape
computations in TIR instead of routing them through NumPy. Constant
tensor Div continues to use the existing generic binary constant-folding
path.
The regression tests cover:
- integer constant folding with negative values to distinguish
truncation from floor division
- a shape-derived PrimExpr Div used as a Slice bound
- integer zero-divisor error handling
Verification:
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_div_integer_constant_zero_divisor_raises_valueerror
tests/python/relax/test_frontend_onnx.py::test_div_integer_constant_folding_truncates_toward_zero
tests/python/relax/test_frontend_onnx.py::test_div_integer_primexpr_folding_truncates_toward_zero
-q`
Fixes#19974
Signed-off-by: viiccwen <vicwen@apache.org>
## 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
This PR removes the suite-wide TIRx compute-capability gate and
localizes hardware skips to CUDA codegen and tile-primitive tests. It
keeps the original test parameterization unchanged, allowing parser,
printer, IR, transform, and other non-hardware TIRx tests to run in
regular CI while device-dependent cases are skipped when SM100 hardware
is unavailable. CUDA codegen helpers use explicit target architectures
where needed, and the run-only benchmark utility tests retain a local
SM100 gate.
This intentionally avoids adding separate compile/run parameter cases. A
follow-up PR can audit slow frontend execution tests and define a
focused runtime regression budget for real TIRx kernels.
Local validation:
- `pytest -n auto -m "not gpu" tests/python/tirx`: 486 passed, 79
skipped.
- `pytest -n auto -m gpu tests/python/tirx` without matching hardware:
1509 skipped.
- Pre-commit passed on all changed files.
This pr switches CI to the Ubuntu 24.04 (noble) images. Bump ci_tag in
ci/jenkins/docker-images.ini to 20260629-192919-24bbfd2e -- the images
built from #19893 (ci_cpu/ci_arm/ci_wasm/ci_gpu on Ubuntu 24.04, whose
default g++ is gcc-13, giving full C++20 support).
Also, this pr re-enables Z3 (AUTO). #19828 temporarily set USE_Z3=OFF
(in CMakeLists.txt and the pyproject wheel build) to dodge a z3-static
build failure. The CI image now ships z3-static (#19835), so this
restores USE_Z3=AUTO: the Z3-backed Analyzer proving is enabled when
z3-static is available and silently skipped otherwise.
While Z3 stayed disabled, PrimExprNode::ty became a method, leaving two
stale field accesses in z3_prover.cc's IsZ3SupportedExpr (only compiled
under TVM_USE_Z3). Fixed expr->ty -> expr->ty().
Verification:
- The Ubuntu 24.04 images (#19893) built successfully for ci_cpu/ci_arm/
ci_wasm/ci_gpu (the GPU image includes ROCm 6.4.4 and the CUDA 24.04
base).
- Re-enabling Z3 was validated with a build-only wheel run: all four
wheels (Linux x86_64/aarch64 manylinux_2_28, macOS arm64, Windows) build
green with z3-static compiled and linked, confirming the earlier
z3-static link concern is resolved on the current toolchain.
`ConstantOfShape` uses its input tensor as shape metadata. When that
input is an initializer and `keep_params_in_input=True`, the Relax ONNX
frontend should recover the initializer value from `params` instead of
treating the input as an opaque runtime value.
This patch applies `get_constant` to the `ConstantOfShape` shape input
before shape handling. It also guards the constant-shape folding path so
it only calls `len(shape)` on `relax.ShapeExpr` values.
The regression test covers an initializer-backed shape input imported
with `keep_params_in_input=True` and checks that the resulting Relax
function has the expected output shape and dtype.
A separate lint follow-up commit removes stale `F821` suppressions from
two DLight files so the repository-wide CI lint is clean.
Verification:
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_constantofshape_initializer_shape_with_keep_params_in_input
-q`
- `pre-commit run --all-files`
Fixes#20001.
---------
Signed-off-by: viiccwen <vicwen@apache.org>
## Summary
Adds support for newer `QuantizeLinear` and `DequantizeLinear`
attributes in the Relax ONNX frontend.
This includes `output_dtype`, `saturate`, and newer opset behavior,
while rejecting unsupported blocked quantization and `precision` cases.
For `QuantizeLinear` and `DequantizeLinear`, opsets 24 and 25 use the
existing converter for currently supported types. Support for
`float8e8m0`, `int2`, and `uint2` are outside this PR’s scope.
## Testing
Added structural and rejection tests for opsets 19, 21, 23, 24, and 25.
This PR fixes an undefined `TX` reference in the DLight GPU GEMV
inner-reduction schedule.
The broadcast epilogue path splits fused epilogue loops and binds the
inner loop to `threadIdx.x`, but it referenced `TX`, which is not
defined in the surrounding scope or passed into the helper. This PR uses
the existing `TR` tile factor, which is already used for the
`threadIdx.x` direction in this schedule.
A regression test is added to cover the GEMV broadcast epilogue path.
Fixes#19969
Tests:
- `python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q`
- `python -m pytest tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
-q`
- `python -m pytest tests/python/s_tir/dlight/test_gpu_reduction.py -q`
Note: The branch is based on `apache/tvm:main`. Local main-branch test
execution on this Windows machine could not be completed because the
available compiled TVM library is from a v0.25 build and does not match
the latest Python sources; the same tests passed in the matching local
v0.25 environment before rebasing the patch to main.
### What changed
Record the original input rank before `Expand` left-pads the input shape
for broadcast validation. The no-op fast path now returns the input
unchanged only when both the padded shape and the original rank match
the target.
A regression test covers expanding `[1]` to `[1, 1]` when the target is
represented as a Relax `ShapeExpr`.
### Why
ONNX `Expand` right-aligns dimensions and may increase tensor rank by
adding leading dimensions. Previously, a rank-expanding broadcast could
look like a no-op after the frontend padded the input shape, causing it
to return the original lower-rank tensor. Downstream operators could
then receive inconsistent ranks.
This fixes the focused bug tracked in #19991 and is part of the
investigation and fixes for #19971. It does not close#19971 because the
attached model exposes additional independent importer issues after this
`Concat` failure is resolved.
Fixes#19991
Part of #19971
### Validation
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_expand -q`
- A/B checked the model attached to #19971: the base revision reproduces
`Concat expects all input tensors to have same ndim`, while this change
advances beyond that `Concat`.
Signed-off-by: viiccwen <vicwen@apache.org>
## Summary
Fix TFLite Relax frontend cast paths that still used removed/nonexistent
cast
APIs.
- Use `relax.op.astype` for FLOAT16 `DEQUANTIZE` constants.
- Use `relax.op.astype` around the existing quantized `AVERAGE_POOL_2D`
converter path.
## Design
This PR only replaces invalid frontend API calls with `relax.op.astype`.
The quantized avgpool regression test calls the converter path directly
because
the top-level TFLite importer still rejects quantized `AVERAGE_POOL_2D`
before
conversion. Enabling that operator globally is out of scope.
## Tests
Added:
- `test_dequantize_float16_uses_astype`
- `test_quantized_avg_pool2d_uses_astype`
Validated with:
```bash
python -m ruff format \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m ruff check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m pytest \
tests/python/relax/test_frontend_tflite.py \
-k "dequantize or avg_pool" -q
```
Result:
```text
ruff format: 2 files left unchanged
ruff check: All checks passed
targeted dequantize/avg_pool tests: 9 passed, 551 deselected
```
I also ran the full TFLite frontend file:
```text
tests/python/relax/test_frontend_tflite.py: 559 passed, 1 failed
```
The remaining failure is unrelated to this PR:
`test_broadcast_to` expects `R.multiply(..., ones)` while the importer
emits
`R.broadcast_to(...)`.
This PR uses `torch.testing.make_tensor` where it provides a clear
testing benefit in the PyTorch exported-program frontend tests.
- Generate boolean masks directly instead of comparing random float
tensors
- Generate parametrized dtypes directly instead of creating integer
tensors and converting them with `.to()`
- Specify the CPU device explicitly
## Summary
Adds ONNX frontend support for `GroupNormalization` by mapping it to the
existing `relax.op.nn.group_norm`.
Supports opset 18 per-group scale/bias expansion, opset 21 per-channel
scale/bias, and `stash_type` cast behavior.
## Testing
Includes structural checks for opset 18, opset 21, rank-3 inputs, and
fp16 `stash_type` paths.
## Why
Fixes#19925. `relax.shape_to_tensor` had no legalization, so it always
lowered to the host packed func `relax.run.shape_to_tensor`, producing a
CPU tensor regardless of the target device.
## How
- Register a legalization that emits the shape values as a `call_tir` TE
kernel, passing symbolic dims via `tir_vars`.
- Fall back to the packed func when the shape values are unknown
(`ShapeStructInfo` without values).
Fixes#19977.
ONNX `BatchNormalization` allows the input/output tensor dtype,
scale/bias dtype, and mean/variance dtype to be separate floating-point
type parameters.
For example, a valid ONNX model may use `float16` data with `float32`
gamma, beta, mean, and variance tensors.
The Relax `batch_norm` operator currently requires all five input
tensors to have the same dtype. The ONNX frontend previously forwarded
the ONNX inputs directly to `relax.nn.batch_norm`, causing import to
fail during normalization
for mixed-dtype ONNX models.
This patch casts the ONNX BatchNormalization parameter tensors (`scale`,
`bias`, `mean`, and `var`) to the data tensor dtype before calling Relax
`batch_norm`.
This preserves the ONNX output dtype, which follows the input data
dtype, while keeping the fix localized to the frontend compatibility
layer.
The regression test builds a minimal ONNX BatchNormalization graph with
`float16` data and `float32` parameters, imports it through the Relax
ONNX frontend, and checks that the generated Relax `batch_norm` call
receives same-dtype inputs.
Verification:
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_batch_norm_mixed_dtype_params
tests/python/relax/test_frontend_onnx.py::test_batch_norm_defaults_to_inference_mode
-q`
Signed-off-by: viiccwen <vicwen@apache.org>
Fixes#19972
ONNX specifies that the second output of TopK, `indices`, has element
type `int64`, and the ONNX TopK operator spec constrains the index
tensor type to `tensor(int64)`:
https://onnx.ai/onnx/operators/onnx__TopK.html
The Relax ONNX frontend previously called `relax.op.topk` without
specifying the output indices dtype, so Relax used its default `int32`
indices.
This can make otherwise valid ONNX graphs fail during import when the
TopK indices are consumed by later integer/index operations that use
ONNX's usual `int64` constants. One example is `TopK -> Div`, where
Relax rejects the binary operation because the imported TopK indices are
`int32` while the divisor is `int64`.
This patch passes `dtype="int64"` when importing ONNX TopK, matching the
ONNX operator spec. It also updates the existing TopK frontend test to
check output dtypes, so the imported indices must match ONNX Runtime's
`int64` output.
Verification:
- `uv run --no-sync python -m pytest
tests/python/relax/test_frontend_onnx.py::test_topk -q`
Signed-off-by: viiccwen <vicwen@apache.org>
The fp8 group-wise scaled GEMM kernels passed a braced DLDataType
initializer as the second argument of the two-argument TVM_FFI_ICHECK_EQ
macro (e.g. TVM_FFI_ICHECK_EQ(a->dtype, DLDataType{kDLFloat8_e4m3fn, 8,
1})). The preprocessor ignores brace grouping and splits on the commas
inside {...}, so it sees four arguments and fails to compile. Wrap the
initializer in parentheses so it is treated as a single macro argument.
thrust.cu calls args[i].cast<DLTensor*>() but did not include
<tvm/ffi/container/tensor.h>, which defines TypeTraits<DLTensor*>;
without it the cast fails template deduction. Add the include.
Both issues break the CUDA runtime build; with them fixed it compiles
cleanly with USE_CUTLASS and USE_THRUST enabled.
A tvm_callback_metal_compile used purely for debugging/inspection may
return the MSL source unchanged. Previously, merely registering the
callback forced the module format to "metallib", so the runtime tried to
load the text source as a binary metallib and failed with "Invalid
library file" (issue #18798).
The callback may now return a (payload, format) pair to declare the
payload format. A bare str/bytes return keeps the legacy metallib
behavior. All kernels of a module share a single declared format, so a
callback that mixes formats across kernels (including a legacy metallib
return alongside a (payload, "metal") return) is rejected at codegen
time instead of producing a module that fails to load.
This PR fixes Relax bucketize lowering to pass the correct integer
output dtype to TOPI searchsorted.
Previously, both LegalizeOps and DispatchSortScan passed the input
tensor dtype as the output dtype. For float input tensors, this caused
TOPI searchsorted to receive a float output dtype, which later failed
during binary-search lowering because bucket indices must be integer
values.
This patch derives the output dtype from bucketize's out_int32
attribute:
- int32 when out_int32=True
- int64 otherwise
A numerical ExportedProgram frontend test is added to cover:
- right=False
- right=True
- out_int32=False
- out_int32=True
- float input values on bucket boundaries
Test:
python -m pytest
tests/python/relax/test_frontend_from_exported_program.py -k "bucketize"
-q
## Why
Inverse index map input variables were always named axis0, axis1, ...,
making generated IR harder to read.
## How
- Name a pass-through inverse input after its source Var instead of
axis{i}.
- Keep the axis{i} fallback for computed (non-pass-through) indices.
- Add test_inverse_preserves_passthrough_var_names covering a transpose
map.
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
This PR fixes the PyTorch ExportedProgram importer lowering for
`torch.masked_select`.
Previously, `masked_select` lowered to:
- flatten data and mask
- `nonzero(mask_flat)`
- `squeeze(axis=[0])`
- `take(data_flat, indices)`
However, the result of `R.nonzero` only carried rank information. The
following `R.squeeze` over the dynamic nonzero output could remain
unhandled during build/VM execution.
This PR inserts a `match_cast` after `R.nonzero` using the exported
output metadata, preserving the dynamic selected-length dimension before
`squeeze`.
A numerical regression test is also added to cover:
PyTorch eager -> torch.export -> Relax import -> build -> VM run ->
output comparison
Testing:
- `python -m pytest -q
tests/python/relax/test_frontend_from_exported_program.py -k
'masked_select'`
Hongyi Jin