| 223 | |
| 224 | class CustomPoolMemoryPlanningPass(MemoryPlanningPass): |
| 225 | def call(self, graph_module: GraphModule) -> PassResult: |
| 226 | for subgm in graph_module.modules(): |
| 227 | if not isinstance(subgm, GraphModule): |
| 228 | continue |
| 229 | for node in subgm.graph.nodes: |
| 230 | # mem_id = 1 placeholder and outputs of mul |
| 231 | # mem_id = 3 for outputs of add |
| 232 | # parent class will copy spec will to alloc nodes |
| 233 | if node.op == "placeholder": |
| 234 | node.meta["spec"].mem_id = 1 |
| 235 | continue |
| 236 | |
| 237 | if node.op != "call_function": |
| 238 | continue |
| 239 | |
| 240 | if node.target == torch.ops.aten.add.out: |
| 241 | node.meta["spec"].mem_id = 3 |
| 242 | elif node.target == torch.ops.aten.mul.out: |
| 243 | node.meta["spec"].mem_id = 1 |
| 244 | |
| 245 | return super().run(graph_module) |
| 246 | |
| 247 | def run( |
| 248 | self, |