## 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
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
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.
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.
- `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>
AttrStmt.node is an ffi::Any field, so converting POD arguments to
PrimExpr makes its representation depend on the caller rather than the
declared container type.
This change preserves values passed through AttrStmt and T.attr, uses
raw zero sentinel nodes consistently, and updates the printer
canonicalization.
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.
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.
### 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.
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
`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(...)`.
## 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.
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
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'`
## Why
- `torch.compile(backend=relax_dynamo(), dynamic=True)` lifts SymInt
scalars as scalar graph inputs
- `from_fx` skips these placeholders, so ops referencing one, e.g.
`view(x.size(0), -1)`, fail with `KeyError`
## How
- Bind sym placeholders to the same-named `tir.Var` from the input
tensors' symbolic shapes; skip as before when none exists
- Add `test_relax_dynamo_dynamic_sym_input_reference`; fails with
`KeyError` without the fix
Remove the Relax-specific Id indirection and use Var/DataflowVar object
identity directly.
Type-changing rewrites now remap definitions, uses, and binding lookups
coherently while preserving reflection, serialization, and the
DataflowVar subtype. Existing tests are adjusted for the API change; no
new test files or test cases are added.
Validation: the compiler and C++ tests build successfully; focused C++
coverage passes 4/4; the affected existing Python matrices pass 549
tests with 2 expected xfails; source censuses, diff checks, and
applicable hooks pass.
## Summary
- Keep the Python test launcher close to plain `pytest -n auto`, move
nightly tests under `tests/nightly/python`, remove obsolete launchers
and collection bookkeeping, and partition CPU/GPU jobs with explicit
`gpu` marker expressions.
- Repair exact-pointer regressions at their owning boundaries: packed
raw-string ABI values, CUDA/Metal matrix intrinsic pointers, internal TE
extern offsets, MetaSchedule scalar annotations, localized
auto-tensorization scope matching, and typed DLTensor fixture fields.
- Preserve typed workspace calls in TIR and cast pointer-returning
external calls in CodeGenC, covered by a plain-TIRx 1024-byte global
workspace that is compiled as C++.
- Finish phasing out value-bearing Relax `R.Prim` annotations by
requiring an explicit dtype, removing obsolete value-based contracts,
and expressing the DISCO rank-dependent slices as explicit scalar
`call_tir` inputs.
- Gate the distributed callback on the optional DISCO runtime, NCCL, and
at least two GPUs so capability-limited jobs skip instead of failing.
- Remove the non-demonstrating pointer probe, use direct TVMScript
comparison for packed strings, and remove the four designated legacy
testing modules.
The seven repaired CPU categories cover packed raw strings (7 failures),
CUDA/Metal matrix access-pointer types (7), internal TE extern offsets
(1), a typed DLTensor fixture (1), MetaSchedule scalar annotations (1),
CodeGenC workspace return casts (12), and localized auto-tensorization
storage-scope matching (19).
## Validation
- Base: `ded6ad8dd212869c881efb5590f8a33fc972728e`
- Head: `a7277e86dbcfe0638c8c252d36760859c4ab4297`
- All 35 locally available original failing node IDs pass across the
focused runs.
- The full focused TE, TIR builtin-lowering, and CodeGenC files pass: 61
tests.
- The complete touched Relax/TVMScript set plus
PlanAndUpdateBufferAllocationLocation passes with 784 passed, 20
skipped, and 1 expected failure.
- The DISCO callback collects and skips when its runtime or two-GPU
environment is unavailable.
- Six direct mapping tests, twelve tensor-core sketches, and the dp4a
sketch pass unchanged.
- The compiler rebuild, branch-wide pre-commit hooks, and full-range
whitespace checks pass.
- The 13 broad CBLAS/TFLite nodes remain dependency-gated; their owning
TE and generated-C regressions compile.
No merge is included in this change.
This PR simplifies Jenkins pytest execution around standard pytest-xdist
behavior.
- Runs each already-filtered CPU/GPU suite once with `-n auto`; the
broad suite keeps load-group scheduling because its order-sensitive
cases require it.
- Removes external sharding, wrapper/profile code, JUnit XML generation
and publication, the skipped-test XML consumer, obsolete suite naming,
and orphaned helpers.
- Retains one inert `task_clear_pytest.sh` entry point only because PR
jobs evaluate their Jenkinsfile from the trusted base branch before
checking out the PR; it performs no cleanup or reporting and can be
removed after this pipeline lands.
- Corrects stale broad-suite paths and explicit target guards, and
migrates a scalar stride test to the current `T.handle` pointer
semantics while preserving its negative lowering check.
- Prevents nested MetaSchedule/XGBoost unit tests from multiplying CPU
fanout without serializing the full suite.
- Builds only the `tvm_runtime` target for the secondary GPU
configuration and removes its unconsumed `gpu2` artifact upload.
The result reduces parallelism to one layer managed by pytest-xdist
while preserving GPU filtering and native failure visibility.
## Rationale
TIRx variables use inherited `ExprNode::ty` as their single semantic
type. Retaining a primitive handle surrogate erases the distinction
between scalar values, typed pointers, and true opaque pointers, then
forces later passes and code generators to reconstruct information that
the IR already owns.
## Changes
- Remove the duplicate reflected `Var::type_annotation` state and
preserve exact `PrimType` or `PointerType` through construction,
visitors, transforms, specialization, builders, printers, and code
generation.
- Keep scalar-only boundaries explicit through `PrimExpr`, `PrimVar`,
and `PrimType`; pointer-capable values remain general `Expr` or `Var`.
- Keep helper boundaries no broader than their contracts: TE tensor
variable indices use `PrimVar`, while expression deep equality recurses
through general `Expr` only where pointer-bearing `Call` arguments
require it and does not generalize private arithmetic subclasses.
- Keep core statement reflection typed as `Expr`, name general
reinterpret targets as `target_ty`, and preserve exact pointer calls in
the general vectorization path with explicit scalarization behavior.
- Delete `PrimType::Handle()` and `PrimType::IsHandle()`. True opaque
pointers use `PointerType::VoidPointerTy()`; TVMScript renders the
canonical global type as `T.handle`, standalone values as `T.handle()`,
and scoped void pointers with a keyword-only storage scope.
- Make `CodeGenSourceBase::SSAGetID` a single `Type` boundary across
source backends, without a separate primitive-type or runtime-dtype
variant.
- Keep WebGPU semantic argument classification type-aware: storage
buffers are identified from `PointerType`, POD arguments from
`PrimType`, and only the final `FunctionInfo` launch ABI is serialized
to `DLDataType`.
- Preserve exact pointer semantics at runtime boundaries, including
access pointers, packed calls and returns, external calls, storage
rewrites, and target-specific lowering.
## Migration guide
- **Variable types:** In C++, replace `var->type_annotation` with
`var->ty`; in Python, replace `var.type_annotation` with `var.ty`. The
result is the exact `Type`: scalar variables carry `PrimType`, while
pointer variables carry `PointerType`.
- **Scalar boundaries:** Use `PrimVar` and `PrimExpr` for variables and
expressions that are semantically scalar. When starting from a general
view, narrow explicitly with `var.as_or_throw<PrimVar>()` or
`expr.as_or_throw<PrimExpr>()`. Keep pointer-capable fields and call
arguments as `Var` or `Expr`. A default-constructed `PrimVar` is
nullable, so construct local scalar variables explicitly, for example
`PrimVar i("i")`.
- **Opaque pointers:** Replace `PrimType::Handle()` with
`PointerType::VoidPointerTy()`. Replace `IsHandle()` tests with explicit
`PointerType` inspection; use `PointerType(element_type, storage_scope)`
when the pointee type is known instead of erasing it to a runtime handle
dtype.
- **TVMScript handles:** Use `arg: T.handle` for a global void-pointer
annotation and `arg = T.handle()` for a standalone value. Use
`T.handle(storage_scope="shared")` for a scoped void pointer. Typed
pointers use forms such as `T.handle("float32")`, `T.handle("float32",
"global")`, or `T.handle("float32", "shared")`. Legacy
`T.handle("void")` input remains parse-compatible, but the printer
canonicalizes it to `T.handle` (or the keyword-only scoped form).
- The separate `tirx.type_annotation` intrinsic used by access-pointer
APIs is unchanged; this migration removes only the duplicate variable
field.
## Validation
- Complete native C++ test executable: 122/122 passed, including
`IRF.CountVar`.
- Relax binding-rewrite suite: 12/12 passed, including transferred-user
bookkeeping.
- Canonical typed/void/scoped TVMScript handle printer and round-trip
checks: 5/5 passed.
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.
## Rationale
After `PrimExpr` and `Expr` share one typed expression hierarchy,
expression types are part of semantic identity. Structurally identical
syntax with different types compare and hash differently, while source
spans remain diagnostic metadata.
## Invariant
`ExprNode::ty` participates in structural equality and hashing by
default. `GlobalVar` and Relax variables retain their symbol identity
rules. `tirx.PrimFunc` compares and hashes authoritative source fields
while excluding its derived type cache until all transformation paths
maintain that cache eagerly. Nested symbolic-shape rendering is isolated
from outer diagnostic configuration so diagnostic context cannot become
script-token content.
## Changes
- include expression types in generic structural equality and hashing
- preserve GlobalVar and Relax variable identity plus definition-safe
SeqExpr traversal
- compare and hash PrimFunc from authoritative fields while excluding
its stale derived type cache
- normalize narrow Relax construction and expected-fixture types exposed
by stricter identity
- isolate nested symbolic-shape token rendering from outer printer
configuration
Expression unification gives TIRX a shared `Expr` surface, but its
visitor and mutator APIs still expose primitive-only signatures. That
mismatch prevents general expressions from flowing through the existing
traversal structure and leaves statement traversal with overlapping
customization hooks.
This refactor generalizes the existing `ExprFunctor`, `ExprVisitor`, and
`ExprMutator` signatures in place to accept and return `Expr`. Statement
visitors and mutators expose a single virtual `VisitExpr(const Expr&)`
hook, while primitive statement reconstruction uses a non-virtual
checked `VisitPrimExpr` helper so invalid narrowing fails at the
boundary. Public pre-order and post-order traversal entry points accept
general `Expr` roots.
The existing specialization, vtable, dispatch registration, and class
structure remain intact; the change adds no parallel functor, fallback
dispatcher, or alternate implementation path.
## Rationale
`SizeVar` encodes nonnegativity in runtime subtype identity, which is
fragile under cloning and remapping. Symbolic integer values should use
one `Var` representation, with nonnegative facts recorded in the
analyzer at the use sites that establish them.
## Changes
- Remove `SizeVar` from the C++, Python, TE, TVMScript, FFI, visitor,
and serialization surfaces, and migrate callers to `Var`.
- Preserve the existing Relax constraint ownership model and use
`MarkGlobalNonNegValue` as the canonical path for global nonnegative
facts.
- Preserve `T.handle()` as the normal opaque-handle form. An optional
dtype constructs a typed pointer, with `T.handle("void")` reserved for
an explicit pointer-to-void.
## Summary
- Make `PrimExpr` a typed C++ view over `Expr` values whose
`ExprNode::ty` is `PrimType`, instead of using a separate runtime node
class as the proof of primitive-ness.
- Use the shared `ir::Call` node for Relax, TIRX, and primitive-valued
calls, while keeping primitive-only APIs explicit at their semantic
boundaries.
- Keep Python on the general `Expr` surface for primitive-typed values
so `isinstance` behavior does not imply a nominal primitive-expression
subclass.
## Design Rationale
The main advantage of this change is that common expression nodes such
as `Call` can be unified without specializing each one to `PrimType`. A
single `ir::Call` can represent a Relax tensor call, a Relax scalar
call, or a primitive-valued intrinsic call; the result type stored in
`ExprNode::ty` determines whether that particular value can be viewed as
`PrimExpr`.
This keeps the IR node hierarchy focused on expression structure rather
than result-type categories. Nodes that are intrinsically primitive,
such as integer and floating-point literals or TIRX primitive operators,
still have strongly typed C++ APIs and data structures. General nodes
whose result type may vary, such as `Call`, remain general `Expr` nodes
and are narrowed to `PrimExpr` only where primitive-only semantics are
required.
The PR also keeps the compatibility surface practical: C++
primitive-only APIs continue to accept `PrimExpr`, Python exposes a
compatibility predicate for checking the primitive typed category, and
visitors/printers use one natural `Call` path rather than duplicating
Relax and primitive call handling. Missing expression types are
represented explicitly with `Type::Missing()` so constructors can leave
type inference to later analysis without relying on nullable `Type`
values.
## Related Issue
closes#19689
## Why
The Relax AffineGrid op only handled 2D (4D theta/grid); 5D 3D inputs
from ONNX failed.
## How
- Generalize struct-info inference to 2D/3D via spatial =
size_sinfo->ndim.
- Branch TOPI affine_grid compute on 2D vs 3D.
- Add the 3D permute path in the frontend and a test_affine_grid_3d
case.
---------
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
RISC-V target tags use the LLVM codegen backend, so their target kind
should remain "llvm". However, the target metadata should still identify
the target as a RISC-V CPU rather than an ARM CPU.
Previously, the RISC-V tag helper used the ARM CPU keys and device
metadata, so Target("riscv/...") expanded with keys ["arm_cpu", "cpu"]
and device "arm_cpu".
```python
{"kind": "llvm", "keys": ["arm_cpu", "cpu"], "device": "arm_cpu"}
```
This is misleading for code that inspects target keys or device metadata
to distinguish CPU families.
This change updates the RISC-V tag helper to use keys ["riscv_cpu",
"cpu"] and device "riscv_cpu", while keeping kind="llvm". It also adds
the SpacemiT K3 RISC-V target tag.
After the tirx refactor, a `T.let` binding no longer implicitly casts
its right-hand side to the annotated dtype. In
`gpu_multinomial_from_uniform`, `row_idx` is annotated `int64` but is
loaded from the `row_indices` buffer, whose dtype is the configurable
`sample_indices_dtype` and may be `int32`. Relying on the let annotation
to widen the value is no longer valid and yields a dtype mismatch.
Wrap the load in an explicit `T.Cast("int64", ...)` so the row index is
always converted to the annotated int64 type regardless of
`sample_indices_dtype`.
Hongyi Jin