Cleanup pass that relies on IR invariants instead of re-checking
already-guaranteed conditions. No new features; this is a
consolidation/cleanup pass only.
## Changes
- **docsifier (`python_doc_printer.cc`)**: the `ExprStringDoc` escape
scope always wraps the printer's fixed in-memory `ostringstream` sink,
which never short-writes and never enters a fail state. Drop the
streambuf-general short-write reporting in `xsputn`, the ctor `good()`
ICHECK, the dtor `rdstate`/`setstate` dance, and the redundant
post-render `good()` ICHECK; keep the one-line `saw_newline()` contract.
- **relax diagnostics (`well_formed.cc`, `block_builder.cc`)**: the ty
diagnostics test `ty.IsMissing()` on a now non-nullable `Type`, so word
them as "is missing" rather than "is nullptr".
- **relax numeric-gradient tests**: derive the device from the build
target via `tvm.device_from_target` inside the helpers instead of
threading a redundant `dev` argument that duplicates `target` at every
call site; annotate the numpy inputs as `np.ndarray`.
- **target/printer tests**: drop assertions that re-check a condition an
earlier assertion in the same test already guarantees.
## Summary
Compiler Targets can carry device-type semantics that runtime
device-name parsing does not preserve.
- add `tvm.device_from_target` for canonical Target-to-Device
translation
- use explicit runtime constructors where the device kind is fixed
- update target-derived utilities, tests, and documentation to use the
explicit boundary
## 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
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.
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>
When processing matmul with transpose (permute_dims operator), the current gradient pass will introduce additional permute_dims ops. That's because the gradient visit all binding from last one to the first one, and matmul and permute_dims are regarded as two separate ops. E.g.
Forward is:
```
out = matmul(a, transpose(b))
```
Then backward is:
```
grad_a = matmul(grad_out, transpose(transpose(b)))
grad_b = transpose(matmul(transpose(a), grad_out))
```
This PR introduces a new pass, GradientSimplifier, and enhances the Gradient pass to simplify these patterns. The example above will be simplified to
```
grad_a = matmul(grad_out, b)
grad_b = matmul(transpose(grad_out), a)
```
Prior to this commit, `relax.transform.LegalizeOps` needed to be
called prior to `relax.build`. This commit adds `LegalizeOps` to the
lowering flow, to simplify the calling steps for an end-user. If the
`IRModule` contains no legalizable functions, a second legalization
pass has no effect.
Some test cases relied on this behavior as an implicit assertion that
operator fusion patterns applied. That is, by omitting `LegalizeOps`,
a successful compilation `relax.build` would only occur if all
legalizable operators have already been removed, and so an incorrect
fusion pattern would result in a failure to build the module. While
these tests would be better expressed by comparing against an expected
fused pattern, updating the tests is outside the scope of this PR. To
allow these tests to keep their implicit assertions, a
`"relax.transform.apply_legalize_ops"` config can be used to disable
the `LegalizeOps` pass.
### Introduction
This PR introduces the high-level reverse-mode automatic differentiation pass `Gradient` for Relax. It's the core component when we are trying training or fine-tuning in Relax IR.
Before upstreaming, this work is actively iterated and maintained in many forks like [mlc](https://github.com/mlc-ai/relax) and [relax-training](https://github.com/ACMClass-TVM-20/relax-training). Now it reaches a relatively stable version and it's time for us to upstream this important work to the unity branch.
The Python side API:
- `Gradient(func_name: str, require_grads: Optional[Union[Var, List[Var]]] = None, target_index: int = 0) -> tvm.ir.transform.Pass`
It will transform the given funcion in the IRModule, and adds a new function that calculates the gradient with regard to the function's output.
### Examples
```
@I.ir_module
class Module:
@R.function
def main(
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
# use R.sum to reduce the tensor to a scalar
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
R.output(lv2)
return lv2
After = relax.transform.Gradient("main")(Module)
```
Then the transformed module `After` will be
```
@I.ir_module
class After:
@R.function
def main(
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
R.output(lv2)
return lv2
@R.function
def main_adjoint(
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
) -> R.Tuple(
R.Tensor((), dtype="float32"),
R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")),
):
with R.dataflow():
# original bindings
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
# bindings w.r.t. intermediate variables
lv2_adjoint: R.Tensor((), dtype="float32") = R.ones((), dtype="float32")
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(
lv2_adjoint, (3, 3)
)
# bindings w.r.t. parameters
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
y_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
R.output(lv2, x_adjoint, y_adjoint)
# return value: (orig_return_values, tuple(adjoints))
return (lv2, (x_adjoint, y_adjoint))
```
Here we specify the target function `main` by its name.
We let the `require_grads` be default value (`None`) so it will calculate all inputs' adjoints (`x_adjoint`, `y_adjoint`) and return them.
We let the `target_index` be default value `0` so it will take the unique return value `lv2` as the target (the scalar we start to differentiate and propagating adjoints) of AD.
Co-authored-by: Yixin Dong <ubospica@gmail.com>
Tianqi Chen