Skip to content

Commit

Permalink
[Relax] Implement relax.transform.TopologicalSort (#16697)
Browse files Browse the repository at this point in the history
* [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
  • Loading branch information
Lunderberg authored Mar 17, 2024
1 parent 1c73491 commit fbfa926
Show file tree
Hide file tree
Showing 4 changed files with 858 additions and 0 deletions.
1 change: 1 addition & 0 deletions python/tvm/relax/transform/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@
StaticPlanBlockMemory,
ToMixedPrecision,
ToNonDataflow,
TopologicalSort,
UpdateParamStructInfo,
UpdateVDevice,
VMBuiltinLower,
Expand Down
23 changes: 23 additions & 0 deletions python/tvm/relax/transform/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,29 @@ def ToNonDataflow() -> tvm.ir.transform.Pass:
return _ffi_api.ToNonDataflow() # type: ignore


def TopologicalSort(order="depth-first", direction="from-inputs") -> tvm.ir.transform.Pass:
"""Sort bindings in relax.Dataflow blocks in the order specified
Parameters
----------
order: str
The order in which bindings should be emitted. Allowed values
are "depth-first" and "breadth-first".
direciton: str
The direction in which the sort should be performed. Allowed
values are "from-inputs" and "from-outputs".
Returns
-------
ret: tvm.ir.transform.Pass
"""
return _ffi_api.TopologicalSort(order, direction) # type: ignore


def RemovePurityChecking() -> tvm.ir.transform.Pass:
"""Activate relax.force_pure on all pure functions in the module
and unwrap all pure override ops into the normal versions.
Expand Down
Loading

0 comments on commit fbfa926

Please sign in to comment.