wehub-resource-sync
183 lines
7.2 KiB
Python
183 lines
7.2 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=abstract-method,unused-argument
|
|
# pylint: disable=missing-function-docstring,missing-module-docstring
|
|
|
|
import tvm
|
|
from tvm.ir import Call, Op
|
|
from tvm.ir.module import IRModule
|
|
|
|
from ..expr import Expr, Function, ShapeExpr
|
|
from ..expr_functor import PyExprVisitor, visitor
|
|
|
|
|
|
def estimate_memory_usage(mod: 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.tirx.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 "
|
|
f"{self.planned_alloc_mem / self.total_alloc_tensor_mem:.1%}."
|
|
)
|
|
return "- Function " + func_name + ":\n" + est
|
|
|
|
if isinstance(mod, Function):
|
|
mod = tvm.IRModule({tvm.ir.GlobalVar("foo"): mod})
|
|
|
|
return MemoryEstimator().estimate(mod)
|