Takes an ExportedProgram and returns the ExportedProgram modified in-place, with the constant tensors as buffers.
(ep)
| 324 | |
| 325 | |
| 326 | def lift_constant_tensor_pass(ep): |
| 327 | """ |
| 328 | Takes an ExportedProgram and returns the ExportedProgram modified in-place, |
| 329 | with the constant tensors as buffers. |
| 330 | """ |
| 331 | if len([node for node in ep.graph.nodes if node.op == "placeholder"]) == 0: |
| 332 | return ep |
| 333 | |
| 334 | graph_signature = ep.graph_signature |
| 335 | buffers = list(graph_signature.buffers) |
| 336 | |
| 337 | fake_mode = _detect_fake_mode_from_gm(ep.graph_module) |
| 338 | |
| 339 | first_user_input = None |
| 340 | lifted_constants = [] |
| 341 | for node in ep.graph.nodes: |
| 342 | if node.op == "placeholder" and node.name in graph_signature.user_inputs: |
| 343 | first_user_input = node |
| 344 | break |
| 345 | |
| 346 | for node in ep.graph.nodes: |
| 347 | if node.op == "get_attr": |
| 348 | constant_tensor = getattr(ep.graph_module, node.target) |
| 349 | if not isinstance(constant_tensor, torch.Tensor): |
| 350 | continue |
| 351 | |
| 352 | constant_tensor_fqn = f"_lifted_tensor_constant{len(buffers)}" |
| 353 | |
| 354 | with ep.graph.inserting_before(first_user_input): |
| 355 | # Insert the constant node before the first user input |
| 356 | const_placeholder_node = ep.graph.placeholder(constant_tensor_fqn) |
| 357 | for k, v in node.meta.items(): |
| 358 | const_placeholder_node.meta[k] = v |
| 359 | if fake_mode is not None: |
| 360 | const_placeholder_node.meta["val"] = fake_mode.from_tensor( |
| 361 | constant_tensor, static_shapes=True |
| 362 | ) |
| 363 | else: |
| 364 | const_placeholder_node.meta["val"] = constant_tensor |
| 365 | const_placeholder_node.meta["val"].constant = constant_tensor |
| 366 | node.replace_all_uses_with(const_placeholder_node) |
| 367 | ep.graph.erase_node(node) |
| 368 | |
| 369 | # Add the constant as a buffer to the graph signature |
| 370 | lifted_constants.append( |
| 371 | InputSpec( |
| 372 | kind=InputKind.BUFFER, |
| 373 | arg=TensorArgument(name=const_placeholder_node.name), |
| 374 | target=constant_tensor_fqn, |
| 375 | persistent=True, |
| 376 | ) |
| 377 | ) |
| 378 | buffers.append(constant_tensor_fqn) |
| 379 | ep.state_dict[constant_tensor_fqn] = constant_tensor |
| 380 | |
| 381 | new_input_specs = [] |
| 382 | for s in graph_signature.input_specs: |
| 383 | if s.kind == InputKind.USER_INPUT and len(lifted_constants) > 0: |