# 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, unused-argument, missing-function-docstring, abstract-method, missing-class-docstring # ruff: noqa: RUF005 """Relax LazyTransformParams pass.""" import tvm from tvm import IRModule, relax from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor @visitor class ForwardCollector(PyExprVisitor): """ Perform a forward pass to collect the following information: out_tuple_map: map from var to its index in the output tuple var_tuple_get_item: list of var that is bound to v = params[i] Parameters ---------- tuple_var: relax.Var The output tuple var input_params: relax.Var The input tuple var """ def __init__(self, tuple_var: relax.Var, input_params: relax.Var) -> None: self.out_tuple_map = {} self.out_tuple_var = tuple_var self.input_params = input_params self.var_tuple_get_item = [] self.is_tuple_get_item_input = False def visit_tuple_getitem_(self, op: relax.TupleGetItem) -> None: if op.tuple_value == self.input_params: self.is_tuple_get_item_input = True else: self.is_tuple_get_item_input = False super().visit_tuple_getitem_(op) def visit_var_binding_(self, binding: relax.VarBinding) -> None: if binding.var == self.out_tuple_var: assert isinstance(binding.value, relax.Tuple) for i, expr in enumerate(binding.value.fields): if expr not in self.out_tuple_map: self.out_tuple_map[expr] = [] self.out_tuple_map[expr].append(relax.prim_value(i)) else: self.is_tuple_get_item_input = False super().visit_var_binding_(binding) if self.is_tuple_get_item_input: self.var_tuple_get_item.append(binding.var) @visitor class LivenessAnalysis(PyExprVisitor): """ Perform a backward pass to collect the following information: var_liveness_end: map from var to the list of var whose liveness is killed by this var binding Parameters ---------- out_tuple_var: relax.Var The output tuple var input_params: set The set of vars that are bound to v = params[i] """ def __init__(self, out_tuple_var: relax.Var) -> None: self.last_appear_in_var_binding = [] self.out_tuple_var = out_tuple_var self.var_liveness_end = {} self.ended_vars = set() def visit_binding_block_(self, block: relax.BindingBlock) -> None: for binding in reversed(block.bindings): self.visit_binding(binding) def visit_var_(self, op: relax.Var) -> None: if op not in self.ended_vars: self.last_appear_in_var_binding.append(op) self.ended_vars.add(op) def visit_var_binding_(self, binding: relax.VarBinding) -> None: if self.out_tuple_var == binding.var: return self.last_appear_in_var_binding = [] super().visit_var_binding_(binding) # param[i] is in output if binding.var not in self.ended_vars: self.last_appear_in_var_binding.append(binding.var) self.ended_vars.add(binding.var) self.var_liveness_end[binding.var] = self.last_appear_in_var_binding class LazyTransformParamsFuncCreator: """ Transform transform_params functions into a lazy version. Parameters ---------- mod: IRModule The module to be transformed """ def __init__( self, fget_item, fset_item, extra_get_item_params, extra_set_item_params, mod: IRModule = None, ) -> None: self.mod = mod self.fget_item = fget_item self.extra_get_item_params = extra_get_item_params self.fset_item = fset_item self.extra_set_item_params = extra_set_item_params self.input_params_set = None self.out_tuple_map = None self.out_tuple_var = None self.memory_free_insertion = None def transform(self, func: relax.Function) -> relax.Function: if "num_input" in func.attrs: num_input = func.attrs["num_input"] else: num_input = 0 seq_expr = func.body self.out_tuple_var = seq_expr.body # Step 1. collect out_tuple_map and input_params_set forward_collector = ForwardCollector(self.out_tuple_var, func.params[num_input]) forward_collector.visit_expr(func) self.out_tuple_map = forward_collector.out_tuple_map # input_params_set is the set of binding var for var = params[i] self.input_params_set = set(forward_collector.var_tuple_get_item) # Step 2. liveness analysis and get where to insert kill_object instruction liveness = LivenessAnalysis(self.out_tuple_var) liveness.visit_expr(func) self.memory_free_insertion = liveness.var_liveness_end # Step 3. rewrite get item and set item if self.fget_item is not None: new_func = LazyInputMutator(self, self.mod).visit_expr(func) new_body = new_func.body if self.fset_item is not None: # The LazyOutputMutator only inspects variable bindings # for replacement. If the output tuple includes elements # that do not have a variable binding, such as # `relax.Const`, these must still produce a call to the # `"set_item"` function. leaf_outputs = { expr: indices for expr, indices in self.out_tuple_map.items() if not isinstance(expr, relax.Var) } if leaf_outputs: new_bindings = [ relax.VarBinding( relax.Var("_", relax.AnyType()), relax.Call( relax.ExternFunc(self.fset_item), [*self.extra_set_item_params, index, expr], None, [relax.AnyType()], ), ) for expr, indices in leaf_outputs.items() for index in indices ] new_body = relax.SeqExpr( [*new_body.blocks, relax.BindingBlock(new_bindings)], new_body.body ) new_body = LazyOutputMutator(self, self.mod).visit_expr(new_body) # Step 4. Add parameters of get_item and set_item (except index) to the function. params = [ *func.params[:num_input], *self.extra_get_item_params, *self.extra_set_item_params, ] # Step 5. Find all shape parameters that should be retained as # parameters. symbolic_vars = relax.analysis.defined_symbolic_vars(func) if symbolic_vars: def unpack_ty(ty): if isinstance(ty, relax.TupleType): for field in ty.fields: yield from unpack_ty(field) else: yield ty # direct iterate over the type annotation for param in func.params[num_input:]: for ty in unpack_ty(param.ty): if isinstance(ty, tvm.ir.PrimType | relax.ShapeType): params.append(relax.Var("symbolic_var_holder", ty)) return relax.Function( params, new_body, relax.AnyType(), attrs=func.attrs, is_pure=False, ).without_attr("relax.force_pure") @mutator class LazyInputMutator(PyExprMutator): def __init__(self, func_creator, mod: IRModule | None = None) -> None: self.func_creator = func_creator super().__init__(mod) def visit_function_(self, func: relax.Function) -> relax.Expr: if "num_input" in func.attrs: num_input = func.attrs["num_input"] else: num_input = 0 params = list(func.params)[num_input:] if len(params) == 1 and isinstance(params[0].ty, relax.TupleType): self.tuple_param = params[0] self.params = {} else: self.tuple_param = None self.params = {var: i for i, var in enumerate(params)} func = relax.Function( func.params[:num_input], func.body, func.ret_ty, is_pure=False, attrs=func.attrs, span=func.span, ).without_attr("relax.force_pure") output = super().visit_function_(func) self.tuple_param = None self.params = {} return output def visit_var_(self, var: relax.Var) -> relax.Expr: if var in self.params: index = self.params[var] get_item_result = self.builder_.emit( relax.Call( relax.ExternFunc(self.func_creator.fget_item), self.func_creator.extra_get_item_params + [relax.prim_value(index)], None, [relax.AnyType()], ) ) match_cast = relax.MatchCast(var, get_item_result, var.ty) self.builder_.emit_normalized(match_cast) del self.params[var] return super().visit_var_(var) def visit_tuple_getitem_(self, node: relax.TupleGetItem) -> relax.Expr: ty = node.ty node = super().visit_tuple_getitem_(node) if self.tuple_param is not None and node.tuple_value.same_as(self.tuple_param): get_item_result = self.builder_.emit( relax.Call( relax.ExternFunc(self.func_creator.fget_item), self.func_creator.extra_get_item_params + [relax.prim_value(node.index)], None, [relax.AnyType()], ) ) return self.builder_.match_cast(get_item_result, ty) else: return node @mutator class LazyOutputMutator(PyExprMutator): def __init__(self, func_creator, mod: IRModule | None = None) -> None: self.func_creator = func_creator self.killed_vars = set() super().__init__(mod) def visit_var_(self, var: relax.Var) -> None: assert var not in self.killed_vars return super().visit_var_(var) def visit_var_binding_(self, binding: relax.VarBinding) -> None: if binding.var == self.func_creator.out_tuple_var: # The function after rewriting returns a empty tuple. func_output = self.builder_.emit(relax.Tuple([])) self.set_var_remap(binding.var, func_output) return super().visit_var_binding_(binding) if binding.var in self.func_creator.memory_free_insertion: for var in self.func_creator.memory_free_insertion[binding.var]: if var in self.func_creator.out_tuple_map: self.killed_vars.add(var) for index in self.func_creator.out_tuple_map[var]: # rewrite set item self.builder_.emit( relax.Call( relax.ExternFunc(self.func_creator.fset_item), self.func_creator.extra_set_item_params + [index, super().visit_var_(var)], None, [relax.AnyType()], ), name_hint="_", ) if var in self.func_creator.input_params_set: self.builder_.emit( relax.op.vm.kill_object(super().visit_var_(var)), name_hint="_" ) @tvm.transform.module_pass(opt_level=0, name="LazyTransformParams") class LazyTransformParams: """ Convert transform_params functions into a lazy version. (Load the input to memory on demand, and immediately free it after the last use.) Note: ToNonDataflow() and RemovePurityTracking() should be invoked before this pass. Parameters ---------- fget_item: str The name of the get_item function. fset_item: str The name of the set_item function. extra_get_item_params: list of relax.Var The parameters of the get_item function except index. The given parameters will be placed before index. For example, if extra_get_item_params is [param1, param2], then the pass will generate call_packed(fget_item, [param1, param2, index]) extra_set_item_params: list of relax.Var The parameters of the set_item function except index and value. The given parameters will be placed before index and value. For example, if extra_set_item_params is [param1, param2], then the pass will generate call_packed(fset_item, [param1, param2, index, value]) """ def __init__( self, fget_item="get_item", fset_item="set_item", extra_get_item_params=None, extra_set_item_params=None, ) -> None: self.fget_item = fget_item self.extra_get_item_params = [] if extra_get_item_params is None else extra_get_item_params assert self.fget_item is not None, "transforming set_item only is not supported" self.fset_item = fset_item self.extra_set_item_params = [] if extra_set_item_params is None else extra_set_item_params def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule: lazy_mutator = LazyTransformParamsFuncCreator( self.fget_item, self.fset_item, self.extra_get_item_params, self.extra_set_item_params, mod, ) builder = relax.BlockBuilder(mod) for gv, _ in mod.functions_items(): if gv.name_hint.endswith("transform_params"): func = mod[gv] if not isinstance(func, relax.Function): continue func = lazy_mutator.transform(func) builder.update_func(gv, func) return builder.get()