"""A compiler pass that lifts TIR-level global allocation to Relax.""" from typing import Dict, List, Tuple # noqa: UP035 import tvm from tvm import relax, tirx from tvm.ir.module import IRModule from tvm.relax.analysis import remove_all_unused from tvm.relax.expr_functor import PyExprMutator, mutator @tvm.transform.module_pass(opt_level=0, name="LiftTIRGlobalBufferAlloc") class LiftTIRGlobalBufferAlloc: """A compiler pass that lifts TIR-level global allocation to Relax.""" def transform_module( self, mod: IRModule, _ctx: tvm.transform.PassContext, ) -> IRModule: """IRModule-level transformation""" return _TIRGlobalAllocRewriter(mod).transform() @mutator class _TIRGlobalAllocRewriter(PyExprMutator): def __init__(self, mod: IRModule): super().__init__(mod) self.mod = mod self.gv2new_tensor_sinfo: Dict[ # noqa: UP006 tvm.ir.GlobalVar, Tuple[tvm.ir.GlobalVar, List[relax.TensorType], tirx.PrimFunc], # noqa: UP006 ] = {} def transform(self) -> IRModule: """Entry point of the transformation""" for g_var, func in self.mod.functions_items(): if isinstance(func, tirx.PrimFunc): updated_func, tensor_sinfo_list = remove_global_buf_alloc(func) if len(tensor_sinfo_list) > 0: new_gv = self.builder_.add_func(updated_func, g_var.name_hint) self.gv2new_tensor_sinfo[g_var] = (new_gv, tensor_sinfo_list, func) self.mod = self.builder_.get() 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) mod = self.builder_.get() return relax.transform.DeadCodeElimination()(mod) def visit_call_(self, call: relax.Call): call = self.visit_expr_post_order(call) if ( call.op != tvm.ir.Op.get("relax.call_tir") or call.args[0] not in self.gv2new_tensor_sinfo ): return call g_var = call.args[0] new_gv, tensor_sinfo, func_before_update = self.gv2new_tensor_sinfo[g_var] assert len(call.ty_args) == 1 if any(_has_symbolic_var(sinfo) for sinfo in tensor_sinfo): tensor_sinfo, success = _resolve_tir_var_mapping(func_before_update, call, tensor_sinfo) if not success: # Cannot resolve TIR var mapping. Fall back to no lifting. self.gv2new_tensor_sinfo.pop(g_var) return call args = list(call.args) args[0] = new_gv if isinstance(call.ty_args[0], relax.TensorType): new_call = relax.Call( call.op, args=args, ty_args=[relax.TupleType(list(call.ty_args) + tensor_sinfo)], attrs=call.attrs, ) emitted_tuple = self.builder_.emit(new_call) return relax.TupleGetItem(emitted_tuple, 0) assert isinstance(call.ty_args[0], relax.TupleType) return relax.Call( call.op, args=args, ty_args=[relax.TupleType(list(call.ty_args[0].fields) + tensor_sinfo)], attrs=call.attrs, ) def remove_global_buf_alloc( func: tirx.PrimFunc, ) -> Tuple[tirx.PrimFunc, List[relax.TensorType]]: # noqa: UP006 """Remove the global buffer allocation for a given TIR PrimFunc.""" assert isinstance(func.body, tirx.SBlockRealize) params = list(func.params) buffer_map = dict(func.buffer_map) tensor_sinfo = [] alloc_buffers = [] insertion_point = len(params) while not isinstance(params[insertion_point - 1].ty, tvm.ir.PointerType): insertion_point -= 1 assert insertion_point >= 1 prev_root_block = func.body.block for buf_alloc in func.body.block.alloc_buffers: if buf_alloc.scope() == "global": param = tirx.Var("var_" + buf_alloc.name, "handle") params.insert(insertion_point, param) insertion_point += 1 buffer_map[param] = buf_alloc tensor_sinfo.append(relax.TensorType(buf_alloc.shape, buf_alloc.dtype)) else: alloc_buffers.append(buf_alloc) if len(tensor_sinfo) == 0: return func, [] assert len(prev_root_block.iter_vars) == 0 assert len(prev_root_block.reads) == 0 assert len(prev_root_block.writes) == 0 assert len(prev_root_block.match_buffers) == 0 assert prev_root_block.name_hint == "root" assert prev_root_block.init is None root_block = tirx.SBlock( iter_vars=[], reads=[], writes=[], name_hint="root", body=prev_root_block.body, alloc_buffers=alloc_buffers, annotations=prev_root_block.annotations, ) updated_func = tirx.PrimFunc( params=params, body=tirx.SBlockRealize(iter_values=[], predicate=True, block=root_block), ret_type=func.ret_type, buffer_map=buffer_map, attrs=func.attrs, ) return updated_func, tensor_sinfo def _has_symbolic_var(tensor_sinfo: relax.TensorType) -> bool: assert isinstance(tensor_sinfo.shape, relax.ShapeExpr) for dim in tensor_sinfo.shape.values: if not isinstance(dim, tirx.IntImm): return True return False def _resolve_tir_var_mapping( func: tirx.PrimFunc, call: relax.Call, tensor_sinfo: List[relax.TensorType], # noqa: UP006 ) -> Tuple[List[relax.TensorType], bool]: # noqa: UP006 """Resolve the TIR symbolic var relationship across sides of PrimFunc and Relax Function""" var_map: Dict[tirx.Var, tirx.Expr] = {} # noqa: UP006 n_arg = len(call.args[1].fields) for i in range(n_arg): buffer_shape = func.buffer_map[func.params[i]].shape arg_shape = call.args[1][i].ty.shape.values assert len(buffer_shape) == len(arg_shape) for v_l, v_r in zip(buffer_shape, arg_shape): if isinstance(v_l, tirx.Var): var_map[v_l] = v_r elif not isinstance(v_l, tirx.IntImm): return [], False ret_tensors = call.ty_args[0] ret_tensors = ( [ret_tensors] if isinstance(ret_tensors, relax.TensorType) else list(ret_tensors.fields) ) for i, ret_tensor in enumerate(ret_tensors): buffer_shape = func.buffer_map[func.params[n_arg + i]].shape ret_tensor_shape = ret_tensor.shape.values assert len(buffer_shape) == len(ret_tensor_shape) for v_l, v_r in zip(buffer_shape, ret_tensor_shape): if isinstance(v_l, tirx.Var): var_map[v_l] = v_r elif not isinstance(v_l, tirx.IntImm): return [], False updated_tensor_sinfo = [] for sinfo in tensor_sinfo: if not _has_symbolic_var(sinfo): updated_tensor_sinfo.append(sinfo) continue new_shape = [] for dim in sinfo.shape.values: new_shape.append(tirx.stmt_functor.substitute(dim, var_map)) updated_tensor_sinfo.append(relax.TensorType(new_shape, sinfo.dtype)) return updated_tensor_sinfo, True