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 introduces Relax AnyType as the primary top/base type spelling,
replacing the previous ObjectType naming for the type that represents
any Relax value.
Changes:
- Add AnyType/AnyTypeNode with relax.AnyType registration and keep
ObjectType/R.Object compatibility aliases.
- Update Relax type analysis, type visitors, opaque function defaults,
and script printer/parser handling to use AnyType/R.Any.
- Migrate affected Python/C++ call sites, docs, and focused tests to the
new spelling.
Validation:
- cmake --build build --parallel 16
- Focused Relax/TVMScript pytest: 704 passed, 1 xfailed
- pre_commit run --files <changed files>
## Summary
- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
## Summary
This PR adds the initial TIRx support needed for low-level programming
of Blackwell-class GPU architectures. As part of the ongoing TIRx
refactor, it introduces TVMScript support for directly scripting
advanced hardware features without relying on scheduling as the primary
programming interface.
The change keeps existing `s_tir` script support intact while making
direct scripting a first-class path for TIRx programs.
## Main Changes
- Add TIRx operator dispatch and layout infrastructure.
- Add TVMScript support for new low-level TIRx operations.
- Add analysis, transform, and lowering support for TIRx IR nodes.
- Add CUDA/Blackwell-oriented codegen and intrinsic coverage.
- Add Python and C++ integration points for TIRx scripting and runtime
support.
## Validation
- `pre-commit run --all-files`
- `ninja -C build -j32`
- `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16`
- `1723 passed, 47 skipped, 32 warnings`
- `CUDA_VISIBLE_DEVICES=2 python -m pytest -v
tests/python/all-platform-minimal-test`
- `37 passed, 105 skipped`
- `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis
tests/python/tirx-base tests/python/tirx-transform -n 16`
- `664 passed, 25 skipped, 9 xfailed, 1 xpassed`
## Local CI Notes
Some full CI-equivalent jobs were not locally reproducible because this
machine is missing parts of the Apache TVM CI environment, including
`llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten,
and ARM/QEMU cross-toolchain components. Metal-specific tests were
skipped locally because no Metal runtime is available.
This PR brings up the tirx namespace. We have been spliting out the
original tir namespace to include high-level component s_tir and this PR
updates the remaining low-level part as tirx namespace
This PR renames tir.Block to SBlock. This clearly indicate the
scheduable property of the block and is a prereq for followup stir
passes refactor.
Main changes:
- Data structure change from Block to SBlock
- Syntax change from T.block to T.sblock
Introduces the below features over texture annotation
- Lowering, codegen and runtime for texture.
- image2d_array_t support - Added depth dimension allows more
allocations using texture instead of falling back to buffer when the
texture limits exceeds.
- A comprehensive set of schedules for Adreno textures.
- Texture packing of arbitrary types up to 128 bit (FP16-NCHW8c,
INT8-NCHW16c ...etc.).
- A clBufferDescriptor debug dump controlled by cmake options.
- Pipeline definition for adreno target.
While covering these features the below interfaces or passes or enhanced
which need a review.
- alloc_tensor: VDevice information is passed across these API's. The
way of texture allocation is ```alloc_storage``` allocates buffer/image
objects as requested followed by alloc_tensor being a view of any scope.
This takes care of optimum utilization backing memory across different
image objects or scopes.
- Constants Saving: Handled by adding memory scope section in
executable. This introduces a new header magic to retain the backward
compatibility.
- Static Memory Planing: Mostly port from Relay static memory planner
with mixed mode allocator.
---------
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Sanjay <sanjs@qti.qualcomm.com>
Prior to this commit, the Relax well-formed checker validated
arguments provided to Relax functions, but did not validate arguments
provided to `R.call_tir`. As a result, incorrect arguments from Relax
to TIR would not be checked until runtime, if at all.
This commit updates the well-formed checker to verify that
`R.call_tir` has received the correct arguments, and has the correct
output shape specified in the `out_sinfo` parameter.
Initial implementation performed the validation as part of
`FNormalize`, to maximize coverage of this check. This increased
end-to-end compilation time by ~10%, and so the check was requested to
be restricted to the well-formed checker. Expensive operator-specific
validation is now performed in the new `FValidate` attribute.
Prior to this commit, the error message produced when validating
`R.call_tir_inplace` included the shape of the argument that will be
mutated in-place. This correctly caught and raised an error when the argument is a
tensor with known shape that is incompatible with the output tensor's
shape. However, this same error message could be also be reached if
the input does not have `TensorStructInfo` at all, which would trigger
a segfault.
This commit updates the validation to print the argument's
`StructInfo` directly, rather than a field from the struct info. This
correctly raises an error for the cases where the argument is not a
tensor, or is a tensor with unknown dimensionality, while still
printing the explicit shape of the mismatched tensor when avalable.
* [UnitTests] Use tvm.ir.assert_structural_equal whenever possible
Prior to commit, many unit tests were implemented as `assert
tvm.ir.structural_equal(output, expected)`. While this is correct, it
doesn't provide much information when the test fails. The
`tvm.ir.assert_structural_equal` method performs the equivalent check,
but displays the exact location where a mismatch occurs.
This commit replaces all use of `assert tvm.ir.structural_equal` with
`tvm.ir.assert_structural_equal`.
* fix unit tests
* [TIR][Analysis] Implemented tir.analysis.is_pure_function
This commit introduces two related utilities,
`tir.analysis.is_pure_function` and `tir.analysis.assert_pure_function`.
In contrast to the existing `tvm::tir::SideEffect`, which checks for
side effects on a for a `PrimExpr`, `is_pure_function` checks for side
effects for the function as a whole.
* [Transform] Implement relax.transform.ComputePrimValue
Prior to this commit, while expressions of type `DataType::Int(64)`
could be computed in the `relax.transform.VMShapeLower`, expressions
of any other type could not. This commit introduces
`relax.transform.ComputePrimValue`, which produces `PrimFunc`
subroutines to compute `PrimExpr` values of any dtype.
This functionality will allow boolean values to be computed based on
the symbolic values known at runtime.
* [Relax] Allow R.Prim('bool') in relax::If and assert_op
Prior to this commit, the condition used for `relax::If` node and the
`"relax.assert_op"` operator was required to be a scalar tensor. This
made it difficult to alter behavior based on a runtime shape
parameter. For example, delegating to a vectorized implementation
based on a whether a tensor shape is divisible by the vector size.
This commit adds support for expressions of type `R.Prim('bool')` as
the conditional for `relax::If` and `"relax.assert_op"`, to allow
these use cases.
* Lint fix
* [Unity] Implement LowerAllocTensor to remove R.builtin.alloc_tensor
The `StaticPlanBlockMemory` transform is provided a module that
expresses all allocations with `R.builtin.alloc_tensor`, and produces
a module that uses `R.memory.alloc_storage` and
`R.memory.alloc_tensor` to express static allocations, while dynamic
allocations continue to use `R.builtin.alloc_tensor`.
Prior to this commit, this mixed output was handled as part of
`VMBuiltinLower`. This commit extracts the lowering of
`R.builtin.alloc_tensor` to a new pass, `LowerAllocTensor`. This pass
runs after `StaticPlanBlockMemory`, and replaces any remaining
`R.builtin.alloc_tensor` with calls to `R.memory.alloc_storage`
and `R.memory.alloc_tensor`.
* Updated unit tests
* Correct order of LowerAllocTensor and KillAfterLastUse
The `R.memory.alloc_storage` produced by `LowerAllocTensor` must be
present in order to be appropriately deleted by `KillAfterLastUse`.
* Set purity as an attribute in the @R.function decorator instead of using R.is_pure() or R.is_impure()
* Remove accidental debug prints
* Need to override pylint unused argument warning in function decorator
* Parser argument no longer needed for find_purity_annotation
This PR implements the tracking of function purity as part of the StructInfo system. This will allow the compiler to enforce that no impure function (one that can possibly have visible side effects) can be called in a DataflowBlock. Tracking this requires noting which operators are pure or impure, which is presently done using an operator attribute called `FPurity` (a simple boolean), and which Relax function calls are pure (via the StructInfo system).
It is difficult to infer the purity of a function in the general case (when there are calls to other Relax functions), so this change does require users to annotate impure functions using a new field on functions, is_pure (in TVMScript, this can be done using R.is_pure() or R.is_impure()). Since most Relax functions are likely to be pure and purity is the default assumption, this will hopefully not be a large imposition on users. We can consider eventually inferring purity in the easier cases, since those are likely to be common.
Note that PackedFuncs are conservatively treated as impure. However, in situations where they are needed inside a dataflow block, a call to a PackedFunc that is, in reality, pure can be done via the new operator call_pure_packed or the existing operator call_dps_packed (it is assumed that any PackedFunc used with it will be pure). (Similarly, a new operator invoke_pure_closure is introduced as a counterpart to invoke_closure for dealing with closure objects, though this really should only come up with the LambdaLifting pass.)
As an "escape hatch" to the purity system, one can use the attribute relax.force_pure, which indicates to the compiler to treat the entire function as pure even if it contains an impure call. Additionally, even though PackedFuncs are normally treated as impure, a user can use call_pure_packed or call_dps_packed to call PackedFuncs in dataflow blocks when appropriate. These can be used to deal with the following situations:
1. A function does side effects but only on a value that will not be exposed anywhere else or on a new value that will be returned. Even though the individual actions are "impure," the overall function fulfills the definition of being pure. relax.force_pure would be useful here.
2. A PackedFunc is, in reality, pure. call_pure_packed or call_dps_packed are useful in this situation.
Changes include:
* Enforcing that impure functions are not used in DataflowBlocks in the well-formed check.
* Enforcing that functions that are not labeled impure do not contain impure calls (unless relax.force_pure is set).
* Implementing the call_pure_packed operator
The current cross-function calls in TVMScript will cause PyLint warnings,
since the GlobalVar will be marked as undefined vars, e.g.:
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
gv0 = R.call_tir(tir_func, x, R.Tensor((128,), dtype="float32")) # <= `tir_func` is not defined in Python syntax.
return gv0
```
This PR changes the behavior into `TestModule.tir_func` instead of direct `tir_func`
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
cls = TestModule # Use `cls` to refer the current Module
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128,), dtype="float32"))
return gv0
```
NOTE: It's a breaking change, the old style is deprecated.
Additionally, this PR contains the following minor fixes:
- mark `R.function` as staticmethod as what we do for `T.prim_func`
- make `I`, `R`, `T`, `cls` be the builtin keywords for the printer
- define names for functions, modules to prevent naming conflict
- checking the var names is valid via regex expression
- fix typos
Given that some latest changes of TVMScript syntax have been merged,
some test files are now containing deprecated uses of TVMScript syntax.
This PR updates the test files with latest TVMScript syntax so that
running the tests will not trigger deprecation warnings.
Co-authored-by: Tianqi Chen <tqchen@users.noreply.github.com>
As we've introduced `arg_sinfo` in CallNode, implicit shape constructor
is not widely used in TVMScript. This PR removes the implicit shape since
it may cause confusion between shape and tuple.
Shushi Hong