r""" Memory planning algorithm suite that runs a list of memory planning algorithms and returns the result of the algorithm that minimizes the total memory usage. Args: graph_module: The graph module to allocate memory for alignment: Memory alignment
(
self,
alignment: int,
specs: Set[TensorSpec],
graph_module: torch.fx.GraphModule,
graph_signature: ExportGraphSignature,
extra_padding: int,
)
| 1038 | self.algo_list: List[Callable[..., MemoryAlgoResult]] = algo_list |
| 1039 | |
| 1040 | def __call__( |
| 1041 | self, |
| 1042 | alignment: int, |
| 1043 | specs: Set[TensorSpec], |
| 1044 | graph_module: torch.fx.GraphModule, |
| 1045 | graph_signature: ExportGraphSignature, |
| 1046 | extra_padding: int, |
| 1047 | ) -> List[int]: |
| 1048 | r""" |
| 1049 | Memory planning algorithm suite that runs a list of memory planning algorithms |
| 1050 | and returns the result of the algorithm that minimizes the total memory usage. |
| 1051 | |
| 1052 | Args: |
| 1053 | graph_module: The graph module to allocate memory for |
| 1054 | alignment: Memory alignment requirement |
| 1055 | graph_signature: Optional graph signature |
| 1056 | alloc_graph_input: Whether to allocate memory for graph input |
| 1057 | alloc_graph_output: Whether to allocate memory for graph output |
| 1058 | allow_overlapping_allocations: Whether to allow overlapping allocations |
| 1059 | algo_list: List of memory planning algorithms to run |
| 1060 | specs: Optional set of TensorSpec objects with updated lifetimes. If None, they will be |
| 1061 | calculated from the graph_module. |
| 1062 | |
| 1063 | Returns: |
| 1064 | List of buffer sizes for each memory hierarchy |
| 1065 | """ |
| 1066 | |
| 1067 | mem_algo_results = {} |
| 1068 | for algo in self.algo_list: |
| 1069 | if isinstance(algo, functools.partial): |
| 1070 | name = algo.func.__name__ |
| 1071 | else: |
| 1072 | name = getattr(algo, "__name__", None) |
| 1073 | |
| 1074 | mem_algo_results[name] = algo( |
| 1075 | alignment, |
| 1076 | specs, |
| 1077 | graph_module, |
| 1078 | graph_signature, |
| 1079 | extra_padding, |
| 1080 | ) |
| 1081 | |
| 1082 | # All the algorithms should have the same number of buffers allocated. |
| 1083 | assert ( |
| 1084 | len( |
| 1085 | { |
| 1086 | len(mem_algo_result.bufsizes) |
| 1087 | for mem_algo_result in mem_algo_results.values() |
| 1088 | } |
| 1089 | ) |
| 1090 | == 1 |
| 1091 | ), "Different memory planning algorithms should have the same number of buffers allocated." |
| 1092 | |
| 1093 | # Find the algorithm that minimizes the total memory usage. |
| 1094 | best_algo = min( |
| 1095 | mem_algo_results, key=lambda k: sum(mem_algo_results[k].bufsizes) |
| 1096 | ) |
| 1097 | logging.debug(f"Best memory planning algo for this model is {best_algo}") |