wehub-resource-sync
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
"""A compiler pass that scatters TupleGetItem for lazy TupleGetItems."""
|
|
|
|
from typing import Dict # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.ir.module import IRModule
|
|
from tvm.relax.analysis import remove_all_unused
|
|
from tvm.relax.expr import Expr, Var
|
|
from tvm.relax.expr_functor import PyExprMutator, mutator
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="ScatterTupleGetItem")
|
|
class ScatterTupleGetItem:
|
|
"""A compiler pass that scatters TupleGetItem for lazy TupleGetItems."""
|
|
|
|
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
|
"""IRModule-level transformation"""
|
|
return _Scatter(mod).transform()
|
|
|
|
|
|
@mutator
|
|
class _Scatter(PyExprMutator):
|
|
def __init__(self, mod: IRModule) -> None:
|
|
super().__init__(mod)
|
|
self.mod = mod
|
|
self.var_map: Dict[Var, Expr] = {} # noqa: UP006
|
|
|
|
def transform(self) -> IRModule:
|
|
"""Entry point"""
|
|
for g_var, func in self.mod.functions_items():
|
|
if isinstance(func, relax.Function):
|
|
updated_func = self.visit_expr(func)
|
|
updated_func = remove_all_unused(updated_func)
|
|
self.builder_.update_func(g_var, updated_func)
|
|
return self.builder_.get()
|
|
|
|
def visit_var_binding_(self, binding: relax.VarBinding):
|
|
super().visit_var_binding_(binding)
|
|
if isinstance(binding.value, relax.TupleGetItem):
|
|
self.var_map[binding.var] = binding.value
|
|
|
|
def visit_dataflow_var_(self, var: relax.DataflowVar) -> Expr:
|
|
if var in self.var_map:
|
|
new_var = self.builder_.emit(self.var_map[var], name_hint=var.name_hint)
|
|
self.set_var_remap(var, new_var)
|
|
self.var_map.pop(var)
|
|
return new_var
|
|
return var
|