wehub-resource-sync
808 lines
27 KiB
Python
808 lines
27 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=no-else-return, invalid-name, unused-argument, import-outside-toplevel
|
|
# ruff: noqa: RUF012
|
|
"""Developer API of constructing Relax AST."""
|
|
|
|
from collections.abc import Callable, Sequence
|
|
from typing import Any, Optional
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
from tvm import relax as rx
|
|
from tvm import tirx
|
|
from tvm.ir.module import IRModule
|
|
from tvm.runtime import Object
|
|
|
|
from . import _ffi_api
|
|
from .expr import BaseFunc, Binding, BindingBlock, Expr, GlobalVar, Tuple, Var
|
|
from .op.base import call_tir, call_tir_with_grad
|
|
from .type import Type
|
|
from .utils import gen_call_tir_inputs
|
|
|
|
|
|
class FunctionScope:
|
|
"""Auxiliary scope for function"""
|
|
|
|
def __init__(self, block_builder, name, params, attrs, is_pure):
|
|
self._bb = block_builder
|
|
self._name = name
|
|
self._params = params
|
|
self._attrs = attrs
|
|
self._is_pure = is_pure
|
|
|
|
# Blocks that have been collected within the function
|
|
self._blocks = []
|
|
# a boolean flag that tracks if emit_func_output has been called
|
|
self._is_emit_func_output_called = False
|
|
|
|
def __enter__(self):
|
|
self._bb._enter_function_scope(self)
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
# __exit__ should properly handle the case where the with block exits with an exception
|
|
# when handling error case in exit, always check if there is already an exception
|
|
# been thrown in the with block
|
|
self._bb._exit_function_scope(exc_type, exc_val, exc_tb)
|
|
|
|
|
|
class DataflowScope:
|
|
"""Auxiliary scope for Dataflow block"""
|
|
|
|
def __init__(self, block_builder):
|
|
self._bb = block_builder
|
|
|
|
def __enter__(self):
|
|
block = self._bb._end_block()
|
|
if len(block.bindings) > 0:
|
|
self._bb._func._blocks.append(block)
|
|
self._bb._begin_dataflow_block()
|
|
|
|
def __exit__(self, ptype, value, trace):
|
|
block = self._bb._end_block()
|
|
if len(block.bindings) > 0:
|
|
self._bb._func._blocks.append(block)
|
|
self._bb._begin_binding_block()
|
|
|
|
|
|
class TestingScope:
|
|
"""Auxiliary scope for testing purposes"""
|
|
|
|
def __init__(self, block_builder, def_vars):
|
|
self._bb = block_builder
|
|
shape_vars = []
|
|
for var in def_vars:
|
|
if isinstance(var, tvm.tirx.Var):
|
|
shape_vars.append(var)
|
|
else:
|
|
raise ValueError("def_vars only can take tirx.Var")
|
|
# setup a dummy var so shape is in scope.
|
|
sparam = rx.Var("sparam", rx.ShapeType(shape_vars))
|
|
self._scope_params = [sparam]
|
|
|
|
def __enter__(self):
|
|
self._bb.begin_scope(self._scope_params)
|
|
self._bb._begin_dataflow_block()
|
|
|
|
def __exit__(self, ptype, value, trace):
|
|
self._bb._end_block()
|
|
self._bb.end_scope()
|
|
|
|
|
|
@tvm_ffi.register_object("relax.BlockBuilder")
|
|
class BlockBuilder(Object):
|
|
"""A builder to build Relax IR for testing and dev.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
m = tirx.Var("m", "int32")
|
|
n = tirx.Var("n", "int32")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
with bb.function([x, y], "func"):
|
|
with bb.dataflow() as df:
|
|
lv0 = bb.emit(rx.add(x, y))
|
|
lv1 = bb.emit(rx.multiply(lv0, y))
|
|
gv0 = bb.emit_output(lv1)
|
|
bb.emit_func_output(gv0)
|
|
mod = bb.get()
|
|
|
|
BlockBuilder can also be used to construct neural networks with nn.Module API
|
|
|
|
.. code-block:: python
|
|
|
|
from tvm.relax.testing import nn
|
|
|
|
n = tirx.Var("n", "int64")
|
|
input_size = 784
|
|
hidden_sizes = [128, 32]
|
|
output_size = 10
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("main"):
|
|
model = nn.Sequential(
|
|
nn.Linear(input_size, hidden_sizes[0]),
|
|
nn.ReLU(),
|
|
nn.Linear(hidden_sizes[0], hidden_sizes[1]),
|
|
nn.ReLU(),
|
|
nn.Linear(hidden_sizes[1], output_size),
|
|
nn.LogSoftmax(),
|
|
)
|
|
data = nn.Placeholder((n, input_size), name="data")
|
|
output = model(data)
|
|
params = [data] + model.parameters()
|
|
builder.emit_func_output(output, params=params)
|
|
mod = bb.get()
|
|
"""
|
|
|
|
__slots__ = ("__dict__",)
|
|
|
|
_stack = []
|
|
|
|
@staticmethod
|
|
def current() -> Optional["BlockBuilder"]:
|
|
"""Returns the current BlockBuilder."""
|
|
if BlockBuilder._stack:
|
|
return BlockBuilder._stack[-1]
|
|
else:
|
|
return None
|
|
|
|
def __init__(self, mod: IRModule = None):
|
|
# Which functions are currently being defined
|
|
self._func_stack: list[FunctionScope] = []
|
|
self.__init_handle_by_constructor__(_ffi_api.BlockBuilderCreate, mod) # type: ignore
|
|
|
|
def _begin_dataflow_block(self) -> None:
|
|
_ffi_api.BlockBuilderBeginDataflowBlock(self) # type: ignore
|
|
|
|
def _begin_binding_block(self) -> None:
|
|
_ffi_api.BlockBuilderBeginBindingBlock(self) # type: ignore
|
|
|
|
def _end_block(self) -> BindingBlock:
|
|
return _ffi_api.BlockBuilderEndBlock(self) # type: ignore
|
|
|
|
@property
|
|
def _func(self):
|
|
if self._func_stack:
|
|
return self._func_stack[-1]
|
|
else:
|
|
raise RuntimeError(
|
|
"Cannot access BlockBuilder._func when outside a bb._function() block"
|
|
)
|
|
|
|
def _enter_function_scope(self, func_scope):
|
|
BlockBuilder._stack.append(self)
|
|
self._func_stack.append(func_scope)
|
|
self.begin_scope(func_scope._params)
|
|
self._begin_binding_block()
|
|
|
|
def _exit_function_scope(self, exc_type, exc_val, exc_tb):
|
|
# record
|
|
is_emit_func_output_called = self._func._is_emit_func_output_called
|
|
# recover to default state
|
|
self._func_stack.pop()
|
|
|
|
assert BlockBuilder._stack
|
|
assert BlockBuilder._stack[-1] is self
|
|
BlockBuilder._stack.pop()
|
|
|
|
# NOTE: we must raise after we recover the state so future
|
|
# block builder scoping functions correctly
|
|
if exc_type is None:
|
|
if not is_emit_func_output_called:
|
|
raise RuntimeError("emit_func_output must be called in a relax function.")
|
|
|
|
def function(
|
|
self,
|
|
name: str,
|
|
params: Var | Tuple | list[Var] | None = None,
|
|
attrs: dict[str, Object] | None = None,
|
|
pure: bool = True,
|
|
private: bool = False,
|
|
) -> FunctionScope:
|
|
"""Annotate a Relax function.
|
|
|
|
Parameters
|
|
----------
|
|
name : str, optional
|
|
The name of the function
|
|
|
|
params : tvm.relax.Var | Tuple | List[tvm.relax.Var], optional
|
|
The parameters of the function.
|
|
If params is None, it means deferring initialization of function parameters
|
|
until emit_func_output.
|
|
|
|
attrs : Dict[str, Object], optional
|
|
The function attrs
|
|
|
|
pure : bool, optional
|
|
Whether the function is annotated as pure.
|
|
|
|
private : bool, optional
|
|
Whether the function is annotated as private.
|
|
If the function is private, it will not have a global symbol attribute.
|
|
If it is not private and not an inner function, then it will have
|
|
a global symbol attribute (mapped to the function's name)
|
|
|
|
Returns
|
|
-------
|
|
ret: FunctionScope
|
|
A FunctionScope for building a Relax function node.
|
|
"""
|
|
if isinstance(params, rx.Var):
|
|
params = [params]
|
|
elif isinstance(params, list | tuple):
|
|
for param in params:
|
|
if not isinstance(param, rx.Var):
|
|
raise TypeError(
|
|
f"each element of function parameters must be of type tvm.relax.Var,\
|
|
but got: {type(param)}"
|
|
)
|
|
if attrs is None:
|
|
attrs = {}
|
|
# The block builder does not permit nesting functions, per above comment,
|
|
# so no further check should be needed
|
|
if not private:
|
|
attrs["global_symbol"] = name
|
|
|
|
return FunctionScope(self, name, params, attrs, is_pure=pure)
|
|
|
|
def testing_scope(self, def_vars: list[tirx.Var]) -> TestingScope:
|
|
"""Start a scope for unit-testing purposes.
|
|
|
|
Parameters
|
|
----------
|
|
def_vars: List[tirx.Var]
|
|
List of symbolic variables that are marked as defined in scope.
|
|
|
|
Returns
|
|
-------
|
|
ret: TestingScope
|
|
A TestingScope to setup builder for emit and other purposes.
|
|
"""
|
|
return TestingScope(self, def_vars)
|
|
|
|
def dataflow(self) -> DataflowScope:
|
|
"""Annotate a Relax dataflow block.
|
|
|
|
Returns
|
|
-------
|
|
ret: DataflowScope
|
|
A DataflowScope for building a Relax dataflow block.
|
|
"""
|
|
return DataflowScope(self)
|
|
|
|
def _normalize_python_tuple(self, expr: Expr | Sequence[Expr]):
|
|
"""Internal utility function to convert to relax.Tuple
|
|
|
|
The `emit`, `emit_output`, and `emit_func_output` can be
|
|
called with python `list` or `tuple` objects. These objects
|
|
should be converted to `relax.Tuple` prior to calling an FFI
|
|
function, as they would otherwise be converted to
|
|
`tvm_ffi.Array`. In addition, any nested tuple objects
|
|
should be converted.
|
|
"""
|
|
if isinstance(expr, list | tuple):
|
|
return Tuple([self._normalize_python_tuple(element) for element in expr])
|
|
elif expr is None:
|
|
from . import op
|
|
|
|
return op.null_value()
|
|
else:
|
|
return expr
|
|
|
|
def emit(self, expr: Expr, name_hint: str = "") -> Var:
|
|
"""Emit an expr.
|
|
This infers the shape and type of the expr, create a variable,
|
|
and bind the expr to the variable.
|
|
|
|
Parameters
|
|
----------
|
|
expr : tvm.relax.Expr
|
|
The Expr to be emitted.
|
|
|
|
name_hint : str
|
|
Name hint for the bound variable.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.relax.Var
|
|
A newly created variable that gets bound to the input expr.
|
|
"""
|
|
expr = self._normalize_python_tuple(expr)
|
|
return _ffi_api.BlockBuilderEmit(self, expr, name_hint) # type: ignore
|
|
|
|
def call_te(self, func: Callable, *args: Any, **kwargs: Any) -> Expr:
|
|
"""Generate a call node according to the te function.
|
|
This function converts arguments from relax expression to te tensor,
|
|
The callback func should return a te tensor or a list of te tensors.
|
|
Please see detailed example in emit_te
|
|
|
|
Parameters
|
|
----------
|
|
func : Callable
|
|
A function that returns a te tensor or a list of te tensors.
|
|
|
|
args : Any, optional
|
|
arguments passed to the function.
|
|
|
|
kwargs : Any, optional
|
|
The keyword arguments passed to the function.
|
|
Note that the following keyword args are reserved:
|
|
|
|
- 'primfunc_name_hint' for passing name hint to the PrimFunc
|
|
that gets generated.
|
|
- 'primfunc_attrs' is reserved for passing func attributes to
|
|
be added to the PrimFunc that gets created.
|
|
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.relax.Call
|
|
A newly created call node
|
|
"""
|
|
|
|
primfunc_name = kwargs.pop("primfunc_name_hint", None)
|
|
tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
|
|
|
|
if not primfunc_name:
|
|
primfunc_name = func.__name__
|
|
gvar = self.add_func(tir_func, primfunc_name)
|
|
|
|
return call_tir(gvar, call_args, output_ty, tir_vars)
|
|
|
|
def call_te_with_grad(
|
|
self,
|
|
func: Callable,
|
|
*args: Any,
|
|
te_grad_name: str,
|
|
te_grad_kwargs: dict[str, Object] | None = None,
|
|
**kwargs: Any,
|
|
) -> Expr:
|
|
"""Generate a call node according to the te function.
|
|
This method will generate a call_tir_with_grad node, i.e. a call_tir node bound with a
|
|
te gradient function (refered by te_grad_name).
|
|
|
|
Parameters
|
|
----------
|
|
func : Callable
|
|
A function that returns a te tensor or a list of te tensors.
|
|
|
|
args : Any, optional
|
|
arguments passed to the function.
|
|
|
|
te_grad_name : str
|
|
The registered name of the te gradient function associated with the call_tir_with_grad
|
|
node. Must be provided as a keyword argument.
|
|
|
|
te_grad_kwargs : Dict[str, Object], optional
|
|
The keyword arguments passed to the te gradient function.
|
|
Optionally provided as a keyword argument. Default: {}.
|
|
|
|
kwargs : Any, optional
|
|
The keyword arguments passed to the function.
|
|
Note that the following keyword args are reserved:
|
|
|
|
- 'primfunc_name_hint' for passing name hint to the PrimFunc
|
|
that gets generated.
|
|
- 'primfunc_attrs' is reserved for passing func attributes to
|
|
be added to the PrimFunc that gets created.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.relax.Call
|
|
A newly created call node
|
|
"""
|
|
|
|
primfunc_name = kwargs.pop("primfunc_name_hint", None)
|
|
tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
|
|
|
|
if te_grad_kwargs is None:
|
|
te_grad_kwargs = {}
|
|
|
|
if not primfunc_name:
|
|
primfunc_name = func.__name__
|
|
gvar = self.add_func(tir_func, primfunc_name)
|
|
|
|
return call_tir_with_grad(
|
|
gvar, call_args, output_ty, te_grad_name, te_grad_kwargs, tir_vars
|
|
)
|
|
|
|
def emit_te(self, func: Callable, *args: Any, **kwargs: Any) -> Var:
|
|
"""Emit a call node according to the te function.
|
|
This function converts arguments from relax expression to te tensor,
|
|
The callback func should return a te tensor or a list of te tensors.
|
|
|
|
Parameters
|
|
----------
|
|
func : Callable
|
|
A function that returns a te tensor or a list of te tensors.
|
|
|
|
args : Any, optional
|
|
arguments passed to the function.
|
|
|
|
kwargs : Any, optional
|
|
The keyword arguments passed to the function.
|
|
Note that the key "primfunc_name_hint" is reserved for passing name hint
|
|
to the PrimFunc that gets generated.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.relax.Var
|
|
A newly created variable that gets bound to the call code.
|
|
|
|
Example
|
|
-------
|
|
|
|
.. code-block:: python
|
|
|
|
bb = rx.BlockBuilder()
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", rx.TensorType([n, m], "float32"))
|
|
y = rx.Var("y", rx.TensorType([n, m], "float32"))
|
|
|
|
def te_func(args, args_dict, msg):
|
|
A = args[0]
|
|
B = args_dict["B"]
|
|
return te.compute((128, 128), lambda i, j: A[i, j] + B[i, j])
|
|
|
|
with bb.function([x, y], "rx_func"):
|
|
out = bb.emit_te(te_func, [x], {"B": y}, msg="hello")
|
|
bb.emit_func_output(out)
|
|
|
|
will result in TVMScript
|
|
|
|
.. code-block:: python
|
|
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def te_func(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle,
|
|
var_compute: T.handle) -> None:
|
|
# function attr dict
|
|
T.func_attr({"tirx.noalias": True})
|
|
m = T.int64()
|
|
n = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [n, m], dtype="float32")
|
|
rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [n, m], dtype="float32")
|
|
compute = T.match_buffer(var_compute, [128, 128], dtype="float32")
|
|
# body
|
|
# with T.sblock("root")
|
|
for i0, i1 in T.grid(128, 128):
|
|
with T.sblock("compute"):
|
|
i, j = T.axis.remap("SS", [i0, i1])
|
|
T.reads([rxplaceholder[i, j], rxplaceholder_1[i, j]])
|
|
T.writes([compute[i, j]])
|
|
compute[i, j] = rxplaceholder[i, j] + rxplaceholder_1[i, j]
|
|
|
|
@R.function
|
|
def rx_func(x: Tensor((n, m), "float32"), y: Tensor((n, m), "float32")) -> Tensor:
|
|
# block 0
|
|
gv = relax.call_tir("te_func", (x, y), R.Tensor((128, 128), "float32"))
|
|
return gv
|
|
|
|
Example
|
|
-------
|
|
|
|
.. code-block:: python
|
|
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
x = relax.Var("x", relax.TensorType([n], "float32"))
|
|
y = relax.Var("y", relax.TensorType([n + 1], "float32"))
|
|
|
|
def te_func(A):
|
|
C = te.compute((n + 1), lambda i: A[i])
|
|
return C
|
|
|
|
with bb.function("rx_func", [x, y]):
|
|
x1 = bb.emit_te(te_func, y)
|
|
bb.emit_func_output(x1)
|
|
|
|
will result in TVMScript
|
|
|
|
.. code-block:: python
|
|
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def te_func(var_rxplaceholder: T.handle, var_compute: T.handle, n: T.int64) -> None:
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [n + T.int64(1)],
|
|
dtype="float32")
|
|
compute = T.match_buffer(var_compute, [n + T.int64(1)], dtype="float32")
|
|
# body
|
|
# with T.sblock("root")
|
|
for i0 in T.serial(0, n + T.int64(1)):
|
|
with T.sblock("compute"):
|
|
i = T.axis.spatial(n + T.int64(1), i0)
|
|
T.reads([rxplaceholder[i]])
|
|
T.writes([compute[i]])
|
|
compute[i] = rxplaceholder[i]
|
|
|
|
@R.function
|
|
def rx_func(x: Tensor((n,), "float32"), y: Tensor(((n + 1),), "float32"))
|
|
-> Tensor(None, "float32", ndim=-1):
|
|
# block 0
|
|
gv = relax.call_tir(te_func, (y,), R.Tensor((n + 1,), "float32"), (n,))
|
|
return gv
|
|
"""
|
|
name_hint = kwargs.pop("name_hint", "")
|
|
return self.emit(self.call_te(func, *args, **kwargs), name_hint=name_hint)
|
|
|
|
def match_cast(self, value: Expr, ty: Type, name_hint: str = "") -> Var:
|
|
"""Emit a MatchCast.
|
|
|
|
Parameters
|
|
----------
|
|
value : tvm.relax.Expr
|
|
The value of the MatchCast to be emitted.
|
|
|
|
ty : Type
|
|
The type to be matched.
|
|
|
|
name_hint : str
|
|
The name of the match cast
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.relax.Var
|
|
A newly created variable that get bounds to be the casted result.
|
|
"""
|
|
return _ffi_api.BlockBuilderEmitMatchCast(
|
|
self,
|
|
value,
|
|
ty,
|
|
name_hint,
|
|
) # type: ignore
|
|
|
|
def emit_output(self, output: Expr | Tuple | list[Expr], name_hint: str = "") -> Var:
|
|
"""Emit output for the current dataflow block or function.
|
|
|
|
Parameters
|
|
----------
|
|
output : Expr | Tuple | List[Expr]
|
|
The output of the current block/function.
|
|
|
|
name_hint : str
|
|
Name hint for the bound variable.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.relax.Var
|
|
The return variable which gets bound to the output.
|
|
"""
|
|
output = self._normalize_python_tuple(output)
|
|
return _ffi_api.BlockBuilderEmitOutput(self, output, name_hint) # type: ignore
|
|
|
|
def emit_func_output(
|
|
self,
|
|
output: Expr | Tuple | list[Expr],
|
|
params: Var | Tuple | list[Var] | None = None,
|
|
) -> GlobalVar:
|
|
"""Emit output for the function.
|
|
|
|
Parameters
|
|
----------
|
|
output : Expr | Tuple | List[Expr]
|
|
The output of the current block/function.
|
|
|
|
params : tvm.relax.Var | Tuple | List[tvm.relax.Var], optional
|
|
The parameters of the function to be built.
|
|
If params is None, it means the params have been initialized in the function with scope.
|
|
|
|
Returns
|
|
-------
|
|
gvar: tvm.ir.GlobalVar
|
|
|
|
A GlobalVar representing the function
|
|
"""
|
|
if self._func._is_emit_func_output_called:
|
|
raise RuntimeError("emit_func_output must be called exactly once in a relax function.")
|
|
self._func._is_emit_func_output_called = True
|
|
|
|
if self._func._params is not None and params is not None:
|
|
raise RuntimeError(
|
|
"function parameters have been initialized in the function with scope."
|
|
)
|
|
|
|
if self._func._params is None and params is None:
|
|
raise RuntimeError("Relax function must have parameter.")
|
|
|
|
if self._func._params is None:
|
|
self._func._params = params
|
|
|
|
if BlockBuilder.current() is not self:
|
|
raise RuntimeError("BlockBuilder.current() must be self.")
|
|
|
|
output = self._normalize_python_tuple(output)
|
|
|
|
block = self._end_block()
|
|
if len(block.bindings) > 0:
|
|
self._func._blocks.append(block)
|
|
|
|
seqe = rx.SeqExpr(self._func._blocks, output)
|
|
|
|
# If the parameters were not provided as part of
|
|
# `bb.function()`, then any variables provided from the params
|
|
# are not in scope. Otherwise, TIR variables used in dynamic
|
|
# inputs are removed as undefined (e.g. Replacing
|
|
# `R.Tensor(["batch_size"])` with `R.Tensor(ndims=1)`).
|
|
self.begin_scope(self._func._params)
|
|
try:
|
|
seqe = self.normalize(seqe)
|
|
finally:
|
|
self.end_scope()
|
|
|
|
# do not specify ret_ty and let constructor deduce
|
|
# from seqe.ty
|
|
func = rx.Function(self._func._params, seqe, is_pure=self._func._is_pure)
|
|
for key, value in self._func._attrs.items():
|
|
func = func.with_attr(key, value)
|
|
self.end_scope()
|
|
return self.add_func(func, self._func._name)
|
|
|
|
def normalize(self, expr: Expr) -> Expr:
|
|
"""Normalize an Expr to complete its shape and type.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The input expr.
|
|
|
|
Returns
|
|
-------
|
|
ret : Expr
|
|
The expr with normalized shape and type.
|
|
"""
|
|
return _ffi_api.BlockBuilderNormalize(self, expr) # type: ignore
|
|
|
|
def get(self) -> tvm.IRModule:
|
|
"""Return intermediate IRModule. For the situation where the IRModule is needed in the
|
|
middle of a building process.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.IRModule
|
|
An IRModule with Relax and TIR functions being built.
|
|
"""
|
|
return _ffi_api.BlockBuilderGetContextIRModule(self) # type: ignore
|
|
|
|
def finalize(self) -> tvm.IRModule:
|
|
"""Finalize the building process and return the result IRModule.
|
|
|
|
Possibly rename GlobalVars in the IRModule to ensure name uniqueness and the invariant:
|
|
every public function has the same name as its "global_symbol" attribute.
|
|
|
|
Note this method should be called only once at the end of the building process, since it may
|
|
invalidate global vars previously returned by this builder.
|
|
See also tvm.relax.transform.NormalizeGlobalVar.
|
|
|
|
Returns
|
|
-------
|
|
ret : tvm.IRModule
|
|
An IRModule with Relax and TIR functions being built.
|
|
"""
|
|
return _ffi_api.BlockBuilderFinalize(self) # type: ignore
|
|
|
|
def get_unique_name(self, name_prefix: str) -> str:
|
|
"""Generate a unique name with a specified prefix.
|
|
|
|
Parameters
|
|
----------
|
|
name_hint : str
|
|
The name prefix.
|
|
|
|
Returns
|
|
-------
|
|
ret : str
|
|
The generated name.
|
|
"""
|
|
return _ffi_api.BlockBuilderGetUniqueName(self, name_prefix) # type: ignore
|
|
|
|
def add_func(self, func: BaseFunc, func_name: str) -> GlobalVar:
|
|
"""Add a Relax function or a TIR PrimFunc to the IRModule being built.
|
|
|
|
Parameters
|
|
----------
|
|
func : BaseFunc
|
|
The function to be added.
|
|
|
|
func_name : str
|
|
The name of the function to be added.
|
|
|
|
Returns
|
|
-------
|
|
gvar : GlobalVar
|
|
The global var bound to the added function.
|
|
"""
|
|
return _ffi_api.BlockBuilderAddFunction(self, func, func_name) # type: ignore
|
|
|
|
def update_func(self, gv: GlobalVar, updated_func: BaseFunc) -> None:
|
|
"""Add a Relax function or a TIR PrimFunc to the IRModule being built.
|
|
|
|
Parameters
|
|
----------
|
|
gv : GlobalVar
|
|
The global var referring the function to be updated.
|
|
|
|
updated_func : BaseFunc
|
|
The updated function.
|
|
"""
|
|
return _ffi_api.BlockBuilderUpdateFunction(self, gv, updated_func) # type: ignore
|
|
|
|
def current_block_is_dataflow(self) -> bool:
|
|
"""Check if the block being built is DataflowBlock or not.
|
|
|
|
Returns
|
|
-------
|
|
ret : bool
|
|
A boolean that indicates if the block being built is DataflowBlock or not.
|
|
"""
|
|
return _ffi_api.BlockBuilderCurrentBlockIsDataFlow(self) # type: ignore
|
|
|
|
def emit_normalized(self, binding: Binding) -> None:
|
|
"""Emit an already normalized binding.
|
|
|
|
Parameters
|
|
----------
|
|
binding: Binding
|
|
The binding to be emitted.
|
|
"""
|
|
_ffi_api.BlockBuilderEmitNormalized(self, binding) # type: ignore
|
|
|
|
def lookup_binding(self, var: Var) -> Expr | None:
|
|
"""Lookup a var in the binding table.
|
|
|
|
Parameters
|
|
----------
|
|
var: Var
|
|
The input var.
|
|
|
|
Returns
|
|
-------
|
|
expr: Expr
|
|
The Expr bound to the input var.
|
|
"""
|
|
return _ffi_api.BlockBuilderLookupBinding(self, var) # type: ignore
|
|
|
|
def begin_scope(self, params: list[Var] | None = None) -> None:
|
|
"""Begin a new scope, with optional parameters that
|
|
are visible within the scope.
|
|
|
|
Parameters
|
|
----------
|
|
params: Optional[List[Var]]
|
|
Parameters that are visible within the scope.
|
|
|
|
Note
|
|
----
|
|
This function should be called when new scope is introduced
|
|
(function, seq) to properly track the variable availability
|
|
and help the best effort deduction.
|
|
"""
|
|
|
|
return _ffi_api.BlockBuilderBeginScope(self, params) # type: ignore
|
|
|
|
def end_scope(self) -> None:
|
|
"""End the current scope. Please see `begin_scope` for details"""
|
|
|
|
return _ffi_api.BlockBuilderEndScope(self) # type: ignore
|