chore: import upstream snapshot with attribution

python/mlc_llm/compiler_pass/estimate_memory_usage.py

@@ -0,0 +1,87 @@
+"""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