wehub-resource-sync
88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
"""Memory usage estimation analysis function for Relax functions."""
|
|
|
|
import json
|
|
from typing import Any, Dict # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import relax, tirx
|
|
from tvm.ir import IRModule, Op
|
|
from tvm.relax.expr_functor import PyExprVisitor, visitor
|
|
|
|
from mlc_llm.support import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="AttachMetadata")
|
|
class AttachMetadataWithMemoryUsage:
|
|
"""Attach a Relax function that returns metadata in a JSON string"""
|
|
|
|
def __init__(self, metadata: Dict[str, Any]): # noqa: UP006
|
|
self.metadata = metadata
|
|
|
|
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
|
"""Entrypoint"""
|
|
|
|
func_name = "_metadata"
|
|
|
|
def _emit_metadata(metadata):
|
|
bb = relax.BlockBuilder()
|
|
with bb.function(func_name, params=[]):
|
|
bb.emit_func_output(relax.StringImm(json.dumps(metadata)))
|
|
return bb.finalize()[func_name]
|
|
|
|
self.metadata["memory_usage"] = _MemoryEstimator().run(mod)
|
|
mod[func_name] = _emit_metadata(self.metadata)
|
|
return mod
|
|
|
|
|
|
@visitor
|
|
class _MemoryEstimator(PyExprVisitor):
|
|
"""The IR visitor which estimates the memory usage of each Relax function."""
|
|
|
|
def __init__(self) -> None:
|
|
self.planned_alloc_mem = 0
|
|
self.planned_mem_num = 0
|
|
self._op_alloc_tensor = Op.get("relax.builtin.alloc_tensor")
|
|
self._op_alloc_storage = Op.get("relax.memory.alloc_storage")
|
|
|
|
def run(self, mod: IRModule) -> Dict[str, int]: # noqa: UP006
|
|
"""Entry point of the visitor."""
|
|
result: Dict[str, int] = {} # noqa: UP006
|
|
for global_var, func in mod.functions_items():
|
|
if isinstance(func, relax.Function):
|
|
self.planned_alloc_mem = 0
|
|
self.planned_mem_num = 0
|
|
self.visit_expr(func)
|
|
result[global_var.name_hint] = self.planned_alloc_mem
|
|
logger.info(
|
|
"[Memory usage] Function `%s`: %.2f MB",
|
|
global_var.name_hint,
|
|
self.planned_alloc_mem / 1024 / 1024,
|
|
)
|
|
return result
|
|
|
|
def visit_call_(self, call: relax.Call) -> None:
|
|
if call.op == self._op_alloc_tensor:
|
|
self._builtin_tensor_alloc(shape=call.args[0], dtype_str=call.args[1].value)
|
|
elif call.op == self._op_alloc_storage:
|
|
self._storage_alloc(size=call.args[0])
|
|
super().visit_call_(call)
|
|
|
|
def _builtin_tensor_alloc(self, shape: relax.Expr, dtype_str: str) -> None:
|
|
assert isinstance(shape, relax.ShapeExpr)
|
|
size = 1
|
|
for dim_len in shape.values:
|
|
if not isinstance(dim_len, tvm.tirx.IntImm):
|
|
return
|
|
size *= dim_len.value
|
|
dtype = tvm.DataType(dtype_str)
|
|
self.planned_mem_num += 1
|
|
self.planned_alloc_mem += size * ((dtype.bits + 7) // 8) * dtype.lanes
|
|
|
|
def _storage_alloc(self, size: relax.Expr) -> None:
|
|
assert isinstance(size, relax.ShapeExpr)
|
|
if isinstance(size.values[0], tirx.IntImm):
|
|
self.planned_mem_num += 1
|
|
self.planned_alloc_mem += size.values[0].value
|