If the constants node is shared by different tagged nodes, like constant_0 ----> op_b (tag_10) |-------------> op_a (tag_11) we make default as constant_0 is duplicated to constant_0_1, constant_0_2, unless the node is tagged with "no_copy" constant_0 ------------------> op_b (
(
tagged_exported_program: ExportedProgram,
tag: str,
)
| 208 | |
| 209 | |
| 210 | def _maybe_duplicate_constant_nodes( |
| 211 | tagged_exported_program: ExportedProgram, |
| 212 | tag: str, |
| 213 | ) -> None: |
| 214 | """ |
| 215 | If the constants node is shared by different tagged nodes, like |
| 216 | constant_0 ----> op_b (tag_10) |
| 217 | |-------------> op_a (tag_11) |
| 218 | |
| 219 | we make default as constant_0 is duplicated to constant_0_1, constant_0_2, unless the node is tagged with "no_copy" |
| 220 | constant_0 ------------------> op_b (tag_10) |
| 221 | constant_0_copy -------------> op_a (tag_11) |
| 222 | |
| 223 | backend can estimate how much they want to duplicate the constant node, either error out or default to duplicate |
| 224 | """ |
| 225 | candidate_nodes = set() |
| 226 | for node in tagged_exported_program.graph.nodes: |
| 227 | if node.meta.get("delegation_tag", "") == tag: |
| 228 | if node.op == "placeholder": |
| 229 | for user in node.users: |
| 230 | users_tag = user.meta.get("delegation_tag", None) |
| 231 | if users_tag != tag: |
| 232 | # If the node is tagged with "no_copy", we stop duplicating it and throw an error |
| 233 | if node.meta.get("no_copy", False): |
| 234 | raise RuntimeError( |
| 235 | f"constant data node ({node}) is tagged with ({tag}) but has user ({user}) which has tag ({users_tag})" |
| 236 | ) |
| 237 | else: |
| 238 | candidate_nodes.add(node.name) |
| 239 | copied_nodes = set() |
| 240 | for candidate_node in candidate_nodes: |
| 241 | # Both tagged exported program and the owning program need to go through the same duplication pass |
| 242 | copied_nodes = copied_nodes.union( |
| 243 | duplicate_constant_node(tagged_exported_program, candidate_node) |
| 244 | ) |
| 245 | candidate_node_with_copies = candidate_nodes.union(copied_nodes) |
| 246 | _assign_new_tag(tagged_exported_program, candidate_node_with_copies) |
| 247 | |
| 248 | |
| 249 | def _get_item_from_executorch_call_delegate(node: torch.fx.Node) -> bool: |
no test coverage detected