## 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 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
## Summary
- Remove `FewShotTuning` pass from Relax transform (C++ implementation,
Python bindings, and test file)
- The pass is unused in the current codebase and can be safely removed
## Files Changed
- `include/tvm/relax/transform.h` — Remove declaration
- `python/tvm/relax/transform/__init__.py` — Remove from imports
- `python/tvm/relax/transform/transform.py` — Remove Python function
- `src/relax/transform/few_shot_tuning.cc` — Delete (C++ implementation)
- `tests/python/relax/test_transform_few_shot_tuning.py` — Delete (test
file)
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.
Fixed 'occured' to 'occurred' in transform.py and 'seperated' to
'separated' in mrvl.py.
Co-authored-by: thecaptain789 <thecaptain789@users.noreply.github.com>
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
This allows user defined callback to specify layouts dynamically based
on call description.
Helpful to alter layouts based on the operator shapes or attributes.
---------
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This PR Updates the NDArray => Tensor.
Both tensor and ndarray are commonly used terms.
Because the term Tensor is getting more common in the context of ML,
we do the rename to stay more aligned with torch.Tensor and DLTensor.
* [FFI][REFACTOR] Establish tvm_ffi as a standalone python module
This PR establishes tvm_ffi as a standalone python module.
The ffi is structured as a minimal pip module that can be
directly install by path or url.
examples/get_started provided a minimal example.
This is a major change as we are decoupling tvm_ffi as a
separate package, users need to install tvm_ffi separately.
Thanks to its minimal dependency, tvm_ffi can be easily installed
even just from the source by pip install ./ffi
This change would enable future improvement for library plugins
to have lightweight dependencies by just working on top of
the tvm_ffi, while the main compiler toolchain and runtime
can be layered on top.
* [FFI] Improve traceback setups
This PR improves traceback related setups
This PR phases out tvm._ffi redirections in favor of new FFI
new functions are now called via tvm.ffi.
We also enabled limited API support for python 3.12+
so the compiled binary can be forward compatible to future
python versions.
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.
This PR adds support for CPU weight prepacking. To be specific, this PR
adds a new pass `AttachAttrLayoutFreeBuffers` to attach layout free buffers
to the weight parameters, so that we can leverage MetaSchedule to optimize
the prepacking process.
After the pass and tuning, we introduce a new pass `SplitLayoutRewritePreproc`
to split the layout rewrite pass into multiple functions, so that we can lift
the parameters transform pass function with existing pass.
This PR introduces a static shape tuning pipeline for Relax. It is designed to work with
the MetaSchedule tuning framework to optimize the performance of the model.
Together with a minor typo fix
* [Docs] Introduce Relax API and move legacy part to standalone page
As the TVM project evolves, the Unity strategy has been the recommended
way to use Apache TVM applications. Hence, we are pushing documentation
for the Relax API to the forefront and moving the legacy part to a
standalone page, which may be removed in the future.
* update for ci
* update for ci
Previously, `R.view` was legalized to extern call to
`runtime.TVMArrayCreateView` during `LegalizeOps`. This call to extern
func can't be properly handled by `StaticBlockPlanMemory` because it
assumes the extern func does not retain the input buffer. Extern func
returning a view of the input would break the ref count of the
buffer. This PR defers the legalization of `R.view` so that it can be
explicitly handled by memory planning.
A new op `R.ensure_aligned` is added as discussed in #16955
* [Relax] Support input axis_separator to allow 2D to 1D conversion
Introduce input_axis_separator in relax.transform_layout op to allow conversion of 2D buffers to 1D buffers.
The conversion from 2D->1D is handled while lowering of transform_layout operator.
Also introducing support for input_axis_separator in AlterOpImpl pass.
* Fix LINT errors
* Fix review comments
* [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.
* [Relax][Transform] Provide callback versions of LazyTransformParams
Prior to this commit, the `LazyTransformParams` function could be used
to load model parameters on demand. However, the function used to
load or set parameters needed to be registered within the global
registry of `PackedFunc`s. This PR provides `LazyGetInput` and
`LazySetOutput` transforms, which perform the lazy-loading through a
`R.Callable` callback argument, rather than through a
globally-registered `PackedFunc`.
* Reverse the order of parameters in fget_param
If `fget_param` accepts the parameter index first, and the parameter
name second, then an implementation with signauture and default values
of `def fget_param(index: int, name: Optional[str]=None)` could be
used as either the callback of `LazyGetInput`, or as the
globally-registered `"get_item"` for the existing
`LazyTransformParams`, which should make it easier to transition
between the two.
* lint fix
* Updates based on review comments
* [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
* [WIP] LiftTransformParams for multiple functions
* pass test
* [In-Progress] Define desired behavior for shared LiftTransformParams
Currently, the `relax.transform.LiftTransformParams` pass produces a
separate `transform_params` function for every function in the
`IRModule`. In most cases, the functions in an `IRModule` all accept
the same set of model weights (e.g. `"prefill"` and `"decode"` in a
transformer model). However, the lifted `*_transform_params`
functions may be different for each inference function.
The goal is to introduce a new optional parameter `shared_transform`
for `LiftTransformParams`. If set, a single parameter transformation
function should be generated for the entire `IRModule`, rather than
one parameter transformation function for each original function.
Because the shared parameter transformation function must be
compatible with all existing functions, it should only contain
parameter transformation steps that are common across all input
functions.
* [TIR] Implemented shared lift transform params
* Comments & skip test.
* Linting.
* Avoid c++20 feature to pass CI.
* Remove unused code.
* Fix interface as suggested.
* Fix docs.
* Fix interface as suggested.
* Move code for readability.
---------
Co-authored-by: Wuwei Lin <wuwei@apache.org>
Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>
* [Relax] Implement relax.transform.TopologicalSort
This commit implements a utility `relax.transform.TopologicalSort`,
which can re-order the bindings that occur in a
`relax.DataflowBlock`. This is not intended for use in a
general-purpose optimization pipeline, but instead as a utility that
may be used as needed in specific cases. For example, normalization
of unit tests that should not depend on the order of variable binding.
* Update docstring according to review comment
* [Transform] Implement relax.transform.ReorderPermuteDimsAfterConcat
This commit implements an optional optimization pass
`relax.transform.ReorderPermuteDimsAfterConcat`, which reorder
expressions of the form `R.concat(R.permute_dims(A),
R.permute_dims(B))` into `R.permute_dims(R.concat(A,B))`.
This pass is intended to be used alongside `CombineParallelMatmul`.
After parallel matmuls are combined, to be lifted out, and optimized
`nn.Linear` kernels to find the `R.matmul(x, R.permute_dims(weights))`
patterns they are looking for.
```python
@R.function
def func(x: R.Tensor, weight_query: R.Tensor, weight_key: R.Tensor, weight_value: R.Tensor):
"""Initial IRModule
The `R.permute_dims` followed by `R.matmul` is the relax
equivalent of `nn.Linear`, and will frequently have optimized
kernels.
"""
weight_query_T = R.permute_dims(weight_query)
query = R.matmul(x, weight_query)
weight_key_T = R.permute_dims(weight_key)
key = R.matmul(x, weight_key)
weight_value_T = R.permute_dims(weight_value)
value = R.matmul(x, weight_value)
@R.function
def func(x: R.Tensor, weight_query: R.Tensor, weight_key: R.Tensor, weight_value: R.Tensor):
"""After `CombineParallelMatmul`
There's now only a single matmul to be performed, which is
generally better than performing three small matmuls. However,
the optimized kernels for `nn.Linear` can no longer be applied,
because the `R.concat` isn't part of the expected pattern.
"""
weight_query_T = R.permute_dims(weight_query)
weight_key_T = R.permute_dims(weight_key)
weight_value_T = R.permute_dims(weight_value)
fused_weight_T = R.concat([weight_query_T, weight_key_T, weight_value_T], axis=1)
fused_qkv = R.matmul(x, fused_weight_T)
query, key, value = R.split(fused_qkv)
@R.function
def func(x: R.Tensor, weight_query: R.Tensor, weight_key: R.Tensor, weight_value: R.Tensor):
"""After `ReorderPermuteDimsAfterConcat`
There's still only a single matmul, and the optimized kernels for
`nn.Linear` can be applied again.
"""
fused_weight = R.concat([weight_query, weight_key, weight_value], axis=0)
fused_weight_T = R.permute_dims(fused_weight)
fused_qkv = R.matmul(x, fused_weight_T)
query, key, value = R.split(fused_qkv)
```
* Expand description of `max_concat` variable as a temporary solution
* [IR] Add utility methods to IRModule
* `IRModule.clone`: Clone the module. While in C++, a module can be
copied using `IRModule::CopyOnWrite()`, copying a module in Python
required passing all members into the `IRModule` initializer. The
`IRModule.clone` method provides an easier way to copy an `IRModule`
from python.
* `IRModule.__delitem__`: Remove a function from the module. This
exposes the C++ method `IRModuleNode::Remove` for use in the python
API. This uses the python `del` keyword, similar to a native python
list. Similar to the existing `IRModule.__getitem__`, this can be
called with either a `GlobalVar` or a python string.
* `IRModule.__contains__`: Check if a function is in the module. This
allows the pythone keyword `in` to check if a module contains a
specific function. Similar to the existing `IRModule.__getitem__`,
this can be called either with a `GlobalVar` (`if gvar in mod`) or
with a python string (`if "function_name" in mod`).
* [Unity][CodeGen] RunCodegen based on externally-exposed functions
Prior to this commit, `relax.transform.RunCodegen` required a list of
entry functions for a module, defaulting to `"main"` if not specified.
The list of entry functions is duplicate information that could be
inferred from the module, and should not be required from the user.
This commit updates `RunCodegen` to treat all externally-exposed
functions as entry points, in the same manner as
`DeadCodeElimination`.
For backwards compatibility, the `entry_functions` argument is still
accepted, and is used to augment the list of externally-exposed
functions.
If `R.matmul(x, R.take(weights, indices))` occurs, with `R.take`
selecting along the output feature dimension, it can be
rearranged to `R.take(R.matmul(x, weights), indices)`.
* 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] Add utility for collecting compile-time bindings
Whether an optimizations should be performed may depend on when the
variables in an expression are known.
For example, consider a LoRA-adjusted model, with base weights `W` of
shape `[m,n]`, LoRA components `A` and `B` with shapes `[r,n]` and
`[m,r]` respectively, and activations `x` with shape `[n,1]`. The
LoRA-adjusted matmul could be computed either as `(W + B*A)*x` or as
`(W*x + B*(A*x))`.
If `A` and `B` are provided at run-time, then computing `(W +
B*(A*x))` requires significantly fewer computations.
* `(W + B*A)*x`: `m*n*(2*r + 3)` operations
1. `B*A`: `2*m*n*r` operations using a naive matmul
2. Adding `W` to (1): `m*n` operations
3. Multiplying `x` by (2): `2*m*n` operations
* `(W*x + B*(A*x))`: (2*m*n + r*(2*n + 2*m + 1))
1. `W*x`: `2*m*n` operations
2. `A*x`: `2*r*n` operations
3. Multiplying `B` by (2): `2*m*r` operations
4. Adding (1) and (3)`: `m` operations
However, if `A` and `B` are known at compile-time, then computing `(W
+ B*A)*x` groups all compile-time values together, allowing them to be
computed earlier (i.e. using `LiftTransformParams`)
* `(W + B*A)*x`: `2*m*n` operations
1. `B*A`: 0 operations, computed at compile-time
2. Adding `W` to (1): 0 operations, computed at compile-time
3. Multiplying `x` by (2): `2*m*n` operations
Since the choice of optimized expression depends on which parameters
can be computed at compile-time, it is useful to have a utility that
identifies values that can be computed at compile-time.
* [Unity] QoL improvements for Dataflow matching
- Update the zero-parameter `WildcardPattern` constructor to produce a
valid instance. Previously, the zero-parameter constructor produced
a null instance of `WildcardPattern`, which resulted in an error
when used. The `WildcardPattern` was expected to be constructed
through the `Wildcard` function instead. Since all other
`DFPattern` child classes could be constructed explicitly, this
could lead to unexpected outcomes.
- Check for `pattern.defined()` when performing a pattern-match. If
a null instance of a pattern is provided, this gives an error
message with more context than the one raised by `DFPatternFunctor`.
- Expose `RewriteCall` for use in C++. Previously, it had only been
exposed through the FFI registry, and had no declaration in a header
file.
* [Unity][Transform] Implement relax.transform.AdjustMatmulOrder
Reorder `x*(A*B)` to `(x*A)*B`. Intended for optimization of LoRA
models, for which `(x*A)*B` has a much smaller memory footprint.
* Fix copy-paste error
* Check for re-orderings from the LHS, skip if cannot prove a benefit
* [Unity][Transform] Implement UpdateParamStructInfo
Provide a convenience method to update parameter struct info,
propagating any changes to internal bindings and return type.
* lint fix
* Update implementation to update params in relax::Function mutator
* Implement a pass that extracts dataflow blocks from consecutive pure operations in binding blocks. Also contains bugfixes for canonicalize bindings
* Trailing whitespace
* current_block_ in CanonicalizeBindings should be optional, also lets us get rid of inside_dataflow_
* Add check that current_block_ is not set before we update it
* Also reset current_block_ for nested SeqExprs (like in If branches)
* Rename pass to ConvertToDataflow, reduce default min size to 2
* Handle case of dataflow bindings coming *after* a dataflow block
* Add additional test cases
* Add test cases for pure and impure inner and outer functions
* Add more non-call test cases
* Rebase fixes
---------
Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>
* [Unity][Transform] Implement Relax function inlining
This commit adds the ability to inline one Relax function into
another. This is provided with two APIs, one which allows explicit
specification of the functions to be inlined, and one which inlines
every private Relax function in an `IRModule`. Function inlining with
explicit replacements is intended for use by an end-user when building
a model function (e.g. inlining of a rotary embedding), and the second
is intended for use as part of a generic optimization pipeline.
* lint fix
* Use DetectRecursion from relax/analysis.h
* Added test case for error on mutually-recursive functions
* [Unity] Implement RemoveUnusedParameters transform
Currently, the `FuseOps` and `FuseTIR` passes have a large amount of
added complexity to identify and handle partial use of tuple
arguments. The handling partial use of tuples could be significantly
simpler if performed in multiple steps.
1. Perform `FuseOps`. Any tuple variables that are used by the fused
function are passed as-is.
2. Expand any parameters that are passed as a tuple. Any unused
tensors that were included in a partially-used tuple will be
converted to unused parameters.
3. Remove any unused parameters. Any unused tensors that were
included in a partially-used tuple will be removed in this
step.
4. Perform `FuseTIR`. No checking for tuple arguments, either partial
or full, is required at this step.
This PR implements `relax.transform.RemoveUnusedParameters`, which is
step (3) in this process.
* Update based on review comments
[Unity] Implement ExpandTupleArguments transform
Currently, the `FuseOps` and `FuseTIR` passes have a large amount of
added complexity to identify and handle partial use of tuple
arguments. The handling partial use of tuples could be significantly
simpler if performed in multiple steps.
1. Perform `FuseOps`. Any tuple variables that are used by the fused
function are passed as-is.
2. Expand any parameters that are passed as a tuple. Any unused
tensors that were included in a partially-used tuple will be
converted to unused parameters.
3. Remove any unused parameters. Any unused tensors that were
included in a partially-used tuple will be removed in this
step.
4. Perform `FuseTIR`. No checking for tuple arguments, either partial
or full, is required at this step.
This PR implements `relax.transform.ExpandTupleArguments`, which is
step (2) in this process.
[Unity] Implement RemoveUnusedOutputs transform
This commit implements `relax.transform.RemoveUnusedOutputs`, an
extension of the current `DeadCodeElimination` pass. If an internal
callee produces multiple outputs, but only some of those outputs are
used, the callee can be rewritten to only produce the outputs that are
used. This is intended to allow more aggressive DCE in the callee, as
reducing the set of outputs can result in intermediate computations
becoming unused.
* enable attrs_getter
* update test case for attrs_getter
* add name checker for test tensorrt
* remove pattern.py
* add unique name for tuple
* change attrs_getter argument
* change attrs getter signature
Tianqi Chen