[Analysis] Memory usage estimation

tlc-pack / relax

Apache License 2.0

193 stars 58 forks source link

[Analysis] Memory usage estimation #412

Closed MasterJH5574 closed 1 year ago

MasterJH5574 commented 1 year ago

This PR brings an analysis function which estimates the memory usage of Relax functions, used for demonstrating the effect of memory planning.

The estimation will return the total memory size needed to be allocated before and after memory planning, as well as the number of tensors / memory blocks to be allocated before and after memory planning.

Note:

The estimation is static -- it does not consider control flows (such as “if” and cross-function calls). It simply accumulates the size of every alloc_tensor and alloc_storage.
We regards “relax.memory.alloc_tensor/storage” as the results produced by memory planning.

slyubomirsky commented 1 year ago

Perhaps the docstring should describe what conditions the estimation expects? Not all programs will be possible to statically estimate (e.g., statically unbounded recursion).

Additionally, it looks like it would have to handle branching as well (or ask that the input programs not include branches).

MasterJH5574 commented 1 year ago

Perhaps the docstring should describe what conditions the estimation expects? Not all programs will be possible to statically estimate (e.g., statically unbounded recursion).

Additionally, it looks like it would have to handle branching as well (or ask that the input programs not include branches).

Thanks for bringing up! I didn’t consider those cases when writing it. Will revise later when getting time :-)

MasterJH5574 commented 1 year ago

Updated doc string as well as the PR description -- mainly stating that the analysis is static and doesn’t consider control flows.

The PR is now ready for review.