## 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.
## 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.
This PR lets Relax expressions directly take `PrimExpr` values without
requiring the explicit `PrimValue` wrapper, continuing the Relax IR
unification work by removing Relax-specific leaf/base expression layers.
Summary:
- Remove `LeafExpr` / `LeafExprNode` and use direct expression-node
checks where needed.
- Converge Relax expression typing onto the shared IR `Expr` base.
- Remove the `PrimValue` node wrapper while keeping `relax.prim_value` /
`R.prim_value` as conversion helpers that return existing `PrimExpr`
values unchanged.
- Register direct `PrimExpr` handling through exact concrete node
dispatch, aligned with the `tirx` expression visitor list and excluding
arith iter-map intermediate nodes.
- Inline the private Python primitive conversion helper into public
`relax.prim_value`.
- Handle direct `PrimExpr` values in frontend scalar paths without
assuming a `.value` field on non-immediate expressions.
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:
- Remove the Relax-specific PrimType node/API and use canonical
ir.PrimType for dtype-only scalar types.
- Update parser, printer, analysis, op inference/legalization, and tests
to avoid value-bearing PrimType semantics.
- Preserve scalar values where needed by reading PrimValue expressions
directly instead of storing values in the type.
## 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
Replace TVM's `Diagnostic` / `DiagnosticContext` machinery with the
tvm-ffi
`visit_error_context` mechanism. Validators throw an `ffi::Error` seeded
with the
offending node; leaf pass executors (`ModulePass` / relax `Function` /
`DataflowBlock`) catch and rethrow `EnrichPassErrorWithContext`, which
appends the
failing pass name and a TVMScript-rendered, underlined source location.
`relax.analysis.well_formed` now throws on the first violation; a new
`check_well_formed` returns a bool, and all C++/Python/test callers are
routed
accordingly. `include/tvm/ir/diagnostic.h` and `src/ir/diagnostic.cc`
are deleted.
The enrichment renders with `num_context_lines=10` so a small function
shows in
full with no skipped-lines marker, while a large module stays bounded.
The TVMScript parser diagnostics
(`python/tvm/script/parser/core/diagnostics.py`)
stay self-contained pure-Python with no `DiagnosticContext` dependency,
and
restore multi-line source rendering: a diagnostic whose offending AST
node spans
multiple source lines now renders every spanned line with its gutter
line number
and an underline covering the span. `tvm.error.DiagnosticError` (used by
the
TVMScript parser) is retained.
A rendered end-to-end enriched-error example is posted as a comment
below.
## 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 enables ruff pyupgrade (UP) rules with py310 target, auto-fixing
~5600 annotation modernizations (PEP 585 generics, PEP 604 unions,
deprecated typing imports).
Also removes from __future__ import annotations from ir/module.py and
rmsnorm.py, bumps requires-python to >=3.10, and removes absolute_import
aliases from topi/contrib files.
This PR phases out legacy target string format in favor of the json
style format that is more well formed. It also simplfies our overall
code in handling multiple formats.
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
Hi Commiters,
This PR is trying to remove duplicated test case
`test_if_branch_var_scope`. Any suggestions would be appreciated if you
are available.
### Root Cause
Both tests `test_var_if_scoping_fail` and `test_if_branch_var_scope`
verify that returning a variable `w` that's only defined within if/else
branches raises an error. They're duplicates and should probably be
consolidated into a single test.
<img width="1388" height="306" alt="image"
src="https://github.com/user-attachments/assets/8e6defd4-fe83-419d-8793-67454ccb561f"
/>
### Solution
Remove duplicated test case `test_if_branch_var_scope`
Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
- Set Python minimal version to 3.9 as Python 3.8 has reached end of life
- Update AST handling to match Python 3.9+ API changes
- Fix deprecation warning: "Support for arbitrary keyword arguments is deprecated and will be removed in Python 3.15"
- Update TVMScript parser components to use modern AST interfaces
- Adjust linting and test infrastructure for Python 3.9+ compatibility
This PR modernizes the FFI foundation of the project and introduce
a new minimal and lightweight module [tvm ffi](https://github.com/apache/tvm/tree/refactor-s3/ffi)
based on our lessons in the past few years. It implements a modern
version of the [Unified Packed and Object RFC](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0097-unify-packed-and-object.md)
that unifies the packed function call and object systems.
Summary of the change:
- A dedicated clean Any/AnyView that can store strong and weak
references of items
- Function(previously PackedFunc) system built on top of the Any/AnyView
- A minimal C API that backs the overall calls. We are stabilizing the
API with a goal to bring clean, stable FFI conventions for both compiled
and registered code
- A rewrite of core python binding and generated code based on the module
- Update existing code and test cases to the new module
- Latest dlpack support
The new module brings many benefits thanks to the cleaner design,
to name a few:
- Any can support both POD types(int) and object types.
- Containers (e.g. Array) can now also contain Any value, e.g. now
`Array<int>` is supported, no need for boxed types
- Error handling now upgrades to object-based, allowing cleaner
traceback across languages
- Map now preserves insertion orders
- Path toward isolated stabilize minimum core ABI/API foundation module
- Type traits based design that cleanly defines how values interact
with Any system
- Automatic conversion of different types based on traits if needed
Because FFI upgrade is at heart of the project, the change touches every
component of the system. Importantly, this is an upgrade of the ABI so the
change is not backward compatible. The code compiled under the old
FFI won't work under the new one. We did provide example ABI translation
(e.g. LegacyTVMArgValueToFFIAny) functions for compatibility.
The PR tries to leave files in their old places while creating redirections.
The goal is to have the first milestone landed and infrastructure in place,
so we can do further refactors to complete features and cleanup legacy code
as trackable PRs. As of now, python binding and compiled code are under the
new convention while RPC and some other bindings still relies on legacy ABI
translation. We will work on upgrades in the coming PRs, including areas such
as reflection, phasing out legacy redirections etc.
* [Relax] Validate StructInfo annotations in well-formed check
Prior to this commit, the Relax well-formed checker verified that each
expression had a non-null `StructInfo` annotation, but did not perform
any validation on the contents of the `StructInfo` annotation.
This commit updates the Relax well-formed check to verify that the
`StructInfo` annotations are accurate by comparing against the
`StructInfo` that would be inferred for an expression. (This only
requires that the information is accurate, not that it is complete.
For example, an expression that is inferred to be
`R.Tensor(shape=[128,8], dtype="float32")` may have annotation of
`R.Tensor(ndim=2, dtype="float32"`, but may not have an annotation of
`R.Tensor(shape=[128,8], dtype="int32")`.)
* lint fix
* lint fix
Prior to this commit, TVMScript required the return value of a Relax
to be specified outside of any `with R.dataflow()` blocks. This
resulted in a common pattern, where the return value of a function was
first called with `R.output(ret_value)`, to mark `ret_value` as a
`tvm::relax::Var` instead of a `tvm::relax::DataflowVar`, followed
immediately by a `return ret_value` statement.
This commit updates the TVMScript parser to allow a `return` statement
inside a `with R.dataflow()` block. This is syntactic sugar that
is equivalent to calling `R.output`, followed by a `return`.
With this change, the following two TVMScript examples are now
equivalent. (Prior to this change, the `return_inside_dataflow`
example would raise an error during parsing.)
```python
@R.function(private=True)
def output_then_return(A: R.Tensor):
with R.dataflow():
B = R.add(A, A)
C = R.multiply(B, B)
R.output(C)
return C
@R.function(private=True)
def return_inside_dataflow(A: R.Tensor):
with R.dataflow():
B = R.add(A, A)
C = R.multiply(B, B)
return C
```
* [TVMScript] Avoid segfault from invalid TVMScript
Prior to this commit, after the `DiagnosticContext` prints its error,
it overwrites the `DiagnosticRenderer` with a NULL renderer. If a
second call to `DiagnosticContext::Render` occurs, it will segfault.
This appears to be intended to prevent double-printing of error
messages, but double-printing error messages is much worse than a
segfault.
In addition, `DiagnosticContext::Render` should only be called once.
There's a common pattern in the parser where it will wrap exceptions
in `DiagnosticError`, but re-raise exceptions that are already a
`DiagnosticError`. This requires every such location to include
`except DiagnosticError: raise`, and can easily be missed.
This PR makes two changes: First, the `DiagnosticRenderer` is updated
to have a no-op callback rather than a NULL callback. Second, the
re-raising of `DiagnosticError` is moved to `Parser.report_error`, so
that it does not need to be handled separately at several independent
locations in the TVMScript parser.
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.
* [Relax] Avoid wrapping TupleStructInfo into a Tuple for R.call_tir
Prior to this commit, the different `R.call_tir*` variations would
wrap the arguments into an in-line `relax.Tuple`, if it is not
already a `relax.Tuple`. While this allows a tensor to be passed into
these functions as a single argument (`R.call_tir(func, arg, ...)`
instead of `R.call_tir(func, [arg], ...)`), the wrapped Relax variable
may already refer to a tuple.
This use of a variable to refer to an argument tuple rather than an
in-line argument tuple is not allowed by Relax. (See discussion on
https://github.com/apache/tvm/pull/15916 for details.) However, by
wrapping a variable `args: R.Tuple(R.Tensor, R.Tensor, ...)` into a
tuple-of-tuples, the error occurs after the expression has already
been generated, and refers to an expression `R.Tuple(R.Tuple(R.Tensor,
R.Tensor, ...))` that doesn't appear anywhere in the user's input.
This can make debugging difficult (see
https://github.com/apache/tvm/issues/17239 for an example).
This commit updates the argument-handling in `R.call_tir` to only
generate an in-line `relax.Tuple` if the arguments do not already have
`relax.TupleStructInfo`. If the argument was provided as a Relax
variable bound to a tuple of arguments, it will still produce an
error. However, that error will occur much earlier, and will
explicitly state that the argument must be a `relax.Tuple` instead of
a `relax.Var`.
* lint fixes
* [TVMScript][Bugfix] Normalize relax::If with function's TIR var
Prior to this commit, the branches of `relax::If` were normalized
using `EraseToWellDefinedInScope`, using a fresh variable scope.
While this had the intended behavior of preventing variables defined
in a single branch from being usable outside of the conditional, it
also caused the conditional's branches to treat function-scope
symbolic variables as if they were undefined.
This commit updates the `tvm::relax::Normalizer` so that `relax::If`
is normalized within an inherited scope. This preserves the previous
behavior for symbolic variables defined within a branch, but allows
shapes within a branch to use symbolic variables defined outside of
the branch.
* [Relax] Canonicalize known symbolic shapes in Relax expressions
Prior to this commit, known constants in Relax functions would be
inlined by the `CanonicalizeBindings` pass, but only if they appeared as Relax
expressions (e.g. `R.const` or `R.prim_value`). Known constants that
appeared as TIR variables (e.g. symbolic shapes) would be kept as
dynamic parameters, even if they were known at compile time.
This commit updates the `CanonicalizeBindings` pass to identify known
values of symbolic shapes, and to use these known values in shape
expressions.
* [Relax][Refactor] Reorganize pattern-matching
A follow-up to https://github.com/apache/tvm/pull/16730. Now that the
implementations for `rewrite_call` and `rewrite_bindings` are in
separate classes, they can be further split out into separate files.
* [Relax][Refactor] Implement Rewriter class for pattern-rewrite
Prior to this commit, the pattern to be matched and the rewrite to be
performed were provided as separate arguments. This commit introduces
a new class `ExprRewriter`, which contains both parts.
This abstraction will make it easier to combine multiple different
rewrite rules, applying them in a single pass.
* lint fixes
* Remove unnecessary change which broke a unit test
* lint fix for import order
* Add docstrings
* lint fix
* Lint fix
* lint fixes
* lint fix
* Update based on review comments
* Add test case for matching against arbitrary dtype
* Fix breakage in unit tests
One unit test that had been relying on invalid shape propagation.
Another unit test that required constructed an ill-formed output to
test against.
* Updated base class name from ExprRewriter to PatternMatchingRewriter
* lint fix
Because Relax struct inference is performed while the function is
being built, all constraints on symbolic variables that are used
for simplifications must be provided to the analyzer. This is not
currently the case, nor is there a clear way to fix this issue.
A follow-up to https://github.com/apache/tvm/pull/16745. For Relax
functions produced in TVMScript, when `R.func_attrs` was not present,
the default was set to `None` instead of an empty dictionary.
* [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
* [Analysis] Allow calls to GlobalVar in @R.function
Prior to this commit, the post-parsing well-formed check performed by
TVMScript allowed a call to `GlobalVar` in a `@R.function`, but only
if it occurred within the context of a `@I.ir_module`. If
`@R.function` appeared on its own, calls to a `GlobalVar` would be
treated as calls to an undefined function.
* Use approrpirate well-formed checks TIR/Relax functions
* Lint fix
* Import order fix
This PR fixes the purity flag of `relax.vm.call_tir_dyn` and another
few "kill" ops. Their purity flags were set to True, which made them
possible to be removed by `remove_all_unused`.
* `relax.vm.call_tir_dyn` works by mutating the input args in place,
which is not pure.
* though the "kill" ops have no actions so far, their semantics
suggest that they are impure.
A regression test is added to prevent the unexpected removal from
happening again.
* Check well-formedness in the parser
* Correct packed funcs in NN frontend
* Support the check_well_formed optional argument to I.ir_module
* Also check well-formedness in TIR
* Enable normalization for individual Relax functions and PrimFuncs
* Use the error raised by the TIR well-formed checker for the message
* Fix tvmscript test failures
* Whitespace
* Fix errors in verify_well_formed test
* Include a more helpful error message
* Fix TIR test failures
* Address well-formed failures in test_tir_specialize
* Correct well-formedness error in test_tir_analysis_oob
* Correct further well-formedness failures
* Remove __tvm_meta__ from test case to avoid parsing error
* Avoid circular import in entryy.py
* Formatting fixes
* lint fix
* Add pylint exceptions
* Fix whitespace
* Fix more failed test cases
* Catch inappropriate use of decl_function instead of segfaulting
* Fix test_lower.py
* Mark purity in test_relax_2d_buffer_allocation.py
* Mark purity in test_dma_builtin.py
* Remove __tvm_meta___ from test_tir_usmp_analysis_extract_bufferinfo.py
* Suppress well-formed check in test_tir_transform_convert_blocks_to_opaque.py
* Remove __tvm_meta__ in test_tir_usmp_algo.py
* Remove __tvm_meta__ from more USMP tests
* Fix incorrect var in test_tir_transform_storage_flatten.py
* Remove all remaining instances of __tvm_meta__
* Fix purity error in test_dataflow_pattern.py
* Fix purity error in test_ast_printer
* Fix test_arith_domain_touched example
* Okay to set check_well_formed to True in test_tir_analysis_identify_mcmcpy
* Define variable in test_tir_analysis_oob
* Typo fix
* Add explanatory comment to test case
* Define the undefined vars in test_tir_transform_common_subexpr_elim
* Exception no longer necessary in test_tir_transform_inject_rolling_buffer
* Remove unnecessary check exemption in test_tir_transform_convert_ssa
* Avoid checking exemption in test_inject_ptx_ldg32
* Note special case in test_distributed_transform_propagate_sharding
* Exempt well-formed error in dlight/test_benchmark
* Exempt well-formedness errors in test_ethosu/, mostly uninitialized vars
* Whitespace
* Include non-CUDA GPUs in IsScheduledOnGPU
* Fix thread binding bug by changing thread binding var dtype
* Include overrides in test_runtime_builtin_paged_attention_kv_cache.py
* add exemptions in test_ethosu/test_replace_conv2d
* Add more ethosu exemptions
* More exemptions for ethosu tests
* Remove unused reference
* Indicate purity in test_transform_rewrite_cuda_graph
* Indicate purity in test_transform_normalize
* Reorder MergeSharedMemoryAllocations in GPU codegen
* Add target parameter for FP8StorageLegalize and FP8ComputeLegalize
* Don't re-import Target in tvm/tir/transform/transform.py
Prior to this commit, TVMScript required all relax expressions to be
part of an explicit assignment or return statement. While this
matches the structure of the C++ IR types, this can be unexpected for
functions that have no return value. For example, needing to assign
the result of `R.print(...)` to a variable.
This commit updates the TVMScript parser/printer to allow relax
expressions to be used as statements, if they have a void return type.
This allows use of `R.print(...)` and `R.assert_op(...)` to be called
without assigning the result to an unused variable.
Prior to this commit, using `R.match_cast` as the last binding would
produce a segfault, as `var_binding->value` was used instead of
`match_cast->value`. In addition, because the last binding of each
branch was removed, any changes to the struct info resulting from the
match cast were silently discarded.
This commit updates the TVMScript parsing of if/else statements to
remove the segfault and maintain the struct info changes produced by
the `R.match_cast`.
Prior to this commit, `R.call_packed` and `R.call_pure_packed` had
different normalization for the `sinfo_args` argument. While
`R.call_packed` checked if the struct info needed to be converted using
`ObjectGeneric.asobject()`, `R.call_pure_packed` did not.
This commit updates the `R.call_pure_packed` to handle `sinfo_args`
in the same manner as `R.call_packed`.
In low-level Relax (after pass `CallTIRewrite`), the `call_packed` nodes
do not always have explicit `sinfo_args`. This PR extents the parser to
support this case.
* Implement basic analyses
* Fix typo
* Add tests for analyses
* Include in-place analysis
* Return the lists instead
* Update python binding
* No need to assume *pure* functions capture all values ever passed to them. Also use pointers instead of non-const refs
* Improve handling of tuples in mystery call case
* Corrections to inplace checking
* Add test case for mystery value
* typo
* Add inplace test case, correct minor issues
* Consider also using larger tensors to store smaller ones
* Check call args against any possible target sinfo, also check tensor sinfo dtype
* Handle output vars and tuple get item
* Add legalization for in-place functions
* No need to update the NoAlias attribute, actually
* Fix TIR transformation, add tests for inline transformation
* Only find candidates from supported ops and list _all_ feasible argument indices
* Implement basic transformation pass
* Use a module pass so wider changes are visible, reorganize
* Have an end-to-end test case for the in-place transformation
* Rebase fixes and use GetBoundValue instead of reimplementing it
* Let's just use 'inplace' everywhere
* Reorganize code and add more documentation
* Include proper bounds check
* Trailing whitespace
* Need a trailing newline
* Remove unused imports
* Add docstrings for exposed inner functions
* Reformat docstrings to appease the linter
* C++ stylistic changes
* Treat args as mystery values by default, do not allow overwriting
* Formatting
* Clarify pass description
* Add check to ensure that testing functions are used only in a testing environment
* Improve size match check readability per review suggestions
* Improve the size match check per review suggestions (use PrimExprs)
* Treat non-dataflow vars as living past the end of the block in all cases
* Clarify notion of size in comment
* Remove commented-out code
* Assume any op that returns a tuple is returning a fresh one (exceptions can be noted later)
* Add full structural equality check in large test case
* Fix parser roundtripping bug with call_tir_inplace
* Refactor tests to ensure maps are nonempty
* Use .empty() where it's more reasonable
* linting changes
* Flipped the check by accident
* Remove debug print
* Factor out data structure for representing matches and match opportunities
* Style fix
* Use the analyzer to handle dynamic cases too
* Whitespace
* Use BlockBuilder APIs more to avoid re-normalizing
* Check for expired vars at start of loop so that the use of continue does not skip that step
---------
Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>
* [Unity][Analysis] Handle PrimStructInfo in EraseToWellDefined
Prior to this commit, the `EraseToWellDefined` pass would update
symbolic variable definitions in `ShapeStructInfo` and
`TensorStructInfo`, but did not in `PrimStructInfo`. This commit
updates the `WellDefinedEraser` to include symbolic variables defined
in `PrimStructInfo`.
* Update collecting of symbolic variables in InferSymbolicVarMap
* CI bump due to flaky unit test
`tests/python/relax/test_frontend_onnx.py::test_attention` fails for
some inputs. Failures on `unity` head occurred 3/100 test cases
* [Unity][TVMScript] Produce var = R.ExternFunc("") statements
Prior to this commit, any `ExternFunc` usage in a relax function would
print the string name of the function on its own line, omitting any
variable definition, and later use of the variable would occur without
a definition. This commit updates the printing of `R.ExternFunc` to
appear as a normal relax variable.
* Preserve special handling as callee, test round-trip
* Updated parser to handle `var = R.ExternFunc(...)` in IRModule
Since this is now a representation that may be produced by the
TVMScript printer, it must also be handled at the parser.
Prior to this commit, a `relax.PrimValue` could have a datatype, but
couldn't have a corresponding `tir.PrimExpr`. As a result, it could
not be used to specify tensor shapes. This makes some expressions
require fallback to `R.Tensor(ndim=ndim)`, even though the shape could
still be inferred.
```python
@R.function
def func(
A: R.Tensor(16, 16),
first_n_rows: R.prim("int64"),
) -> R.Tensor([first_n_rows, 16]):
# ^^^^^^^^^^^^
# R.Tensor requires a PrimExpr, not relax.Expr
#
# Operations may require PrimExpr
# vvvvvvvvvvvv
out = R.op.strided_slice(axis=[0], begin=[0], end=[first_n_rows])
return out
```
This commit adds a `Optional<PrimExpr> value` field to the
`PrimStructInfo`. This field acts similarly to the `PrimExpr` fields
already used in `ShapeStructInfo`, and may contain symbolic variables.
```python
@R.function
def func(
A: R.Tensor(16, 16),
# TIR definitions in signature allow in-line definitions,
# similar to R.Tensor and R.Shape. R.Prim takes `dtype` or
# `value` kwarg to distinguish between in-line symbolic variable
# and string representation of dtype.
first_n_rows: R.prim(value="first_n_rows_tir"),
) -> R.Tensor(["first_n_rows_tir", 16]):
# Body contains a TIR variable definition, which may be used
# in function calls, inferred shape annotations.
first_n_rows_tir = T.int64()
out = R.op.strided_slice(axis=[0], begin=[0], end=[first_n_rows])
return out
```
Use distinct PrimStructInfo arguments for dtype/value
Update TVMScript printer
Parser updates, Support R.Prim(value=...) annotations in function signature
* Added unit tests for new functionality in API, parser, printer
* Add unit tests for bind_symbolic_vars
* Add test cases to valid bind_symbolic_vars
In the same way as T.macro, R.macro support hygiene (via "hygienic"
keyword), but unlike TIR macros, Relax macros produce expressions,
and have to return a value: the last statement in R.macro must be
"return".
* Make it an error to mark a function private and specify a global symbol
* Fix improper use of global symbol
* Update privacy annotations in other tests
This PR implements the privacy annotation proposal. Namely, the @R.function decorator now has an optional private attribute. If a function is marked as private, then it will not have a global symbol attached to it and thus will not be externally accessible. By default, functions are not private, so the parser does insert a global symbol for them.
* 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
Tianqi Chen