def METHOD_NAME(mod: Union[IRModule, Function]) -> str:
    """Analysis function that estimates the memory usage of Relax functions
    in an IRModule. The estimation includes the total memory size needed to
    be allocated before and after memory planning.
    The result might be over-estimated, as the estimation is static, which
    does not consider control flows (such as "if" and cross-function calls).
    It simply accumulates the size of every alloc_tensor and alloc_storage.
    This analysis function is used to demonstrate the effect of memory
    planning.
    Parameters
    ----------
    mod : Union[IRModule, Function]
        The input IRModule whose functions inside are to be analyzed.
        If the input is a Function, we will wrap it with a IRModule, with
        the function named "main".
    Returns
    -------
    est : str
        The estimation information, in the form of a string.
    Notes
    -----
    We regards "relax.memory.alloc_tensor/storage" as the results produced by memory planning.
    """
    @visitor
    class MemoryEstimator(PyExprVisitor):
        """The IR visitor which estimates the memory usage of each Relax function.
        Attributes
        ----------
        total_alloc_tensor_mem : int
            The total memory size of alloc_tensor, in bytes.
        total_const_size_tensor_num : int
            The number of constant-size tensors.
        total_dyn_size_tensor_num : int
            The number of dynamic-size tensors.
        planned_alloc_mem : int
            The total memory size of memory.alloc_storage after memory planning, in bytes.
        planned_mem_num : int
            The number of memory.alloc_storages.
        """
        total_alloc_tensor_mem: int
        total_const_size_tensor_num: int
        total_dyn_size_tensor_num: int
        planned_alloc_mem: int
        planned_mem_num: int
        builtin_alloc_tensor_op = Op.get("relax.builtin.alloc_tensor")
        memory_alloc_tensor_op = Op.get("relax.memory.alloc_tensor")
        memory_alloc_storage_op = Op.get("relax.memory.alloc_storage")
        def estimate(self, mod: IRModule) -> str:
            estimation: str = ""
            for global_var, func in mod.functions.items():
                if not isinstance(func, Function):
                    continue
                self.cleanup()
                self.visit_expr(func)
                estimation += self.generate_est_string(global_var.name_hint) + "\n"
            if estimation != "":
                estimation = "Memory usage estimation:\n" + estimation
            return estimation
        def cleanup(self) -> None:
            self.total_alloc_tensor_mem = 0
            self.total_const_size_tensor_num = 0
            self.total_dyn_size_tensor_num = 0
            self.planned_alloc_mem = 0
            self.planned_mem_num = 0
        def visit_call_(self, call: Call) -> None:  # pylint: disable=arguments-differ
            if call.op == self.builtin_alloc_tensor_op:
                self.accumulate_builtin_tensor_alloc(
                    shape=call.args[0], dtype_str=call.args[1].value
                )
            elif call.op == self.memory_alloc_tensor_op:
                self.accumulate_tensor_alloc(shape=call.args[2], dtype_str=call.args[3].value)
            elif call.op == self.memory_alloc_storage_op:
                self.accumulate_storage_alloc(size=call.args[0])
        def calculate_size(self, shape: Expr, dtype_str: str) -> int:
            if not isinstance(shape, ShapeExpr):
                raise TypeError(
                    "The shape of relax.builtin.alloc_tensor and "
                    "relax.memory.alloc_tensor is expected to be ShapeExpr"
                )
            size: int = 1
            for dim_len in shape.values:
                if not isinstance(dim_len, tvm.tir.IntImm):
                    self.total_dyn_size_tensor_num += 1
                    return -1
                size *= dim_len.value
            dtype = tvm.DataType(dtype_str)
            return size * ((dtype.bits + 7) // 8) * dtype.lanes
        def accumulate_builtin_tensor_alloc(self, shape: Expr, dtype_str: str) -> None:
            size = self.calculate_size(shape, dtype_str)
            if size == -1:
                return
            self.total_const_size_tensor_num += 1
            self.total_alloc_tensor_mem += size
            self.planned_mem_num += 1
            self.planned_alloc_mem += size
        def accumulate_tensor_alloc(self, shape: Expr, dtype_str: str) -> None:
            size = self.calculate_size(shape, dtype_str)
            if size == -1:
                return
            self.total_const_size_tensor_num += 1
            self.total_alloc_tensor_mem += size
        def accumulate_storage_alloc(self, size: Expr) -> None:
            if not isinstance(size, ShapeExpr):
                raise TypeError(
                    "The size of relax.memory.alloc_storage is expected to be ShapeExpr"
                )
            self.planned_mem_num += 1
            self.planned_alloc_mem += size.values[0].value
        def generate_est_string(self, func_name: str) -> str:
            est = (
                f" * Without memory planning, there are {self.total_const_size_tensor_num} "
                "constant-size memory allocation(s) with total size "
                "{0:.4} GB".format(self.total_alloc_tensor_mem / 2**30)
            )
            if self.total_dyn_size_tensor_num > 0:
                est += f", and {self.total_dyn_size_tensor_num} dynamic-size allocation(s)"
            est += (
                f".\n * With memory planning, there are {self.planned_mem_num} constant-size "
                "memory allocation(s) with total size "
                "{0:.4} GB.\n".format(self.planned_alloc_mem / 2**30)
            )
            if self.total_alloc_tensor_mem != 0:
                est += " * Memory planning reduces constant memory size to " "{0:.1%}.".format(
                    self.planned_alloc_mem / self.total_alloc_tensor_mem
                )
            return "- Function " + func_name + ":\n" + est
    if isinstance(mod, Function):
        mod = tvm.IRModule({tvm.ir.GlobalVar("foo"): mod})
    return MemoryEstimator().estimate(mod)