49 Commits

Author SHA1 Message Date
Tianqi Chen c717c5b217 [IR][Relax][TIRx] Unify Var identity (#20004) 2026-07-15 17:12:40 +08:00
Tianqi Chen 275114b327 [REFACTOR][IR] Unify PrimExpr with Expr typed view (#19910)
## 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.
2026-07-01 18:55:33 -04:00
Tianqi Chen 120812e9ac [REFACTOR][Relax] Phase out PrimValue and Relax expression wrappers (#19891)
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.
2026-06-26 07:18:04 -04:00
Tianqi Chen f3f5a3e42a [REFACTOR][RELAX] Rename Relax base type to AnyType (#19889)
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>
2026-06-25 13:29:10 -04:00
Tianqi Chen 0082836d2d [REFACTOR][RELAX] Phase out Relax PrimType (#19858)
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.
2026-06-22 11:27:20 -04:00
Tianqi Chen 1bb5cf6102 [REFACTOR][IR] Unify StructInfo and Type (#19853)
## 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
2026-06-21 10:12:12 -04:00
Tianqi Chen d883f5064f [REFACTOR] Remove runtime/object.py shim and route Object via tvm_ffi (#19440)
## Summary

TVM-side cleanup that drops the `python/tvm/runtime/object.py` shim and
routes `tvm.runtime.Object` directly to `tvm_ffi.Object`. The
`tvm.runtime.Object` re-export is preserved (now a re-export of
`tvm_ffi.Object`) so external callers keep working.

The load-bearing `__object_repr__` install — which wires TVM IR objects
up to the rich C++ `ReprPrinter` registered through
`init_ffi_api("node", ...)` — moves into
`python/tvm/runtime/_ffi_node_api.py`.
That module is already imported as a side-effect-only module from
`python/tvm/runtime/__init__.py`, so the override fires at the right
time (after `init_ffi_api` registers the C++ printer).

`_ffi_node_api.AsRepr` itself is **kept**: `tvm_ffi`'s default repr is
primitive (`ClassName(ptr)`); TVM IR objects need the rich printer
registered via `init_ffi_api("node", ...)`. `AsRepr` is what bridges
that printer back into Python `repr(obj)` and is also the runtime-only
fallback when `libtvm.so` is unavailable.

The 7 in-tree importers of the deleted shim (plus one straggler in
`runtime/disco/session.py`) are switched to either
`from tvm.runtime import Object` or `from tvm_ffi import Object`,
depending on which pattern the file already uses.

## Test plan

- [x] `python -c "import tvm; print(repr(tvm.IRModule({})))"` produces
  TVMScript-style output (rich repr preserved).
- [x] `pytest tests/python/all-platform-minimal-test/ -x` — 75 passed,
  77 skipped (matches baseline).
- [x] `pytest tests/python/tirx-base/ -x` — 273 passed, 2 skipped.
- [x] `pre-commit run --files <changed files>` — all hooks pass.
- [ ] CI green.
2026-04-25 12:20:01 -04:00
Tianqi Chen 141c22fd8a [Refactor] Bring up tirx namespace (#18913)
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
2026-03-19 21:27:54 -07:00
Tianqi Chen 9a8320acbd [LINT][PYTHON] Modernize annotations with ruff UP rules (#18830)
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.
2026-02-27 21:29:47 -05:00
Tianqi Chen 33dcea1686 [REFACTOR][LINT] Modernize ruff config (#18810)
This PR removes the extra lint violations from the codebase so lint
aligns with the latest style
2026-02-23 07:29:21 -05:00
Tianqi Chen aa2e609136 [LINT] Modernize lint to use pre-commit hooks (#18807)
This PR migrates existing lint to use pre-commit hooks
2026-02-22 11:03:21 -05:00
Guan-Ming (Wesley) Chiu 99c1d9245f [Relax] Remove obsolete TODO comments (#18649)
- Remove TODO about `ctx->ReportWarning()` since LOG(WARNING) is the
standard pattern
- Remove TODO about skipping type params/constraints since the current
FuncType implementation only contains arg_types and ret_type fields.
`type_params` and `type_constraints` don't exist in the codebase.
- Fixed typo in warning message
2026-01-11 21:45:07 +08:00
Tianqi Chen 543e64dbb1 [FFI][REFACTOR] Cleanup tvm_ffi python API and types (#18277)
This PR cleans up the python API to make things more consistent
with existing python array api and torch.

Device update
- device_id => index, to be consistent with torch
- device_type => dlpack_device_type() returns int
- added type property same as torch.device

API updates:

- Move the convenient method like cpu() out into tvm runtime to keep device minimal
- tvm_ffi._init_api => tvm_ffi.init_ffi_api
- tvm_ffi.register_func => tvm_ffi.register_global_func
2025-09-07 10:38:50 -04:00
Tianqi Chen 3c36ce2ec6 [FFI][REFACTOR][ABI] Rename NDArray to Tensor (#18275)
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.
2025-09-06 14:33:59 -07:00
Tianqi Chen a7a0168be5 [FFI][REFACTOR] Establish tvm_ffi python module (#18226)
* [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
2025-08-24 15:46:20 -07:00
Tianqi Chen 10e66934ec [REFACTOR] Phase out the RelaxExpr.checked_type in favor of struct_info (#18078) 2025-06-18 16:39:06 -04:00
Tianqi Chen 4289efa0d5 [REFACTOR][PYTHON] Phase out tvm._ffi and Limited API support (#18020)
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.
2025-05-28 16:52:36 -04:00
Tianqi Chen c81fccaa2f [REFACTOR] Phase out te.Schedule c++ components (#17662)
* cleanup schedule c++

* remove vitis ai

* remove VERILATOR

* remove aocl and sdaccel

* remove opengl

* remove microdev and antlr

* remove frontends

* fix

* Cleanup relay related legacy components

* fix

---------

Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
2025-02-17 17:03:09 -05:00
Tianqi Chen ccaa534b2c [REFACTOR] Phase out relay python components (#17656)
This PR starts the step 0 to phase out relay from the current
development main branch.  This PR focuses on the python
components of relay, autotvm, auto_scheduler. To make the change
manageable, we will also do followup steps on te.Schedule and
c++ components in followup PRs.

To continue support community members who depends on
legacy flows, the [v0.19.0](https://github.com/apache/tvm/tree/v0.19.0)
branch will continue contain these components.


As noted in [discussion on phasing out legacy components](https://discuss.tvm.apache.org/t/phasing-out-legacy-components/17703/30),
this would help us to do two purposes:

- By removing outdated or redundant elements, we can significantly
reduce complexity and improve maintainability.
- Unify our focus: Concentrating our efforts on the new unity flow
will allow for more efficient development and innovation.

It is also a good opportunity for us to revisit and reduce CI time.
The past relay legacy flow contains a lot of end to end tests that
requires hardware resources to run and causing long CI time.
Moving onwards, we can focus more on unit-tests that focuses
on structural equality and runs within seconds, while be mindful
about tests that requires hardware resources (by restricting them
to specific folders and CI nightly in some cases).

---

Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
2025-02-15 13:48:28 -05:00
Eric Lunderberg 02f48828e4 [FFI] Re-introduce the boxed primitive values (#17257)
* Revert "Revert "[FFI][RUNTIME] Introduce runtime boxed types for int/float/bool" (#17252)"

This reverts commit 11be832620.

* [FFI] Re-introduce the boxed primitive values

Initially introduced in https://github.com/apache/tvm/pull/16183,
these changes were reverted in
https://github.com/apache/tvm/pull/17252 due to performance
degredation in some Relax models.  This could occur when a model
contained a large number of calls to `"vm.builtin.tuple_getitem"`,
which may occur when model weights are provided as a tuple.

This PR re-applies the changes from
https://github.com/apache/tvm/pull/16183, but with the performance
degredation resolved.  The root cause was unnecessary type-checking
when converting from an untyped `tvm::ArrayNode*` to the typed
`tvm::Array<T>`, in the case where `T` is `ObjectRef`.

* Correct typo from T to U
2024-08-12 08:36:17 -04:00
Tianqi Chen 11be832620 Revert "[FFI][RUNTIME] Introduce runtime boxed types for int/float/bool" (#17252)
Revert "[FFI][RUNTIME] Introduce runtime boxed types for int/float/bool (#16183)"

This reverts commit 5f22be4d83.
2024-08-07 12:19:13 -04:00
Eric Lunderberg 5f22be4d83 [FFI][RUNTIME] Introduce runtime boxed types for int/float/bool (#16183)
* [Container] Support non-nullable types in Array::Map

Prior to this commit, the `Array::Map` member function could only be
applied to nullable object types.  This was due to the internal use of
`U()` as the default value for initializing the output `ArrayNode`, where
`U` is the return type of the mapping function.  This default
constructor is only available for nullable types, and would result in
a compile-time failure for non-nullable types.

This commit replaces `U()` with `ObjectRef()` in `Array::Map`,
removing this limitation.  Since all items in the output array are
overwritten before returning to the calling scope, initializing the
output array with `ObjectRef()` does not violate type safety.

* [FFI] Separate runtime types from IR types for int/float/bool

Prior to this commit, `int`, `float`, and `bool` arguments from Python
were converted to `IntImm`, `FloatImm`, and `Bool`.  These are
subtypes of `PrimExpr`, and should only be used at compile-time.  By
automatically applying this conversion as part of the FFI, these types
are required to be present whenever a primitive is converted to a
`tvm::ObjectRef`.

This can become especially fragile for an end-user when storing
objects into a TVM container.  Because TVM containers require all
contents to be `ObjectRef` subclasses, an automatic conversion may be
applied on storing into a container, resulting in an unexpected type
being retrieved from the container.  For example, this currently
occurs in Relax when extracting a `R.Prim` from a `R.Tuple`.

This commit introduces a `Box<T>` type for storage of boxed primitives
at runtime, distinct from the IR types.

* Primitive arguments provided to a PackedFunc that requires an
  `ObjectRef` will be converted to the corresponding boxed type.
  (e.g. Passing a Python `int` to a C++ function accepting `ObjectRef`
  produces a `Box<int64_t>`.

* Boxed primitives provided to a PackedFunc that requires an unboxed
  primitive will be converted to the corresponding primitive.

* PackedFunc return values of `ObjectRef` are converted to the
  corresponding primitive, if present.  (e.g. If a `tuple_getitem`
  with static return type `ObjectRef` returns a `Box<int64_t>`, it
  will be unwrapped to a python `int`.)

Together, these three rules provide backwards compatibility for
existing PackedFunc definitions, while avoiding exposing the user to
any container-induced type conversions betweeen primitive types and
`ObjectRef`.

* Fix unit test failure after merge

* Fix breakage in new unit test
2024-08-05 09:19:20 -04:00
Wuwei Lin 5d5edd2fd8 [Relax] Integrate cuDNN attention (#17157)
* [Relax] Integrate cuDNN attention

* update cmake

* lint

* lint

* cudnn frontend

* lint

* lint

* fix test

* skip test
2024-07-22 12:36:06 -07:00
Wuwei Lin 63f9cd6523 [Relax] Alloc BYOC workspace with R.builtin.alloc_tensor (#17110)
* [Relax] Alloc BYOC workspace with R.builtin.alloc_tensor

This makes the allocation go through memory planning and make it
compatible with cuda graph.

* lint

* lint
2024-06-26 08:43:12 -07:00
Steven S. Lyubomirsky a763b22119 [Unity][Transform] Replace eligible operators with in-place versions in dataflow blocks (#16129)
* 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>
2024-01-17 17:05:54 -05:00
Ivan Sidorenko bc2415bcc8 [Unity][BYOC] Support IGEMM in cuBLASLt (#16134)
This commit adds support of i8i8i32 matmul and its offloading to cuBLAS.
2023-11-16 04:58:59 +09:00
Eric Lunderberg e506bffa72 [Unity] Implement FNormalize attribute for operators (#16067)
* [Unity] Implement FNormalize attribute for operators

Some Relax operators have requirements regarding their AST that are
stronger than are checked by the C++ types being used.  These are
similar to checks that are present in the `tvm::relax::WellFormed`
utility, such as checks forbidding the use of undefined variables,
which are also stronger than required by the underlying C++ types.
However, because every operator may have unique requirements, it would
be unreasonable to expect a writer of a `relax::ExprMutator` to be
aware of and to maintain all such requirements.

This PR introduces an operation operator attribute `FNormalize`.  If
defined, this function is used to apply an operator-specific
normalization.

* If no change is required, `FNormalize` should return the input
  argument unmodified.

* `FNormalize` is only responsible for normalization of the operator
  itself.  The expression it returns may be unnormalized (e.g. contain
  nested expressions).

* `FNormalize` receives the `BlockBuilder` as an argument, to allow
  context-dependent normalization.

  For example, an operator whose normalization requires in-line
  expressions may use `BlockBuilder::LookupBinding` to perform
  variable replacement.

* `FNormalize` is applied after `FInferStructInfo`.  `FNormalize` may
  assume that the `relax::Call` passed to `FNormalize` has
  well-defined struct info.

  * Corollary: `FInferStructInfo` may not assume that its
    `relax::Call` argument has been passed through `FNormalize`.

    This is a reasonable requirement, because (1) shape inference
    should depend only on the struct info of arguments and not the
    values themselves, and (2) this only impacts operators that use
    `FNormalize`.

* `FNormalize` should not be used to apply simplifications, and should
  be limited to cases where the same computation may be expressed in
  multiple manners.

  For example, replacing a by-variable tuple with an in-line tuple in
  `R.call_tir` is a form of normalization, but replacing `R.add(arg,
  R.const(0))` with `arg` is a form of simplification.

  This separation is to ensure that `FNormalize` has minimal overhead,
  as some simplifications may have large computational costs, and
  `FNormalize` is applied as part of all `ExprMutator` usage.  A later
  PR will introduce an attribute `FSimplify`, along with a dedicated
  pass to apply simplifications.

* Use of `FNormalize` is suppressed while parsing TVMScript.
  TVMScript must be able to generate test cases that trigger specific
  failure modes, and that may include producing un-normalized relax
  IR.  In addition, TVMScript must be stable when passed through a
  round-trip from IR to text to IR.

* Disable C++ lint on explicit zero-parameter constructor

* Avoid double-lookup with map.count(op) then map[op]
2023-11-06 19:03:04 -05:00
Eric Lunderberg 49a3a51ff1 [Unity] Ensure one VM register for each relax binding (#15855)
* [Unity] Ensure one VM register for each relax binding

Prior to this commit, if a relax variable were assigned to itself,
through either `VarBinding` or `MatchCast` nodes, the two relax
variables would share the same register in the VM.  As a result, any
upstream transform that deletes an object with `R.memory.kill_tensor`
or `R.memory.kill_storage` must be aware of this VM behavior, and to
only output one such instruction for each set of aliased registers.

This commit updates the VM to produce one register for each aliased
relax variable.  The trivial bindings can be removed applying
`relax.transform.CanonicalizeBindings`, instead of being implicitly
de-duplicated at the codegen level.

This PR is a follow-up to https://github.com/apache/tvm/pull/15854,
with a better long-term solution, but which may have knock-on effects
that must also be resolved.  In addition, this adds a usage example
for the bug reported in https://github.com/apache/tvm/pull/15852, to
avoid re-occurrence of similar issues.

* Add unit test to validate that the alias is preserved

* Use the new GetBoundValue utility function

* Update VMTIRCodeGen to also avoid de-duplication of bindings

* Move the callback definition to tvm.relax.testing.vm namespace
2023-10-31 11:00:39 -05:00
Junru Shao 7486476b6f [Unity] Deterministic Ordering when Iterating IRModule::functions (#16020)
Prior to this PR, visiting order of the following for-loop is
non-deterministic:

```python
mod: tvm.ir.IRModule
for gv, func in mod.functions.items():
  ...
```

This is because `IRModule` stores those functions inside a hash map
`Map<GlobalVar, BaseFunc>`, which is based on pointer equality.
This behavior is usually innocent and harmless given many previous
workloads only have one "main" function as the primary entry point, e.g.
a single Relax Function and a bunch of TIR functions. However, it is not
true in LLM usecases where `prefill`, `decoding` are both equally
important entrypoints, and in those cases, it is possible that the
names of generated TIR functions from LegalizeOps are essentially
different from each run, which is, again, harmless in basic usecases,
but it does make debugging more challenging.

This PR corrects this behavior by sorting the functions alphabetically
by their names.
2023-10-31 08:51:25 -04:00
Anirudh Sundar Subramaniam d2972f3b42 [Unity] Add instruments to relay translator (#15601)
* [Unity] Add instruments to relay translator

Sometimes its useful to instrument relay passes that are run during
relay to relax translation, and this patch adds a new argument to relay
translator to accept an instruments list argument that gets passed onto
the PassContext used while running relay prefix passes

* Fix test case
2023-08-27 10:21:14 -05:00
Eric Lunderberg 1a1af0c700 [Unity][NN] Allow nn.Module to generate subroutines (#15453)
* [Relax][IRBuilder] Allow subroutine definition in BlockBuilder

Prior to this commit, a user needed to first define any subroutines
used within a module, then define the functions that use those
subroutines.  This prevents some use cases, such as
`tvm.relax.testing.nn.Module` implementations that generate a
subroutine when required.

This commit updates the `tvm.relax.block_builder` to maintain a stack
of functions currently being defined, such that subroutines can be
defined as needed, resuming collection into the main function
afterwards.

* [Unity][NN] Allow nn.Module to generate subroutines

Prior to this commit, `nn.Module` generated the relax expression
within the body of `relax.BlockBuilder.current()`.  For large models,
this can result in extremely large function bodies, making it
difficult to identify regions within the module.

This commit adds an option `nn.Module.define_subroutine`, which
defaults to `False` (current behavior).  If set to `True`; either
within a subclass, within an instance, or globally in `nn.Module`;
function calls into the module will produce a subroutine representing
the module's execution, and a call into that subroutine.  For example,
calling a `class Linear(nn.Module)` would produce a `def linear(arg,
weights)` function definition and a `module.linear(arg, weights)`
function call.

To ensure correct shape propagation, a subroutine is generated for
each unique set of argument shapes passed to the `nn.Module` subclass.

* Backwards compatibility workarounds

If a `def forward` expected a python-style tuple object, it may now
receive a `relax.Tuple`, or a `relax.Var` annotated with
`relax.TupleStructInfo`.  This commit adds backwards compatibility
type checks to allow common python operations to keep the same
result in these cases.

* Update accessor of _blocks

* [Relax][NN] Implemented subroutines for relax.frontend.nn

* [Bugfix][Relax] In-scope late-bound parameters

Prior to this commit, the normalization of a function body was
performed with only the early-binding parameters in-scope, but not the
late-binding parameters.  Therefore, if a parameter has a dynamic
shape parameter that determines a dynamic output shape, the output
shape would only be inferred correctly if the parameter is provided as
part of `block_builder.function(name, params=[...])`.  If the
parameter is instead provided as part of
`block_builder.emit_func_output(output, params=[...])`, the output
would be erroneously annotated with `R.Tensor(ndims=ndims)` instead of
`R.Tensor(shape=dynamic_shape)`.

This commit updates `BlockBuilder.emit_func_output` to ensure the
parameters are in-scope, regardless of when they are provided.

* Updated implementation in relax.testing.nn

* Lint fixes
2023-08-07 10:03:06 -05:00
Yixin Dong 98ef29a852 [Unity][Training] Enhance gradient system (#15230)
This PR adds these features to the gradient system:
- Checkpointing for gradient pass
    - `tvm.relax.testing.nn.checkpoint`
    - `tvm.relax.op.grad.start_checkpoint` and `tvm.relax.op.grad.start_checkpoint`
    - Support in the Gradient pass
- Fix several minor problems in op_gradient
2023-07-18 14:52:56 -07:00
Eric Lunderberg 0328e98b94 [Unity][NN] Allow nn.Placeholder/Parameter prior to BlockBuilder (#15025)
* [Unity][NN] Allow nn.Placeholder/Parameter prior to BlockBuilder

Prior to this commit, use of `nn.Placeholder` or `nn.Parameter`
outside of a `with block_builder.function('name'):` scope resulted in
an error.  This commit updates the behavior to allow declaration prior
to entering the `with` block.  This can be useful for declaring a
model object, which is then used to define several related functions.

The scope was required so that `relax.BlockBuilder.current()` could
de-duplicate variable names.  While two distinct variables in Relax
may have identical names, for user readability it is convenient to
have all names be unique within a Relax function.  This commit
maintains the de-duplication of names if a `nn.Placeholder` or
`nn.Parameter` is defined within an active `relax.BlockBuilder`, that
context may be used to provide a unique name.

* Lint fix
2023-06-09 10:33:15 +05:30
Eric Lunderberg 3dd9345024 [Unity] Allow name_hint in additional locations (#15027)
* [Unity] Allow name_hint argument to BlockBuilder.emit_te

* [Unity][NN] Pass name_hint from nn.emit to BlockBuilder.emit
2023-06-06 07:56:03 -05:00
Steven S. Lyubomirsky 7aff6bfbb0 [Unity][IR] Purity Tracking (#14394)
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
2023-05-18 15:20:15 -04:00
Lite Ye 0c9aa58907 [Unity][BYOC] Add shape validation for bias arg in cuBLAS (#14809)
Add shape check for bias arg in cuBLAS backend

And use the correct type key for cublas json runtime

Fix lint
2023-05-10 09:50:12 +09:00
Junru Shao d837fd0738 [MERGE] Fix lint after lint image upgrade 2023-05-07 21:00:22 -04:00
Tianqi Chen ee53619661 [Unity] Improve and reduces possible memory leak RPC debug (#14662)
This PR improves and reduces possible memory leak
during lib comparator and rpc debug.
2023-04-18 17:37:34 -04:00
Ruihang Lai bcad5a9abe [Unity][VM] LibComparator using dtype from input (#14623)
Previously the LibComparator omitted the dtype argument when creating
a NDArray, which leads to the dtype always float32. This PR fixes this
behavior.
2023-04-14 07:51:01 -04:00
Ruihang Lai 6bcd375369 [Unity] Torch-like NN module enhancement (#14499)
This PR
* supports `emit` to the Relax nn module so that we can emit arbitrary
Expr instead of only having `emit_te` available,
* fixes module parameter fetching to ignore non-parameter fields of an
module (the previous behavior is to throw error for unrecognized
fields).
2023-04-04 16:16:22 -04:00
masahi e54e04d520 [Unity][BYOC] Add cuBLAS backend (#14291)
* stub

* fixed build

* test stub

* basic gemm working

* transposed gemm work

* wip

* bias and epilogue work

* support fp16 and transposed bias

* support batched gemm

* clean up

* access arguments properly

* expose ExtractArgIdx to python and use it in cutlass byoc

* put matmul ir into common testing file

* updated for the latest rev

* pylint
2023-04-04 11:00:32 +09:00
Anirudh Sundar Subramaniam 20cc2ff300 [Unity] Add support to append relay op attrs in translator (#14356)
This PR adds an argument to the relay to relax translator to append the
relay Op Attrs as function attributes to the generated TIR PrimFuncs

This information is really useful in Relax as the absence of this
prevents us from being able schedule efficiently for ops that are
heavily sensitive to the attributes.

For example, the `groups` attribute to `conv2d` op is needed to
differentiate between regular conv2d and depthwise conv2d.
2023-04-01 15:31:44 -04:00
Tianqi Chen 67e183c33f [Unity][DEBUG] Add Instrument (#14302)
This PR adds an instrumentation option to the relax VM.
The instrument will be called before/after each call
instruction if specified.

We also include a testing utility that leverages uses
instrument. LibCompareVMInstrument leverages the instrument
to compare implementations on another backend.

Also updated a few places in web runtime to improve debugging.
2023-04-01 15:31:44 -04:00
ibsidorenko 71c23b7a11 [Unity] Consider target context for Relay to Relax conversion (#14269)
Some of Relay passes are target specific (for example, AlterOpLayout).
This commit adds target information for Relay -> Relax conversion.
2023-04-01 15:31:44 -04:00
Steven S. Lyubomirsky f17a7285bb [Unity][Debugging] AST printer (#14152)
This PR transfers over the AST printer from tlc-pack/relax. The AST printer is a debugging tool that prints out a Relax AST in a precise and human-readable format, which can be helpful for debugging the parser or various passes.

Co-authored-by: Yuchen Jin <yuchenj@cs.washington.edu>
Co-authored-by: Lesheng Jin <34279105+LeshengJin@users.noreply.github.com>
Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-authored-by: Ruihang Lai <ruihangl@cd.cmu.edu>
Co-authored-by: Tianqi Chen <tqchen@users.noreply.github.com>
2023-04-01 15:31:37 -04:00
Yuchen Jin b792751aaa [Unity] Relay -> Relax translator (#14026)
This PR implements a Relay to Relax translator, which allows us to import Relay workloads to Relax for benchmarking and development purposes (tests and examples are added).
2023-04-01 15:31:37 -04:00
Yuchen Jin b3d913a737 [Unity] e2e Relax minimum build flow (#13961)
This PR introduces the e2e Relax lowering flow (`relax.vm.build`). Tests for each pass in the flow are added.

Co-Authored-by: Altan Haan <altanh@cs.washington.edu>
Co-Authored-by: Andrew Liu <andrewlliu@gmail.com>
Co-Authored-by: Hongyi Jin <3231950289@qq.com>
Co-Authored-by: Jiawei Liu <jaway.liu@gmail.com>
Co-Authored-by: Junru Shao <junrushao1994@gmail.com>
Co-Authored-by: Prakalp Srivastava <prakalp@octoml.ai>
Co-Authored-by: Ruihang Lai <ruihangl@cs.cmu.edu>
Co-Authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-Authored-by: Steven S. <Lyubomirsky slyubomirsky@octoml.ai>
Co-Authored-by: Sunghyun Park <49998730+sunggg@users.noreply.github.com>
Co-Authored-by: Tianqi Chen <tianqi.tchen@gmail.com>
Co-Authored-by: Yong Wu <yongcale@gmail.com>
Co-Authored-by: Ziheng Jiang <ziheng@apache.org>
2023-04-01 15:31:36 -04:00
Yuchen Jin 574084d340 [Unity] Relax VM codegen (#13954) 2023-04-01 15:31:36 -04:00
Yuchen Jin 70d80cc126 [Unity] Relax VM (#13878)
This PR implements a flexible register-based VM to execute relax programs with dynamic shape and control flow. Design: https://github.com/tlc-pack/relax/wiki/Relax-VM-Design.

Co-Authored-by: Ziheng Jiang <ziheng@apache.org>
Co-Authored-by: Ruihang Lai <ruihangl@cs.cmu.edu>
Co-Authored-by: Sunghyun Park <49998730+sunggg@users.noreply.github.com>
Co-Authored-by: Junru Shao <junrushao1994@gmail.com>
Co-Authored-by: Prakalp Srivastava <prakalp@octoml.ai>
Co-Authored-by: Yong Wu <yongcale@gmail.com>
Co-Authored-by: Steven S. Lyubomirsky <slyubomirsky@octoml.ai>
Co-Authored-by: Tianqi Chen <tianqi.tchen@gmail.com>
Co-Authored-by: Hongyi Jin <3231950289@qq.com>
2023-04-01 15:31:31 -04:00