Partitioner that tags add ops for delegation with target_device=cuda:0.
| 46 | |
| 47 | @final |
| 48 | class _DeviceAwarePartitioner(Partitioner): |
| 49 | """Partitioner that tags add ops for delegation with target_device=cuda:0.""" |
| 50 | |
| 51 | def __init__(self) -> None: |
| 52 | super().__init__() |
| 53 | self.delegation_spec = DelegationSpec( |
| 54 | BackendWithCompilerDemo.__name__, |
| 55 | [ |
| 56 | CompileSpec("max_value", bytes([4])), |
| 57 | CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, b"cuda:0"), |
| 58 | ], |
| 59 | ) |
| 60 | |
| 61 | def partition(self, exported_program) -> PartitionResult: |
| 62 | partition_tags: Dict[str, DelegationSpec] = {} |
| 63 | partition_list = generate_pattern_op_partitions( |
| 64 | exported_program.graph_module, |
| 65 | op_support=any_chain(_AddOperatorSupport()), |
| 66 | ) |
| 67 | for partition in partition_list: |
| 68 | for node in partition.nodes: |
| 69 | tag = f"tag{partition.id}" |
| 70 | node.meta["delegation_tag"] = tag |
| 71 | partition_tags[tag] = self.delegation_spec |
| 72 | return PartitionResult( |
| 73 | tagged_exported_program=exported_program, |
| 74 | partition_tags=partition_tags, |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | class ModuleAddWithDevice(nn.Module): |