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>
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.
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
`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.
### 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.
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
Fixes#17745. `has_reshape_pattern` builds an inverse index map that
divides by each iter extent, so a zero-extent iter crashed with
divide-by-zero.
## How
- Skip the fused-var check when any block iter has zero extent; such
blocks touch no elements, so they are not reshape patterns.
- Added `test_reshape_pattern_zero_extent` in
`tests/python/relax/test_analysis.py`.
## 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.
## 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.
Replace the nested `DocToPythonScript` invocation in Relax
dependent-shape printing with an expression-string Doc rendered during
the active `PythonDocPrinter` traversal.
This keeps recursive IR-to-doc conversion inside the active docsifier,
preserves naming, precedence, escaping, source paths, and printer
configuration, and avoids a nested top-level renderer or a new public
rendering entry point.
Expression-string escaping is streamed directly into the final output so
wrapper and nested source spans retain exact escaped byte offsets.
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
## 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>
FlashInfer 0.6.3 changes the paged-attention plan/run ABI: the
prefill/decode plans take new arguments (e.g. window_left,
fixed_split_size, disable_split_kv; the decode plan now dispatches dtype
through empty q/kv tensors), the runs add enable_pdl and drop the
explicit stream, and the kernels consume separate key/value paged caches
read through tensor strides rather than one combined tensor. This
updates
the runtime attention backend (paged MHA, ragged, decode and MLA) to the
new signatures and to the Array<int64_t> plan-info representation.
FlashInfer 0.6.3 reads tensors from `data` directly and does not honor
the DLPack `byte_offset` field. mlc's auxiliary index tensors
(qo_indptr,
kv_indptr, page_indptr, page_indices, length_info) are views packed into
a shared workspace and so carry a non-zero byte_offset; passed as-is the
kernels read the wrong addresses (e.g. a ragged prefill processed only
the first query row). Three zero-copy DLPack view helpers address this:
`ZeroByteOffsetView` folds byte_offset into the data pointer,
`PagedKVCacheView` exposes the combined (num_pages, 2, ...) page tensor
as separate strided key/value caches, and `SliceLastDimView` slices the
last dimension for MLA.
This also completes the MLA FlashInfer path, which previously shipped
only the test and module generator. The MLA run splits the query into
nope/pe parts and the paged cache into ckv/kpe parts, and the ragged
self-attention is given its own uncompressed head dims and per-query kv
head count via a 5-element backend spec, since they differ from the
compressed MLA cache.
The MHA and MLA FlashInfer KV-cache tests are re-enabled as regression
coverage, guarded on FlashInfer availability (inline-RoPE is skipped as
unsupported by FlashInfer).
Structural diagnostics can identify a field below an object that
TVMScript renders without exposing that field. An underline alone then
points at the nearest visible parent and hides the full internal
location. This change keeps Script string-returning while making the
diagnostic context self-contained.
When the requested path is <root>.dtype, Script now returns this string
by default:
```text
Access path: <root>.dtype
Note: The underlined object is the nearest visible parent of this path.
T.int32
^^^^^^^
```
render_invisible_path_info defaults to true. Calls without target paths
are unchanged, and callers can set it to false to retain the legacy
underline-only string.
The implementation reuses the printer span-selection logic to capture
the deepest visible path and assembles the minimal
access-path/note/script block in C++. Pass-error enrichment uses the
same Script path. Production Python remains unchanged; focused Python
tests assert the complete strings for default-on, explicit-false,
hidden, exact-visible, unavailable-visible, pass-error, and
structural-equality cases.
## Why
After #19890 moved Relax dtype to optional, two spots still used the
deprecated `DLDataType{kDLOpaqueHandle, 0, 0}` void sentinel.
## How
- Drop the now-unreachable unknown-dtype branch in IsBoolType (handled
earlier by IsUnknownDtype()).
- Replace the raw void-sentinel comparison in ones/zeros with
PrimType::IsVoid().
---------
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Akaash Parthasarathy