wehub-resource-sync
1909 lines
65 KiB
Python
1909 lines
65 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name
|
|
# ruff: noqa: E731, E741, F401, F811
|
|
"""Relax transformation passes."""
|
|
|
|
import functools
|
|
import inspect
|
|
import types
|
|
import warnings
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from typing import Optional, Union
|
|
|
|
import numpy as np # type: ignore
|
|
import tvm_ffi
|
|
from tvm_ffi import Array
|
|
|
|
import tvm.ir
|
|
from tvm.relax import Expr, Type, Var
|
|
from tvm.relax.dpl import DFPattern
|
|
from tvm.runtime import Object, Tensor
|
|
from tvm.tirx import IndexMap, PrimFunc
|
|
|
|
from ..expr import Var
|
|
from . import _ffi_api
|
|
from .legalize_ops.common import LegalizeFunc
|
|
|
|
|
|
@tvm_ffi.register_object("relax.FunctionPass")
|
|
class FunctionPass(tvm.ir.transform.Pass):
|
|
"""A pass that works on each tvm.relax.Function in a module. A function
|
|
pass class should be created through `function_pass`.
|
|
"""
|
|
|
|
|
|
@tvm_ffi.register_object("relax.DataflowBlockPass")
|
|
class DataflowBlockPass(tvm.ir.transform.Pass):
|
|
"""A pass that works on each tvm.relax.DataflowBlock in a module."""
|
|
|
|
|
|
def Gradient(
|
|
func_name: str, require_grads: Var | list[Var] | None = None, target_index: int = 0
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Reverse-mode automatic differentiation.
|
|
|
|
This pass will differentiate one function in the IRModule. Now the input function must have only
|
|
one dataflow block (ConvertToDataflow may need to be called first).
|
|
|
|
For a given function specified by `func_name`, it generates a new function with the name
|
|
`func_name + "_adjoint"`. The new function computes the gradient of the **differentiation
|
|
target** with respect to the arguments specified by `require_grads` of the original function.
|
|
|
|
If the function has only one return value, the return value will be specified as target. If the
|
|
function has more than one return values, the target will be specified as the target_index-th
|
|
return value. The target must be a scalar (0-dim tensor).
|
|
|
|
The new function will be like:
|
|
|
|
.. code-block:: python
|
|
|
|
@R.function
|
|
def main_adjoint(original_parameters):
|
|
with R.dataflow():
|
|
# the bindings of the original function
|
|
...
|
|
# calculating the gradients
|
|
...
|
|
R.output(original_outputs, grad_1, grad_2, ...)
|
|
return (original_return_value, (grad_1, grad_2, ...))
|
|
|
|
This AD pass also supports checkpointing as described in
|
|
"Training deep nets with sublinear memory cost." - Chen, Tianqi, et al. (2016).
|
|
See tvm.relax.testing.nn.checkpoint for more details.
|
|
|
|
Parameters
|
|
----------
|
|
func_name : str
|
|
The name of the specific function.
|
|
|
|
require_grads : Optional[Union[relax.Var, List[relax.Var]]]
|
|
The relax variables whose adjoints is needed. Must be parameters of the given function and
|
|
should not be duplicate. If it is not specified, adjoints of all parameters would be
|
|
computed.
|
|
|
|
target_index : int
|
|
If the specified function has more than one return values, specify the index of the return
|
|
value as the target. If it is not specified, the first return value will be the target.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
The Pass.
|
|
|
|
Examples
|
|
--------
|
|
The following code shows how to use this pass:
|
|
|
|
.. code-block:: python
|
|
|
|
@I.ir_module(s_tir=True)
|
|
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)
|
|
|
|
The module after the Gradient pass will be:
|
|
|
|
.. code-block:: python
|
|
|
|
@I.ir_module(s_tir=True)
|
|
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))
|
|
|
|
The second example is returning multiple values and specifying the target with `target_index`:
|
|
|
|
.. code-block:: python
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
lv2: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
|
|
R.output(lv1, lv2)
|
|
return (lv1, lv2)
|
|
|
|
After = relax.transform.Gradient("main", target_index=1)(Module)
|
|
|
|
The module after the Gradient pass will be:
|
|
|
|
.. code-block:: python
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
lv2: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
|
|
R.output(lv1, lv2)
|
|
return (lv1, lv2)
|
|
|
|
@R.function
|
|
def main_adjoint(
|
|
x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")
|
|
) -> R.Tuple(
|
|
R.Tuple(R.Tensor((), dtype="float32"), 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((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
lv2: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
|
|
# bindings w.r.t. intermediate variables
|
|
# gradient of intermediate variables that is not related to the target will not
|
|
# be calculated
|
|
lv2_adjoint: R.Tensor((), dtype="float32") = R.ones((), dtype="float32")
|
|
# bindings w.r.t. parameters
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = R.zeros((3, 3), dtype="float32")
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(
|
|
lv2_adjoint, (3, 3)
|
|
)
|
|
R.output(lv1, lv2, x_adjoint, y_adjoint)
|
|
# return value: (orig_return_values, tuple(adjoints))
|
|
return ((lv1, lv2), (x_adjoint, y_adjoint))
|
|
"""
|
|
if require_grads is not None and not isinstance(require_grads, list | tvm_ffi.Array):
|
|
require_grads = [require_grads]
|
|
|
|
return _ffi_api.Gradient(func_name, require_grads, target_index) # type: ignore
|
|
|
|
|
|
def ToNonDataflow() -> tvm.ir.transform.Pass:
|
|
"""Transform all dataflow structure to non-dataflow version.
|
|
|
|
Returns
|
|
-------
|
|
ret: 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.
|
|
|
|
This effectively means that there will be no more purity tracking,
|
|
useful for low-level code generation.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
The Pass.
|
|
|
|
Note
|
|
----
|
|
Should be used after ToNonDataflow()
|
|
"""
|
|
return _ffi_api.RemovePurityChecking() # type: ignore
|
|
|
|
|
|
def DataflowUseInplaceCalls() -> tvm.ir.transform.Pass:
|
|
"""
|
|
Pass that changes calls to operators that can be done in-place
|
|
(generally, these are elementwise operations) into in-place implementations.
|
|
Supported operators will be replaced by calls to `call_tir_inplace` that invoke
|
|
in-place PrimFunc implementations of those operators (which are based on the legalizations of
|
|
those operators).
|
|
|
|
Note: ConvertToDataflow may need to be called first to provide dataflow blocks.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
The pass
|
|
"""
|
|
return _ffi_api.DataflowUseInplaceCalls()
|
|
|
|
|
|
def LambdaLift() -> tvm.ir.transform.Pass:
|
|
"""A pass that lifts local functions into global.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.LambdaLift()
|
|
|
|
|
|
def LazyGetInput() -> tvm.ir.transform.Pass:
|
|
"""A pass that requests inputs lazily.
|
|
|
|
In many cases, the size of the model weights exceeds the available
|
|
memory on a GPU. In these cases, a function that accepts all
|
|
model weights as arguments would not be able to be called. In
|
|
these cases, parameters must be loaded as they are required by the
|
|
function, and unloaded once they are no longer needed.
|
|
|
|
This pass mutates a function such that all model weights
|
|
(arguments after the first `func.attrs["num_input"]` arguments)
|
|
are loaded on demand. Rather than accepting the weights as
|
|
function arguments, the function accepts a callback argument,
|
|
which can load each parameter as needed. The callback accepts two
|
|
arguments, first the index of the model weight, and second the
|
|
name of the parameter. The callback should return the parameter
|
|
as specified.
|
|
|
|
.. code-block:: python
|
|
|
|
@R.function
|
|
def before(A: R.Tensor([16,32],"float32")):
|
|
...
|
|
|
|
@R.function
|
|
def after(fget_param: R.Callable([R.Prim('int64'), R.Any], R.Any)):
|
|
A_untyped = fget_param(0, R.str('A'))
|
|
A = R.match_cast(A_untyped, R.Tensor([16,32], "float32")
|
|
...
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
|
|
"""
|
|
return _ffi_api.LazyGetInput()
|
|
|
|
|
|
def LazySetOutput() -> tvm.ir.transform.Pass:
|
|
"""A pass that sets function outputs when available
|
|
|
|
In many cases, the size of the model weights exceeds the available
|
|
memory on a GPU. In these cases, a function that produces all
|
|
model weights as a single return value would not be able to be
|
|
called. In these cases, parameters must be returned as they are
|
|
produced, unloaded from the GPU (or saved to disk), before
|
|
producing additional outputs.
|
|
|
|
This pass mutates a function such that all outputs from a function
|
|
are returned when they are available. The function accepts an
|
|
additional callback argument, which is called with each output of
|
|
the function. The callback accepts two arguments, first the index
|
|
of the output tuple that was produced (or zero if the output is
|
|
not a tuple), and second the value itself.
|
|
|
|
.. code-block:: python
|
|
|
|
@R.function
|
|
def before(args):
|
|
...
|
|
return (A, B)
|
|
|
|
@R.function
|
|
def after(args, fset_param: R.Callable([R.Prim('int64'), R.Any])):
|
|
...
|
|
fset_param(0, A)
|
|
...
|
|
fset_param(1, B)
|
|
...
|
|
return ()
|
|
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
|
|
"""
|
|
return _ffi_api.LazySetOutput()
|
|
|
|
|
|
def ConvertToDataflow(min_size: int = 2) -> tvm.ir.transform.Pass:
|
|
"""A pass that converts consecutive dataflow operations
|
|
inside binding blocks into dataflow blocks.
|
|
|
|
Note: ConvertToDataflow may need to be called first.
|
|
|
|
Parameters
|
|
----------
|
|
min_size: int
|
|
The minimum number of consecutive dataflow bindings
|
|
the pass needs to extract a new block.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
The pass.
|
|
"""
|
|
return _ffi_api.ConvertToDataflow(min_size)
|
|
|
|
|
|
def CallTIRRewrite() -> tvm.ir.transform.Pass:
|
|
"""Perform explicit tensor allocation for call_tir and call_dps_packed.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.CallTIRRewrite() # type: ignore
|
|
|
|
|
|
def Normalize() -> tvm.ir.transform.Pass:
|
|
"""Transforming Relax IR to normal form, i.e., the expressions are normalized(no nesting
|
|
and hence the AST is in ANF), and all `ty` fields of expressions are available.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.Normalize() # type: ignore
|
|
|
|
|
|
def NormalizeGlobalVar() -> tvm.ir.transform.Pass:
|
|
"""Possibly rename the GlobalVar in an IRModule to ensure these properties:
|
|
|
|
1. (Invariant) First ensure every public function has the same name as its "global_symbol"
|
|
attribute
|
|
2. To ensure 1., we may need to rename private functions with conflicting names;
|
|
3. Finally, the name of every GlobalVar is unique in the IRModule.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.NormalizeGlobalVar() # type: ignore
|
|
|
|
|
|
def CanonicalizeBindings() -> tvm.ir.transform.Pass:
|
|
"""
|
|
Canonicalizes variable definitions
|
|
(e.g., if there is y = x and z = y, it replaces uses of y and z with x).
|
|
Also simplifies match cast nodes (eliminating redundant checks)
|
|
and tuple indices.
|
|
|
|
Best combined with constant folding and the elimination of unused definitions.
|
|
|
|
Note: If a dataflow var is used only in a binding to the dataflow block
|
|
output var (i.e., a non-dataflow var), this pass will also remove the dataflow var
|
|
and replaces the output var's binding with the dataflow var's direct definition.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.CanonicalizeBindings() # type: ignore
|
|
|
|
|
|
def EliminateCommonSubexpr(call_only=False) -> FunctionPass:
|
|
"""Eliminate common subexpressions within functions.
|
|
|
|
Note: For nested functions, this pass performs CSE *within* those functions
|
|
|
|
Parameters
|
|
----------
|
|
call_only : bool
|
|
If True, enable eliminating only call nodes.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass that eliminates common subexpressions.
|
|
"""
|
|
return _ffi_api.EliminateCommonSubexpr(call_only) # type: ignore
|
|
|
|
|
|
def UpdateVDevice(new_vdevice: tvm.ir.VDevice, index: int) -> tvm.ir.transform.Pass:
|
|
"""Update virtual device.
|
|
|
|
Parameters
|
|
----------
|
|
new_vdevice : tvm.ir.VDevice
|
|
The new virtual device.
|
|
index : int
|
|
The device index indicates the device on which the update will be performed.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
The registered pass that modifies the virtual device.
|
|
"""
|
|
return _ffi_api.UpdateVDevice(new_vdevice, index) # type: ignore
|
|
|
|
|
|
def RewriteDataflowReshape() -> tvm.ir.transform.Pass:
|
|
"""Convert all reshape-like call_tir to VM reshape operator call.
|
|
The VM reshape operator calls will be further lowered to a CreateView
|
|
operation at runtime, instead of doing real data copy.
|
|
Here "reshape-like" includes reshape, expand_dims, flatten, etc.
|
|
|
|
Note: Operates only in dataflow blocks. ConvertToDataflow may need to be called first.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.RewriteDataflowReshape() # type: ignore
|
|
|
|
|
|
def StaticPlanBlockMemory() -> tvm.ir.transform.Pass:
|
|
"""The static memory planning pass on BindingBlock level.
|
|
The pass will reuse allocated memory to its best effort, in order to
|
|
reduce the total amount of allocated memory size.
|
|
|
|
The pass "supports" dynamic shape in the way of TIR variable upper bound
|
|
annotation. We can optionally annotate the attribute "tir_var_upper_bound"
|
|
to Relax functions. The attribute value is a dict from strings to integers,
|
|
denoting the name of TIR variables to the upper bound values of the TIR vars.
|
|
Note: The annotated upper bound attribute only applies to TIR vars in the
|
|
function signature for clarity.
|
|
|
|
For example, we can annotate a Relax function with
|
|
:code:`R.func_attr({"tir_var_upper_bound": {"n": 1024}})`.
|
|
It means the maximum value of variable that names "n" in the function
|
|
signature will have upper bound 1024. And we will use 1024 as its value
|
|
during memory planning.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.StaticPlanBlockMemory() # type: ignore
|
|
|
|
|
|
def LowerAllocTensor() -> tvm.ir.transform.Pass:
|
|
"""Lower remaining instances of R.builtin.alloc_tensor
|
|
|
|
The static memory planner removes static instances of
|
|
`R.builtin.alloc_tensor`, replacing with `R.memory.alloc_storage`
|
|
and `R.memory.alloc_tensor`. However, `R.builtin.alloc_tensor`
|
|
still remains for any dynamic allocations.
|
|
|
|
This transform replaces any remaining `R.builtin.alloc_tensor`
|
|
instances with `R.memory.alloc_storage` and
|
|
`R.memory.alloc_tensor`. If no `R.builtin.alloc_tensor` are
|
|
present, this pass has no effect.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.LowerAllocTensor() # type: ignore
|
|
|
|
|
|
def KillAfterLastUse() -> tvm.ir.transform.Pass:
|
|
"""Drop all tensor/storage objects after last use
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.KillAfterLastUse() # type: ignore
|
|
|
|
|
|
def ComputePrimValue() -> tvm.ir.transform.Pass:
|
|
"""Compute all R.prim_value instances
|
|
|
|
While high-level relax can include expressions in terms of its
|
|
symbolic variables, these expressions cannot natively be computed
|
|
within relax. In order to provide values for symbolic expressions
|
|
(e.g. `R.prim_value(N*N)`, where `N` is a symbolic variable), this
|
|
pass generates a PrimFunc in which the expression can be computed.
|
|
The relax graph is then updated to include a call to that
|
|
PrimFunc, in place of the original `R.prim_value(expr)`.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.ir.transform.Pass
|
|
|
|
"""
|
|
return _ffi_api.ComputePrimValue() # type: ignore
|
|
|
|
|
|
def LowerRuntimeBuiltin() -> tvm.ir.transform.Pass:
|
|
"""Lowering generic intrinsic to VM intrinsics.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.LowerRuntimeBuiltin() # type: ignore
|
|
|
|
|
|
def VMBuiltinLower() -> tvm.ir.transform.Pass:
|
|
"""Lowering generic intrinsic to VM intrinsics.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
warnings.warn(
|
|
"tvm.relax.transform.VMBuiltinLower has been renamed to 'LowerRuntimeBuiltin'. "
|
|
"This wrapper is for backwards compatibility, and will be removed in a later update."
|
|
)
|
|
return _ffi_api.LowerRuntimeBuiltin() # type: ignore
|
|
|
|
|
|
def VMShapeLower(*, emit_err_ctx: bool = True) -> tvm.ir.transform.Pass:
|
|
"""Lower the symbolic shape and argument and match-cast structinfo matching.
|
|
|
|
Parameters
|
|
----------
|
|
emit_err_ctx: Optional[bool]
|
|
Whether emit err context string, can be turned off for testing purposes.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.VMShapeLower(emit_err_ctx) # type: ignore
|
|
|
|
|
|
def AttachGlobalSymbol() -> tvm.ir.transform.Pass:
|
|
"""Attach global_symbol to Relax functions and TIR Primfuncs for codegen.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.AttachGlobalSymbol() # type: ignore
|
|
|
|
|
|
def BindParams(
|
|
func_name: str,
|
|
params: dict[str | Var, tvm.runtime.Tensor | np.ndarray],
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Bind params of function of the module to constant tensors.
|
|
|
|
Parameters
|
|
----------
|
|
func_name: str
|
|
The function name to be bound
|
|
|
|
params: Dict[Union[str,relax.Var], Union[tvm.runtime.Tensor, np.ndarray]]
|
|
The map from parameter or parameter name to constant tensors.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
tvm_params = {}
|
|
for k, v in params.items():
|
|
if isinstance(v, np.ndarray):
|
|
v = tvm.runtime.tensor(v)
|
|
assert isinstance(v, tvm.runtime.Tensor | tvm.relax.Constant), (
|
|
f"param values are expected to be TVM.Tensor,"
|
|
f"numpy.ndarray or tvm.relax.Constant, but got {type(v)}"
|
|
)
|
|
tvm_params[k] = v
|
|
|
|
return _ffi_api.BindParams(func_name, tvm_params) # type: ignore
|
|
|
|
|
|
def BindSymbolicVars(
|
|
binding_map: Mapping[str | tvm.tirx.Var, tvm.tirx.Expr],
|
|
func_name: str | None = None,
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Bind params of function of the module to constant tensors.
|
|
|
|
Parameters
|
|
----------
|
|
binding_map : Mapping[Union[str, tvm.tirx.Var], tvm.tirx.Expr]
|
|
The map from symbolic varname to integer.
|
|
|
|
func_name : Optional[str]
|
|
The function name to be bound. If None (default), all
|
|
functions within the module will be updated.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
# Relax uses int64 for symbolic variables, but the FFI
|
|
# converts python integers into int32.
|
|
binding_map = {
|
|
key: tvm.tirx.const(value, "int64") if isinstance(value, int) else value
|
|
for key, value in binding_map.items()
|
|
}
|
|
return _ffi_api.BindSymbolicVars(binding_map, func_name) # type: ignore
|
|
|
|
|
|
def RunCodegen(
|
|
target_options: dict | None = None,
|
|
entry_functions: list[str] | None = None,
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Produce the runtime::Module with an annotated codegen and global symbol.
|
|
|
|
Parameters
|
|
----------
|
|
target_options: Optional[dict]
|
|
Pairs of a target name and compilation options
|
|
entry_functions: Optional[List[str]]
|
|
The set of entry functions to start from.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass to remove unused functions.
|
|
"""
|
|
if entry_functions is None:
|
|
entry_functions = []
|
|
|
|
# enable cutlass byoc registries
|
|
# pylint: disable=unused-import,import-outside-toplevel
|
|
from tvm.contrib import cutlass as _cutlass
|
|
|
|
return _ffi_api.RunCodegen(target_options, entry_functions) # type: ignore
|
|
|
|
|
|
def FoldConstant() -> tvm.ir.transform.Pass:
|
|
"""Fold constant expressions within dataflow blocks.
|
|
|
|
Note: ConvertToDataflow may need to be called first to provide dataflow blocks.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.FoldConstant() # type: ignore
|
|
|
|
|
|
def ExpandTupleArguments() -> tvm.ir.transform.Pass:
|
|
"""Expand tuple arguments to internal functions
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.ExpandTupleArguments() # type: ignore
|
|
|
|
|
|
def RemoveUnusedParameters() -> tvm.ir.transform.Pass:
|
|
"""Remove unused arguments to internal functions
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.RemoveUnusedParameters() # type: ignore
|
|
|
|
|
|
def RemoveUnusedOutputs() -> tvm.ir.transform.Pass:
|
|
"""Remove unused outputs from internal functions
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.RemoveUnusedOutputs() # type: ignore
|
|
|
|
|
|
def InlinePrivateFunctions() -> tvm.ir.transform.Pass:
|
|
"""Inline all private relax functions
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.InlinePrivateFunctions() # type: ignore
|
|
|
|
|
|
def AnnotateTIROpPattern() -> tvm.ir.transform.Pass:
|
|
"""Annotate Op Pattern Kind for TIR functions
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.AnnotateTIROpPattern() # type: ignore
|
|
|
|
|
|
def FuseOps(fuse_opt_level=-1) -> tvm.ir.transform.Pass:
|
|
"""This pass groups bindings in a dataflow block of Relax functions and generate a new grouped
|
|
Relax function for each group, according to the fusion algorithm described in the pass
|
|
implementation. By grouping bindings into new Relax functions, we substitute the bindings in
|
|
the function being manipulated into function calls to the new grouped function.
|
|
|
|
A follow-up pass named "FuseTIR" will generate a TIR PrimFunc for each grouped function.
|
|
|
|
Note: ConvertToDataflow may need to be called first to provide dataflow blocks.
|
|
|
|
Parameters
|
|
----------
|
|
fuse_opt_level : int
|
|
The level of fuse optimization. -1 indicates that the level will be
|
|
inferred from pass context.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for operator fusion.
|
|
"""
|
|
return _ffi_api.FuseOps(fuse_opt_level) # type: ignore
|
|
|
|
|
|
def FuseTIR() -> tvm.ir.transform.Pass:
|
|
"""Fuse primitive relax function into a larger TIR function if possible
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for tirx fusion.
|
|
"""
|
|
return _ffi_api.FuseTIR() # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("relax.transform.PatternCheckContext")
|
|
class PatternCheckContext(Object):
|
|
"""
|
|
The input of check function `FusionPattern.check`.
|
|
|
|
Parameters
|
|
----------
|
|
matched_expr: Expr
|
|
The expression that's matched with the FusionPattern.pattern.
|
|
|
|
annotated_expr: Mapping[str, Expr]
|
|
A map which contains all expressions matched by the sub patterns in
|
|
FusionPattern.annotation_patterns.
|
|
|
|
matched_bindings: Mapping[Var, Expr]
|
|
Map from variable to its value. It contains variables from bindings that is
|
|
being fused by FuseOpsByPattern.
|
|
|
|
var_usages: Mapping[Var, Sequence[Var]]
|
|
A map mapping variable definitions to a set of uses. It has all variables
|
|
used in the function.
|
|
|
|
value_to_bound_var: Mapping[Expr, Var]
|
|
Map from value to its bound variable. It doesn't have variables after the
|
|
matched expression.
|
|
"""
|
|
|
|
matched_expr: Expr
|
|
annotated_expr: Mapping[str, Expr]
|
|
matched_bindings: Mapping[Var, Expr]
|
|
var_usages: Mapping[Var, Sequence[Var]]
|
|
value_to_bound_var: Mapping[Expr, Var]
|
|
|
|
|
|
@tvm_ffi.register_object("relax.transform.FusionPattern")
|
|
class FusionPattern(Object):
|
|
"""
|
|
The pattern used by `FuseOpsByPattern`. It's mainly DFPattern but with other
|
|
information to help during the fusion pass.
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
The name of pattern. Usually it starts with the name of backend, like 'cutlass.matmul'.
|
|
|
|
pattern: DFPattern
|
|
The dataflow pattern that will be used to match expressions that can be handled
|
|
by external backends.
|
|
|
|
annotation_patterns: Mapping[str, DFPattern]
|
|
The map which is used to extract important expressions from the pattern match
|
|
result. All DFPattern in this map should be part of the `pattern`.
|
|
|
|
check: Callable[[PatternCheckContext], bool]
|
|
The function to check whether the match result is accepted.
|
|
"""
|
|
|
|
name: str
|
|
pattern: DFPattern
|
|
annotation_patterns: Mapping[str, DFPattern]
|
|
check: Callable[[PatternCheckContext], bool]
|
|
attrs_getter: Callable[[dict[str, Expr]], dict[str, str]]
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
pattern: DFPattern,
|
|
annotation_patterns: Mapping[str, DFPattern] | None = None,
|
|
check: Callable[[PatternCheckContext], bool] | None = None,
|
|
attrs_getter: Callable[[dict[str, Expr]], dict[str, str]] | None = None,
|
|
):
|
|
if annotation_patterns is None:
|
|
annotation_patterns = {}
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.FusionPattern, name, pattern, annotation_patterns, check, attrs_getter
|
|
) # type: ignore
|
|
|
|
|
|
def FuseOpsByPattern(
|
|
patterns: list[FusionPattern | tuple],
|
|
bind_constants: bool = True,
|
|
annotate_codegen: bool = False,
|
|
entry_functions: list[str] | None = None,
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Apply pattern matching to each function in the given module, and group matched expressions
|
|
into a new function.
|
|
|
|
The end result is similar to FuseOps, but fusion is driven completely by the provided patterns.
|
|
|
|
Note: Only operates within dataflow blocks. ConvertToDataflow may need to be called first.
|
|
|
|
Parameters
|
|
----------
|
|
patterns : List[Union[FusionPattern, Tuple]]
|
|
A list of patterns to be matched. The order of the patterns determines the order of priority
|
|
in which they are matched. Higher-priority patterns should come earlier in the list.
|
|
|
|
In addition to FusionPattern, a tuple can be passed as item of this list. The pattern
|
|
will be constructed through :code:`FusionPattern(*item)`
|
|
|
|
bind_constants : bool
|
|
Whether or not to keep bound constants in the grouped function.
|
|
|
|
annotate_codegen : bool
|
|
If True, wrap each created composite function with another function, whose body consists
|
|
only of a call to the composite function, and annotate the outer function with "Codegen"
|
|
and "global_symbol" attributes. The "Codegen" attribute is set as the prefix of the
|
|
corresponding pattern name. For example, "dnnl" if the pattern name is "dnnl.conv2d_relu".
|
|
|
|
This must be True if the created composite functions are intended to be offloaded to
|
|
an external backend without using the MergeCompositeFunctions pass.
|
|
|
|
entry_functions : Optional[List[str]]
|
|
The set of entry functions to start from.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for pattern-based fusion.
|
|
|
|
"""
|
|
converted_patterns = []
|
|
for pattern in patterns:
|
|
if isinstance(pattern, tuple):
|
|
converted_patterns.append(FusionPattern(*pattern))
|
|
elif isinstance(pattern, FusionPattern):
|
|
converted_patterns.append(pattern)
|
|
else:
|
|
raise ValueError(f"Invalid pattern: {pattern}")
|
|
|
|
return _ffi_api.FuseOpsByPattern(
|
|
converted_patterns,
|
|
bind_constants,
|
|
annotate_codegen,
|
|
entry_functions or [],
|
|
) # type: ignore
|
|
|
|
|
|
def MergeCompositeFunctions() -> tvm.ir.transform.Pass:
|
|
"""Group one or multiple composite functions created by FuseOpsByPattern into a new function.
|
|
The new function will be annotated with "Codegen" and "global_symbol" attributes, and it
|
|
is intented to be offloaded to an external backend.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for merging composite functions.
|
|
"""
|
|
return _ffi_api.MergeCompositeFunctions() # type: ignore
|
|
|
|
|
|
def AttachAttrLayoutFreeBuffers() -> tvm.ir.transform.Pass:
|
|
"""Attach layout free buffers to the tirx::PrimFunc.
|
|
|
|
This pass is used to attach layout free buffers to the tirx::PrimFunc according to
|
|
the function usage in the relax function. Currently, the layout free buffers are the model
|
|
weights and relax constants.
|
|
|
|
Note that we recommend applying CanonicalizeBindings before this pass.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for attaching layout free buffers.
|
|
"""
|
|
return _ffi_api.AttachAttrLayoutFreeBuffers() # type: ignore
|
|
|
|
|
|
def SplitLayoutRewritePreproc() -> tvm.ir.transform.Pass:
|
|
"""Split the TIR layout rewrite into multiple TIR functions.
|
|
This pass is used in the prepack weight after meta_schedule tuning.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for splitting TIR layout rewrite.
|
|
"""
|
|
return _ffi_api.SplitLayoutRewritePreproc() # type: ignore
|
|
|
|
|
|
def LiftTransformParams(shared_transform: bool | list[str] = False) -> tvm.ir.transform.Pass:
|
|
"""Lift transformation of the parameters of a function.
|
|
|
|
When some inputs of the function is marked as 'parameters' (the model weights), this pass
|
|
identifies the transformation of the parameters and lifts them to a separate function called
|
|
`transform_params`. `transform_params` takes a tuple of the original parameters as input and
|
|
returns a tuple of the transformed parameters. The original function will be rewritten to accept
|
|
a tuple of transformed parameters as input.
|
|
|
|
Users are expected to invoke the `transform_params` function in runtime and pass the transformed
|
|
parameters to the original function as input.
|
|
|
|
Parameters
|
|
----------
|
|
shared_transform: Union[bool, List[str]]
|
|
|
|
Indicates how the parameter transformation function will be produced
|
|
|
|
- `False` (default): A separate parameter transformation function will be
|
|
produced for each function with the `"num_input"` attribute.
|
|
|
|
- `True`: A single parameter transformation function will be produced,
|
|
containing the preprocessing steps common across all functions with
|
|
the `"num_input"` attribute.
|
|
|
|
- List[str]: A single parameter transformation function will be produced,
|
|
containing the preprocessing steps common across each function whose
|
|
name is in the list. Passing a list of all functions with the `"num_input"`
|
|
attribute or an empty list is equivalent to passing `True`.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for lifting transformation of parameters.
|
|
"""
|
|
return _ffi_api.LiftTransformParams(shared_transform) # type: ignore
|
|
|
|
|
|
def BundleModelParams(param_tuple_name: str | None = None) -> tvm.ir.transform.Pass:
|
|
"""Bundle several model parameters into a single parameter tuple
|
|
|
|
For each function, if the function has the attribute "num_input",
|
|
separate between run-time parameters and compile-time weights.
|
|
Run-time parameters (e.g. activations) are the first `num_input`
|
|
parameters, and the remainder are compile-time weights.
|
|
|
|
Parameters
|
|
----------
|
|
param_tuple_name: Optional[str]
|
|
|
|
The name of the tuple parameter. If unspecified, defaults to
|
|
"model_params".
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for bundling model parameters.
|
|
"""
|
|
return _ffi_api.BundleModelParams(param_tuple_name) # type: ignore
|
|
|
|
|
|
def LegalizeOps(
|
|
customize_legalize_map: dict[str, LegalizeFunc] | None = None,
|
|
skip_ops: list[str] | None = None,
|
|
enable_warning: bool = False,
|
|
):
|
|
"""Legalize high-level operator calls in Relax functions to call_tir
|
|
with corresponding low-level TIR PrimFuncs.
|
|
|
|
For each high-level operator, we register the way of legalizing it as a
|
|
function, which takes a context BlockBuilder and the Call being legalized
|
|
as input, and returns the legalized call. Here the input BlockBuilder is
|
|
mainly used for adding the PrimFunc created by call_te into the context
|
|
IRModule.
|
|
|
|
The legalization function for each operator is registered as an attribute (with
|
|
attribute key `FLegalize`) of the operator.
|
|
|
|
This pass provides customizability for users to use their own legalization
|
|
function for operators. The pass takes an optional customized map,
|
|
with the key to be the operator name (`str`) and value to be the function
|
|
(`LegalizeFunc`). The default legalization function will be overridden by the customized
|
|
one.
|
|
|
|
Parameters
|
|
----------
|
|
customize_legalize_map : Optional[Dict[str, LegalizeFunc]]
|
|
The customized operator legalization function map. The customized function will override
|
|
the default one.
|
|
|
|
skip_ops : Optional,List[str]]
|
|
List of ops that need to be skipped from legalization
|
|
|
|
enable_warning : bool
|
|
A boolean value indicating if to print warnings for CallNode whose op's
|
|
legalization function is not registered. By default we don't print
|
|
warnings.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass
|
|
|
|
Examples
|
|
--------
|
|
The following code shows how to use this pass:
|
|
|
|
.. code-block:: python
|
|
|
|
# Define the pass input IRModule
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
|
|
) -> R.Tensor((2, 3), "float32"):
|
|
z: R.Tensor((2, 3), "float32") = R.add(x, y)
|
|
r: R.Tensor((2, 3), "float32") = R.multiply(y, z)
|
|
return r
|
|
|
|
# Define the customized legalization function for "relax.add"
|
|
def customize_legalize_add(bb: relax.BlockBuilder, call: relax.Call) -> relax.Expr:
|
|
from tvm import topi
|
|
return bb.call_te(topi.add, call.args[1], call.args[0])
|
|
|
|
# Apply the pass with the customized function to the module.
|
|
mod = LegalizeOps({"relax.add": customize_legalize_add})(Module)
|
|
|
|
Print out the result by `mod.show()`, we can see the IRModule after
|
|
legalization becomes
|
|
|
|
.. code-block:: python
|
|
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
|
|
) -> R.Tensor((2, 3), "float32"):
|
|
z = R.call_tir(add, (y, x), (2, 3), dtype="float32")
|
|
r = R.call_tir(multiply, (y, z), (2, 3), dtype="float32")
|
|
return r
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def add(
|
|
A: T.Buffer((2, 3), "float32"),
|
|
B: T.Buffer((2, 3), "float32"),
|
|
T_add: T.Buffer((2, 3), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(2, 3):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
|
|
T.writes(T_add[v_ax0, v_ax1])
|
|
T_add[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1]
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def multiply(
|
|
A: T.Buffer((2, 3), "float32"),
|
|
B: T.Buffer((2, 3), "float32"),
|
|
T_multiply: T.Buffer((2, 3), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(2, 3):
|
|
with T.sblock("T_multiply"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
|
|
T.writes(T_multiply[v_ax0, v_ax1])
|
|
T_multiply[v_ax0, v_ax1] = A[v_ax0, v_ax1] * B[v_ax0, v_ax1]
|
|
"""
|
|
|
|
return _ffi_api.LegalizeOps(customize_legalize_map, skip_ops, enable_warning) # type: ignore
|
|
|
|
|
|
def RealizeVDevice() -> tvm.ir.transform.Pass:
|
|
"""Propagate virtual device information.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass
|
|
"""
|
|
|
|
return _ffi_api.RealizeVDevice() # type: ignore
|
|
|
|
|
|
def MetaScheduleApplyDatabase(
|
|
work_dir: str | None = None, enable_warning: bool = False
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Apply the best schedule from tuning database.
|
|
|
|
Parameters
|
|
----------
|
|
work_dir : Optional[str]
|
|
work directory to deduce default database if database is not provided
|
|
(it will be ignored when an user passes database)
|
|
enable_warning : bool
|
|
A boolean value indicating if to print warnings for TIR functions not
|
|
showing up in the database. By default we don't print warning.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass
|
|
"""
|
|
return _ffi_api.MetaScheduleApplyDatabase(work_dir, enable_warning) # type: ignore
|
|
|
|
|
|
def MetaScheduleTuneTIR(
|
|
work_dir: str,
|
|
max_trials_global: int,
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Tune TIR with MetaSchedule.
|
|
|
|
Parameters
|
|
----------
|
|
work_dir: str
|
|
work directory
|
|
max_trials_gloabl: int
|
|
maximum number of total trials allowed for tuning
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.MetaScheduleTuneTIR(work_dir, max_trials_global) # type: ignore
|
|
|
|
|
|
def MetaScheduleTuneIRMod(
|
|
params: dict[str, Tensor],
|
|
work_dir: str,
|
|
max_trials_global: int,
|
|
max_trials_per_task: int | None = None,
|
|
op_names: list[str] | None = None,
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Tune Relax IRModule with MetaSchedule.
|
|
|
|
Parameters
|
|
----------
|
|
params: Dict[str, Tensor]
|
|
model params
|
|
work_dir: str
|
|
work directory
|
|
max_trials_gloabl: int
|
|
maximum number of total trials allowed for tuning
|
|
max_trials_per_task: int
|
|
maximum number of trials per task
|
|
op_names: Optional[List[str]]
|
|
A list of operator names to specify which op to tune. When it is None, all operators
|
|
are tuned.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
return _ffi_api.MetaScheduleTuneIRMod(
|
|
params, work_dir, max_trials_global, max_trials_per_task, op_names
|
|
) # type: ignore
|
|
|
|
|
|
def DecomposeOpsForInference(func_name: str | None = None) -> tvm.ir.transform.Pass:
|
|
"""Decompose composite operators that are composed by other operators during inference.
|
|
For example, the result of batch norm (a triple) will be simplified. Attention, tensor_to_shape,
|
|
etc. can be also decomposed into a number of simplified operators as well.
|
|
|
|
Parameters
|
|
----------
|
|
func_name: Optional[str]
|
|
The name of the specified function. If not specified, the pass will run in
|
|
all functions.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass
|
|
"""
|
|
return _ffi_api.DecomposeOpsForInference(func_name) # type: ignore
|
|
|
|
|
|
def DecomposeOpsForTraining(func_name: str | None = None) -> tvm.ir.transform.Pass:
|
|
"""Decompose composite operators that are composed by other operators during training.
|
|
For example, the result of batch norm (a triple) will be simplified. Attention, tensor_to_shape,
|
|
etc. can be also decomposed into a number of simplified operators as well.
|
|
|
|
Parameters
|
|
----------
|
|
func_name: Optional[str]
|
|
The name of the specified function. If not specified, the pass will run in
|
|
all functions.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass
|
|
"""
|
|
return _ffi_api.DecomposeOpsForTraining(func_name) # type: ignore
|
|
|
|
|
|
def AlterOpImpl(
|
|
op_impl_map: dict[str, PrimFunc],
|
|
op_buffer_transforms: dict[str, list[IndexMap | Callable]],
|
|
op_buffer_axis_separators: dict[str, list[str | Callable]], # str=IndexMap.AXIS_SEPARATOR
|
|
op_buffer_input_axis_separators: dict[str, list[str | Callable]], # str=IndexMap.AXIS_SEPARATOR
|
|
):
|
|
"""Replace all PrimFunc's which have matching 'operator_name' attribute, with replacement
|
|
PrimFunc that could possibly have different layouts on i/o buffers. The layout
|
|
transformations on i/o buffers is present in the op_buffer_transforms map. Inserts the layout
|
|
transformations in the call sites of PrimFuncs being replaced to transform i/o
|
|
tensors into expected layout by new PrimFunc.
|
|
|
|
Parameters
|
|
----------
|
|
op_impl_map: Dict[str, PrimFunc]
|
|
op_kind to PrimFunc map
|
|
op_buffer_transforms: Dict[str, List[Union[IndexMap, Callable]]
|
|
op_kind to layout transformation map for each of the buffers
|
|
op_buffer_axis_separators: Dict[str, List[Union[IndexMap.AXIS_SEPARATOR, Callable]]]
|
|
op_kind to axis_separator for each index_map
|
|
op_buffer_input_axis_separators: Dict[str, List[Union[IndexMap.AXIS_SEPARATOR, Callable]]]
|
|
op_kind to axis_separator for input index_map
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
"""
|
|
for operator_name, transform_list in op_buffer_transforms.items():
|
|
l = []
|
|
for transform in transform_list:
|
|
# Extract the index_map
|
|
if isinstance(transform, Callable):
|
|
transform = IndexMap.from_func_with_separators(transform)[0]
|
|
elif isinstance(transform, Array | tuple) and isinstance(transform[0], IndexMap):
|
|
transform = transform[0]
|
|
l.append(transform)
|
|
op_buffer_transforms[operator_name] = l
|
|
|
|
return _ffi_api.AlterOpImpl(
|
|
op_impl_map,
|
|
op_buffer_transforms,
|
|
op_buffer_axis_separators,
|
|
op_buffer_input_axis_separators,
|
|
) # type: ignore
|
|
|
|
|
|
def ConvertLayout(
|
|
desired_layouts: dict[str, list[str]],
|
|
layout_cb: Callable | None = None,
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Automatic layout conversion pass.
|
|
|
|
Parameters
|
|
----------
|
|
desired_layouts : Dict[str, List[str]]
|
|
The desired layout of conv2d ops is a map from the name of the op to the desired layout
|
|
of the desired feature map, weight and output. For example, if we want to convert the
|
|
layout of conv2d from NCHW to NHWC, we can set the desired layout of conv2d to be
|
|
``{"relax.nn.conv2d": ["NHWC", "OHWI"]}``.
|
|
layout_cb : Callable
|
|
A user defined call back function that can dynamically handle operator layouts
|
|
based on Call description. desired_layouts will be ignored if layout_cb is defined.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for layout conversion.
|
|
"""
|
|
return _ffi_api.ConvertLayout(desired_layouts, layout_cb) # type: ignore
|
|
|
|
|
|
def DeadCodeElimination(entry_functions: list[str] | None = None) -> tvm.ir.transform.Pass:
|
|
"""Remove dead code in the IRModule.
|
|
Currently it removes:
|
|
|
|
1. Unused local VarBindings
|
|
(those where the bound var is unused and no impure operation is used).
|
|
2. Unused Relax functions in the module.
|
|
We detect the call chain from the entry function, and remove all unused functions.
|
|
|
|
Any binding blocks that are left empty will be removed by the normalizer.
|
|
|
|
Notes
|
|
-----
|
|
For function-wise DCE, use py:func:`tvm.relax.analysis.remove_all_unused`.
|
|
|
|
Parameters
|
|
----------
|
|
entry_functions: Optional[List[str]]
|
|
The set of entry functions to start from.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass.
|
|
"""
|
|
if entry_functions is None:
|
|
entry_functions = []
|
|
return _ffi_api.DeadCodeElimination(entry_functions) # type: ignore
|
|
|
|
|
|
def ToMixedPrecision(
|
|
out_dtype="float32", fp16_input_names: list[str] | None = None
|
|
) -> tvm.ir.transform.Pass:
|
|
"""Automatic mixed precision pass. Currently the pass assumes the input module to be fp32
|
|
only, and will automatically cast fp32 to fp16 for certain ops.
|
|
|
|
Note: Mainly operates within dataflow blocks. ConvertToDataflow may need to be called first.
|
|
|
|
Parameters
|
|
----------
|
|
out_dtype : str
|
|
The output data type of gemm/conv, which is the data type of the accumulator.
|
|
fp16_input_names : List[str]
|
|
The names of function parameters whose dtype should become fp16. The function signature
|
|
would change accordingly.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for mixed precision.
|
|
"""
|
|
return _ffi_api.ToMixedPrecision(out_dtype, fp16_input_names) # type: ignore
|
|
|
|
|
|
def SplitCallTIRByPattern(patterns: list[PrimFunc], fcodegen: Callable) -> tvm.ir.transform.Pass:
|
|
"""Split a PrimFunc into 2 parts: the first part is a TIR PrimFunc which is
|
|
matched with some pattern, and the second part is the rest of the original
|
|
PrimFunc. It will call fcodegen to generate the code for the matched pattern
|
|
to replace it with a ExternFunc call.
|
|
|
|
Parameters
|
|
----------
|
|
patterns : List[PrimFunc]
|
|
The list of patterns to match.
|
|
|
|
fcodegen: Callable[[List[MatchResult]], List[Object]]
|
|
The function to generate the code for the matched patterns.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The registered pass for splitting call_tir.
|
|
"""
|
|
return _ffi_api.SplitCallTIRByPattern(patterns, fcodegen) # type: ignore
|
|
|
|
|
|
def UpdateParamType(ty_func: Callable[[Var], Type | None]):
|
|
"""Update parameter types.
|
|
|
|
Internal bindings and the function return type are updated using Relax's
|
|
type inference rules. Errors resulting from type inference are propagated
|
|
to the user.
|
|
|
|
Parameters
|
|
----------
|
|
ty_func: Callable[[Var], Optional[Type]]
|
|
|
|
A function that is called once for each function parameter,
|
|
and returns the updated type to be used for it. If the
|
|
function returns `None`, the parameter is not modified.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The corresponding pass.
|
|
|
|
"""
|
|
return _ffi_api.UpdateParamType(ty_func) # type: ignore
|
|
|
|
|
|
def AdjustMatmulOrder():
|
|
"""Reorder `x*(A*B)` to `(x*A)*B`
|
|
|
|
Useful for optimizing LoRA computations, where `matmul(x,
|
|
LoraA*LoraB)` may be computed as `matmul(matmul(x, LoraA),
|
|
LoraB)`, reducing the total memory usage.
|
|
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The corresponding pass.
|
|
"""
|
|
|
|
return _ffi_api.AdjustMatmulOrder() # type: ignore
|
|
|
|
|
|
def ExpandMatmulOfSum():
|
|
"""Expand `matmul(x, A+B)` to `matmul(x,A) + matmul(x,B)`
|
|
|
|
If either operand can be fully computed at compile-time (only
|
|
depends on function parameters after kNumInput), this expansion is
|
|
suppressed.
|
|
|
|
Useful for optimizing LoRA computations, where `matmul(x, Base +
|
|
LoraA*LoraB)` may be expanded to `matmul(x, Base) + matmul(x,
|
|
LoraA*LoraB)`, allowing it to optimized with `CombineParallelMatmul`.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The corresponding pass.
|
|
"""
|
|
|
|
return _ffi_api.ExpandMatmulOfSum() # type: ignore
|
|
|
|
|
|
def ReorderPermuteDimsAfterConcat():
|
|
"""Reorder `concat(permute_dims(A), permute_dims(B))` into `permute_dims(concat(A,B))`
|
|
|
|
Useful for optimizing computations after `CombineParallelMatmul`.
|
|
The patterns for optimized `nn.Linear` implementations look for
|
|
`matmul(activations, permute_dims(weights))`. After
|
|
`CombineParallelMatmul`, the `matmul(activations,
|
|
concat(permute_dims(A), permute_dims(B)))` no longer matches this
|
|
pattern. Rearranging into `matmul(activations,
|
|
permute_dims(concat(A,B)))` restores the pattern match.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The corresponding pass.
|
|
"""
|
|
|
|
return _ffi_api.ReorderPermuteDimsAfterConcat() # type: ignore
|
|
|
|
|
|
def ReorderTakeAfterMatmul():
|
|
"""Reorder `matmul(x, take(weights, indices))` to `take(matmul(x,weights),indices)`
|
|
|
|
Useful for optimizing LoRA computations, where several LoRAs may
|
|
be batched together.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The corresponding pass.
|
|
"""
|
|
|
|
return _ffi_api.ReorderTakeAfterMatmul() # type: ignore
|
|
|
|
|
|
def CombineParallelMatmul(check=None):
|
|
"""Combine multiple matmul operators sharing the same LHS matrix into one,
|
|
followed by slicing. When all matmul branches in a tree have the same set of fused ops,
|
|
the fused ops are applied to the combined matmul output before slicing.
|
|
|
|
Currently, only a limited set of fused ops is supported. It includes bias add,
|
|
relu, gelu, gelu_tanh and silu activation.
|
|
|
|
Parameters
|
|
----------
|
|
check : Callable[[Var, List[Var], List[Var], Dict[Var, Expr]], bool]
|
|
A function to filter out unwanted branches, with the signature
|
|
(input, [rhs], [bias], binding) -> bool.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.transform.Pass
|
|
The corresponding pass.
|
|
"""
|
|
if check is None:
|
|
check = lambda *_: True
|
|
return _ffi_api.CombineParallelMatmul(check) # type: ignore
|
|
|
|
|
|
def RewriteCUDAGraph() -> tvm.ir.transform.Pass:
|
|
"""Rewrite a Relax module for executing with CUDA graph. This pass identifies the regions that
|
|
can be executed with CUDA graph and lifts them into new functions for runtime graph capturing.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
The registered pass for rewriting cuda graph
|
|
"""
|
|
return _ffi_api.RewriteCUDAGraph() # type: ignore
|
|
|
|
|
|
def AllocateWorkspace() -> tvm.ir.transform.Pass:
|
|
"""Allocate a workspace, represented by a tensor of size big enough for all external
|
|
functions that require a temporary storage, and append it to the arguments of external
|
|
functions.
|
|
|
|
An external function can specify its workspace requirement by the kWorkspaceSize attribute.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
The registered pass for allocating workspace.
|
|
"""
|
|
return _ffi_api.AllocateWorkspace() # type: ignore
|
|
|
|
|
|
def SpecializePrimFuncBasedOnCallSite() -> tvm.ir.transform.Pass:
|
|
"""This pass updates the var_buffer mapping of PrimFunctions from the call_tir info.
|
|
Primarily used to update the VDevice information if any changes occurred from the caller.
|
|
This pass recreates the buffers and updates the map.
|
|
|
|
Returns
|
|
-------
|
|
ret: tvm.ir.transform.Pass
|
|
The registered pass for specializing PrimFuncs based on call site.
|
|
"""
|
|
return _ffi_api.SpecializePrimFuncBasedOnCallSite() # type: ignore
|
|
|
|
|
|
def _wrap_class_function_pass(pass_cls, pass_info):
|
|
"""Wrap a python class as function pass."""
|
|
|
|
class PyFunctionPass(FunctionPass):
|
|
"""Internal wrapper class to create a class instance."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
inst = pass_cls(*args, **kwargs)
|
|
|
|
# it is important not to capture self to
|
|
# avoid a cyclic dependency
|
|
def _pass_func(func, mod, ctx):
|
|
return inst.transform_function(func, mod, ctx)
|
|
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.MakeFunctionPass,
|
|
_pass_func,
|
|
pass_info, # type: ignore
|
|
)
|
|
self._inst = inst
|
|
|
|
def __getattr__(self, name):
|
|
# fall back to instance attribute if there is not any
|
|
return self._inst.__getattribute__(name)
|
|
|
|
functools.update_wrapper(PyFunctionPass.__init__, pass_cls.__init__)
|
|
PyFunctionPass.__name__ = pass_cls.__name__
|
|
PyFunctionPass.__doc__ = pass_cls.__doc__
|
|
PyFunctionPass.__module__ = pass_cls.__module__
|
|
return PyFunctionPass
|
|
|
|
|
|
def function_pass(
|
|
pass_func=None,
|
|
opt_level=None,
|
|
name=None,
|
|
required=None,
|
|
traceable=False,
|
|
) -> Callable | FunctionPass:
|
|
"""Decorate a function pass.
|
|
|
|
This function returns a callback when pass_func
|
|
is provided. Otherwise, it returns the created function pass using the
|
|
given optimization function.
|
|
|
|
Parameters
|
|
----------
|
|
pass_func : Optional[Callable[(Function, Module, PassContext) -> Function]]
|
|
The transformation function or class.
|
|
|
|
opt_level : int
|
|
The optimization level of this function pass.
|
|
|
|
name : Optional[str]
|
|
The name of the function pass. The name could be empty. In this case, the
|
|
name of the optimization function will be used as the pass name.
|
|
|
|
required : Optional[List[str]]
|
|
The list of passes that the function pass is dependent on.
|
|
|
|
traceable: Boolean
|
|
Boolean variable whether the function pass is traceable
|
|
|
|
Returns
|
|
-------
|
|
create_function_pass : Union[Callable, FunctionPass]
|
|
|
|
A decorator will be returned if pass_func is not provided,
|
|
otherwise return the decorated result.
|
|
The returned decorator has two behaviors depending on the input:
|
|
A new FunctionPass will be returned when we decorate a pass function.
|
|
A new FunctionPass class will be returned when we decorate a class type.
|
|
|
|
Examples
|
|
--------
|
|
The following code block decorates a function pass class.
|
|
|
|
.. code-block:: python
|
|
|
|
@relax.transform.function_pass(opt_level=1)
|
|
class TestReplaceFunc:
|
|
def __init__(self, new_func):
|
|
self.new_func = new_func
|
|
|
|
def transform_function(self, func, mod, ctx):
|
|
# just for demo purposes
|
|
# transform func to new_func
|
|
return self.new_func
|
|
|
|
@R.function
|
|
def f1(x: Tensor[(m, n), "float32"]):
|
|
return x
|
|
|
|
@tvm.script.ir_module
|
|
class InputMod:
|
|
@R.function
|
|
def f2(x: Tensor[(m, n), "float32"]):
|
|
gv0 = relax.add(x, x)
|
|
return gv0
|
|
# fpass is now a special pass that replaces every
|
|
# function to f1
|
|
fpass = TestReplaceFunc(f1)
|
|
# now every function in InputMod is replaced by f1
|
|
updated_mod = fpass(InputMod)
|
|
|
|
|
|
The following code creates a function pass by decorating
|
|
a user defined transform function.
|
|
|
|
.. code-block:: python
|
|
|
|
@relax.transform.function_pass(opt_level=2)
|
|
def transform(func, mod, ctx):
|
|
# my transformations here.
|
|
return func
|
|
|
|
function_pass = transform
|
|
assert isinstance(function_pass, relax.transform.FunctionPass)
|
|
assert function_pass.info.opt_level == 2
|
|
|
|
# Given a module m, the optimization could be invoked as the follwoing:
|
|
updated_mod = function_pass(m)
|
|
# Now transform should have been applied to every function in
|
|
# the provided module m. And the updated module will be returned.
|
|
"""
|
|
|
|
if opt_level is None:
|
|
raise ValueError("Please provide opt_level for the function pass.")
|
|
|
|
required = required if required else []
|
|
if not isinstance(required, list | tuple):
|
|
raise TypeError("Required is expected to be the type of " + "list/tuple.")
|
|
|
|
def create_function_pass(pass_arg):
|
|
"""Internal function that creates a function pass"""
|
|
fname = name if name else pass_arg.__name__
|
|
info = tvm.transform.PassInfo(opt_level, fname, required, traceable)
|
|
if inspect.isclass(pass_arg):
|
|
return _wrap_class_function_pass(pass_arg, info)
|
|
if not isinstance(pass_arg, types.FunctionType | types.LambdaType):
|
|
raise TypeError("pass_func must be a callable for Function pass")
|
|
return _ffi_api.MakeFunctionPass(pass_arg, info) # type: ignore
|
|
|
|
if pass_func:
|
|
return create_function_pass(pass_func)
|
|
return create_function_pass
|
|
|
|
|
|
def _wrap_class_dataflowblock_pass(pass_cls, pass_info):
|
|
"""Wrap a python class as dataflowblock pass"""
|
|
|
|
class PyDataflowBlockPass(DataflowBlockPass):
|
|
"""Internal wrapper class to create a class instance."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
inst = pass_cls(*args, **kwargs)
|
|
|
|
# it is important not to capture self to
|
|
# avoid a cyclic dependency
|
|
def _pass_func(func, mod, ctx):
|
|
return inst.transform_dataflowblock(func, mod, ctx)
|
|
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.MakeDataflowBlockPass,
|
|
_pass_func,
|
|
pass_info, # type: ignore
|
|
)
|
|
self._inst = inst
|
|
|
|
def __getattr__(self, name):
|
|
# fall back to instance attribute if there is not any
|
|
return self._inst.__getattribute__(name)
|
|
|
|
functools.update_wrapper(PyDataflowBlockPass.__init__, pass_cls.__init__)
|
|
PyDataflowBlockPass.__name__ = pass_cls.__name__
|
|
PyDataflowBlockPass.__doc__ = pass_cls.__doc__
|
|
PyDataflowBlockPass.__module__ = pass_cls.__module__
|
|
return PyDataflowBlockPass
|
|
|
|
|
|
def dataflowblock_pass(
|
|
pass_func=None, opt_level=None, name=None, required=None, traceable=False
|
|
) -> Callable | DataflowBlockPass:
|
|
"""Decorate a dataflowblock pass.
|
|
|
|
This function returns a callback when pass_func
|
|
is provided. Otherwise, it returns the created dataflowblock pass using the
|
|
given optimization function.
|
|
|
|
Parameters
|
|
----------
|
|
pass_func : Optional[Callable[(DataflowBlock, Module, PassContext) -> DataflowBlock]]
|
|
The transformation function or class.
|
|
|
|
opt_level : int
|
|
The optimization level of this dataflowblock pass.
|
|
|
|
name : Optional[str]
|
|
The name of the dataflowblock pass. The name could be empty. In this case, the
|
|
name of the optimization function will be used as the pass name.
|
|
|
|
required : Optional[List[str]]
|
|
The list of passes that the dataflowblock pass is dependent on.
|
|
|
|
traceable: Boolean
|
|
Boolean variable whether the dataflowblock pass is traceable
|
|
|
|
Returns
|
|
-------
|
|
create_dataflowblock_pass : Union[Callable, DataflowBlockPass]
|
|
|
|
A decorator will be returned if pass_func is not provided,
|
|
otherwise return the decorated result.
|
|
The returned decorator has two behaviors depending on the input:
|
|
A new DataflowBlockPass will be returned when we decorate a pass function.
|
|
A new DataflowBlockPass class will be returned when we decorate a class type.
|
|
|
|
Examples
|
|
--------
|
|
The following code block decorates a dataflowblock pass class.
|
|
|
|
.. code-block:: python
|
|
|
|
@relax.transform.dataflowblock_pass(opt_level=1)
|
|
class TestReplaceBinding:
|
|
# Simple test function to replace the first VarBinding to another.
|
|
|
|
def __init__(self):
|
|
# create a new VarBinding
|
|
m, n = tirx.Var("m", "int64"), tirx.Var("n", "int64")
|
|
lv0 = relax.Var("lv1", relax.TensorType([m, n], "float32"))
|
|
val = relax.const(np.random.rand(24, 56))
|
|
self.new_binding = relax.VarBinding(lv0, val)
|
|
|
|
def transform_dataflowblock(self, block, mod, ctx):
|
|
# just for demo purposes
|
|
# Replace the first binding in the DataflowBlock
|
|
new_bindings = [self.new_binding, block.bindings[1]]
|
|
new_block = relax.expr.DataflowBlock(new_bindings, block.span)
|
|
return new_block
|
|
|
|
@tvm.script.ir_module
|
|
class InputMod:
|
|
@R.function
|
|
def f1(x: Tensor[(m, n), "float32"]):
|
|
with relax.dataflow():
|
|
lv0 = relax.multiply(x, x)
|
|
gv0 = relax.add(x, x)
|
|
relax.output(gv0)
|
|
return gv0
|
|
# block_pass is now a special pass that replaces every
|
|
# first binding to the constant value binding
|
|
block_pass = TestReplaceBinding()
|
|
# now every first binding in DataflowBlock of InputMod
|
|
# is replaced by new_binding
|
|
updated_mod = block_pass(InputMod)
|
|
|
|
|
|
The following code creates a dataflowblock pass by decorating
|
|
a user defined transform function.
|
|
|
|
.. code-block:: python
|
|
|
|
@relax.transform.dataflowblock_pass(opt_level=2)
|
|
def transform(block, mod, ctx):
|
|
# my transformations here.
|
|
return block
|
|
|
|
block_pass = transform
|
|
assert isinstance(block_pass, relax.transform.DataflowBlockPass)
|
|
assert block_pass.info.opt_level == 2
|
|
|
|
# Given a module m, the optimization could be invoked as the follwoing:
|
|
updated_mod = block_pass(m)
|
|
# Now transform should have been applied to every DataflowBlock in
|
|
# the provided module m. And the updated module will be returned.
|
|
"""
|
|
|
|
if opt_level is None:
|
|
raise ValueError("Please provide opt_level for the dataflowblock pass.")
|
|
|
|
required = required if required else []
|
|
if not isinstance(required, list | tuple):
|
|
raise TypeError("Required is expected to be the type of " + "list/tuple.")
|
|
|
|
def create_dataflowblock_pass(pass_arg):
|
|
"""Internal function that creates a dataflowblock pass"""
|
|
fname = name if name else pass_arg.__name__
|
|
info = tvm.transform.PassInfo(opt_level, fname, required, traceable)
|
|
if inspect.isclass(pass_arg):
|
|
return _wrap_class_dataflowblock_pass(pass_arg, info)
|
|
if not isinstance(pass_arg, types.FunctionType | types.LambdaType):
|
|
raise TypeError("pass_func must be a callable for DataflowBlock pass")
|
|
return _ffi_api.MakeDataflowBlockPass(pass_arg, info) # type: ignore
|
|
|
|
if pass_func:
|
|
return create_dataflowblock_pass(pass_func)
|
|
return create_dataflowblock_pass
|