Check if a node is marked with a given field in node.meta["custom"]
| 56 | |
| 57 | |
| 58 | class NodeFlagIsSetChecker(OperatorSupportBase): |
| 59 | """ |
| 60 | Check if a node is marked with a given field in node.meta["custom"] |
| 61 | """ |
| 62 | |
| 63 | def __init__(self, field: str) -> None: |
| 64 | super().__init__() |
| 65 | self.field = field |
| 66 | |
| 67 | def check_field(self, node: torch.fx.Node) -> bool: |
| 68 | if "custom" not in node.meta: |
| 69 | return False |
| 70 | |
| 71 | custom_map = node.meta["custom"] |
| 72 | if self.field not in custom_map: |
| 73 | return False |
| 74 | |
| 75 | return custom_map[self.field] |
| 76 | |
| 77 | def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: |
| 78 | if node.op == "placeholder" or node.op == "output": |
| 79 | return False |
| 80 | |
| 81 | # Check if the node itself is tagged |
| 82 | if self.check_field(node): |
| 83 | return True |
| 84 | |
| 85 | # Check if any direct user of this node is tagged |
| 86 | for user in node.users: |
| 87 | if self.check_field(user): |
| 88 | return True |
| 89 | |
| 90 | return False |
| 91 | |
| 92 | |
| 93 | class FlagBasedPartitioner(Partitioner): |