Implementation:Mlc ai Mlc llm Lift Global Buffer Alloc

Overview

LiftTIRGlobalBufferAlloc is a TVM compiler pass that lifts TIR-level global buffer allocations from inside TIR PrimFuncs up to the Relax level. When a TIR function internally allocates global-scope buffers (e.g., for intermediate results), this pass converts them into additional output parameters. The corresponding Relax-level call_tir invocations are updated to provide these buffers as additional outputs, allowing the Relax memory planner to manage their allocation and potentially reuse memory across function calls.

File: python/mlc_llm/compiler_pass/lift_global_buffer_alloc.py

Purpose

TIR PrimFuncs may contain alloc_buffer statements with "global" scope for workspace buffers. These allocations are invisible to the Relax-level memory planner, which means:

• Memory cannot be shared across different TIR function invocations
• The planner cannot account for these allocations in its optimization decisions

By lifting these allocations to the Relax level, the pass enables better memory reuse and more accurate memory planning.

Class: LiftTIRGlobalBufferAlloc

@tvm.transform.module_pass(opt_level=0, name="LiftTIRGlobalBufferAlloc")
class LiftTIRGlobalBufferAlloc:
    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
        return _TIRGlobalAllocRewriter(mod).transform()

Class: _TIRGlobalAllocRewriter

@mutator
class _TIRGlobalAllocRewriter(PyExprMutator):
    def __init__(self, mod: IRModule):
        super().__init__(mod)
        self.mod = mod
        self.gv2new_tensor_sinfo: Dict[
            tvm.ir.GlobalVar,
            Tuple[tvm.ir.GlobalVar, List[relax.TensorStructInfo], tir.PrimFunc],
        ] = {}

The gv2new_tensor_sinfo dictionary maps original TIR function global variables to tuples of:

• The new global variable (for the updated function)
• A list of TensorStructInfo for the lifted buffers
• The original function (for resolving symbolic variables)

transform

The transformation proceeds in two phases:

def transform(self) -> IRModule:
    # Phase 1: Rewrite TIR functions to remove global buffer allocations
    for g_var, func in self.mod.functions_items():
        if isinstance(func, tir.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)

    # Phase 2: Update Relax call_tir invocations
    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)

Phase 1 scans all TIR PrimFuncs and extracts global buffer allocations. Phase 2 updates all Relax functions to pass the lifted buffers as additional outputs. Finally, dead code elimination removes the now-unreferenced original TIR functions.

visit_call_

The Relax call rewriter updates call_tir calls targeting modified TIR functions:

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]

    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:
            self.gv2new_tensor_sinfo.pop(g_var)
            return call

When the lifted buffers have symbolic shapes, the pass must resolve TIR symbolic variables to Relax-level expressions. If resolution fails, the pass falls back to keeping the original (un-lifted) function.

The output struct info is updated to include the lifted buffers:

• For single-tensor output: wraps in a TupleStructInfo and returns TupleGetItem(result, 0)
• For tuple output: extends the existing tuple with additional tensor entries

Function: remove_global_buf_alloc

This is the core TIR-level transformation that converts global buffer allocations into function parameters:

def remove_global_buf_alloc(func: tir.PrimFunc) -> Tuple[tir.PrimFunc, List[relax.TensorStructInfo]]:
    params = list(func.params)
    buffer_map = dict(func.buffer_map)
    tensor_sinfo = []
    alloc_buffers = []

    # Find insertion point (after all handle params)
    insertion_point = len(params)
    while params[insertion_point - 1].dtype != "handle":
        insertion_point -= 1

    for buf_alloc in func.body.block.alloc_buffers:
        if buf_alloc.scope() == "global":
            param = tir.Var("var_" + buf_alloc.name, "handle")
            params.insert(insertion_point, param)
            insertion_point += 1
            buffer_map[param] = buf_alloc
            tensor_sinfo.append(relax.TensorStructInfo(buf_alloc.shape, buf_alloc.dtype))
        else:
            alloc_buffers.append(buf_alloc)

Key details:

• New parameters are inserted at a specific point: after all existing handle-type parameters but before any non-handle parameters
• Each global buffer becomes a new handle parameter with a "var_" prefix in its name
• Non-global scope buffers (e.g., "shared", "local") are preserved as internal allocations
• The root block is reconstructed without the lifted allocations

Function: _has_symbolic_var

Checks if a TensorStructInfo has any symbolic (non-constant) shape dimensions:

def _has_symbolic_var(tensor_sinfo: relax.TensorStructInfo) -> bool:
    for dim in tensor_sinfo.shape.values:
        if not isinstance(dim, tir.IntImm):
            return True
    return False

Function: _resolve_tir_var_mapping

Resolves TIR symbolic variables to Relax-level expressions by matching buffer shapes between the TIR function parameters and the Relax call arguments:

def _resolve_tir_var_mapping(func, call, tensor_sinfo):
    var_map: Dict[tir.Var, tir.PrimExpr] = {}

    # Map from TIR input buffer shapes to Relax argument shapes
    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].struct_info.shape.values
        for v_l, v_r in zip(buffer_shape, arg_shape):
            if isinstance(v_l, tir.Var):
                var_map[v_l] = v_r

    # Also map from output buffer shapes
    ret_tensors = call.sinfo_args[0]
    for i, ret_tensor in enumerate(ret_tensors):
        buffer_shape = func.buffer_map[func.params[n_arg + i]].shape
        # ... similar mapping ...

    # Substitute symbolic vars in lifted buffer shapes
    updated_tensor_sinfo = []
    for sinfo in tensor_sinfo:
        new_shape = [tir.stmt_functor.substitute(dim, var_map) for dim in sinfo.shape.values]
        updated_tensor_sinfo.append(relax.TensorStructInfo(new_shape, sinfo.dtype))
    return updated_tensor_sinfo, True

The function builds a mapping from TIR symbolic variables to concrete Relax shape expressions by examining both input and output buffer shapes. This mapping is then applied to the lifted buffer shapes using tir.stmt_functor.substitute.

Transformation Summary

Before (TIR)	After (TIR)
Global buffer allocated inside PrimFunc	Global buffer becomes a function parameter

Before (Relax)	After (Relax)
call_tir(func, args, out_sinfo=T)	TupleGetItem(call_tir(func_new, args, out_sinfo=(T, lifted_buf_1, ...)), 0)

Dependencies

• tvm -- Core TVM framework (relax, tir)
• tvm.relax.analysis.remove_all_unused -- Removes unused bindings
• tvm.relax.transform.DeadCodeElimination -- Removes unreferenced functions
• tvm.relax.expr_functor -- PyExprMutator and @mutator decorator