Skip to content

Make consumed lifted_constants as delegated nodes #4774

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 19, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions backends/qualcomm/partition/qnn_partitioner.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,18 +141,35 @@ def generate_partitions(
op_support=self.op_support_checker,
)

def tag_nodes(self, partitions: List[Partition]) -> None:
def tag_nodes(
self, partitions: List[Partition], edge_program: torch.export.ExportedProgram
) -> None:
for partition in partitions:
for node in partition.nodes:
delegation_tag = f"qnn_{partition.id}"
node.meta["delegation_tag"] = delegation_tag
self.partition_tags[delegation_tag] = self.delegation_spec

# need to take care of consumed constants
consumed_constants = (
*edge_program.graph_signature.inputs_to_buffers,
*edge_program.graph_signature.inputs_to_parameters,
)
for node in edge_program.graph_module.graph.nodes:
# find placeholders as lifted_constants
if node.op != "placeholder" or len(node.users) != 0:
continue

if node.name in consumed_constants:
# does no harm to merge them into last partition,
# since they will all be removed in following stage
node.meta["delegation_tag"] = delegation_tag

# override
def partition(self, edge_program: torch.export.ExportedProgram) -> PartitionResult:
partitions = self.generate_partitions(edge_program)
if len(partitions) != 0:
self.tag_nodes(partitions)
self.tag_nodes(partitions, edge_program)
tag_constant_data(edge_program)
for node in edge_program.graph_module.graph.nodes:
if hasattr(node, "meta"):
Expand Down
Loading