diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 8db23a465ef..77f99ddd1d5 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -41,6 +41,7 @@ PartitionResult, ) from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer +from executorch.exir.dialects._ops import ops as exir_ops from torch.export.exported_program import ExportedProgram from torch.fx.passes.infra.partitioner import Partition from torch.fx.passes.operator_support import OperatorSupportBase @@ -113,7 +114,7 @@ def __init__( for soc_model in target_socs ] - def is_node_supported(self, _, node: torch.fx.Node) -> bool: + def is_node_supported(self, _, node: torch.fx.Node) -> bool: # noqa: C901 if node.op != "call_function" or node.target in not_supported_operator: return False @@ -168,6 +169,20 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool: node, self.soc_info.htp_info.vtcm_size_in_mb, self.phase ) + if node.target in constant_operator: + # Built as a static tensor, so every value it depends on must be known AoT + # (full_like's first argument only supplies the shape). + shape_only = ( + node.args[0] + if node.target == exir_ops.edge.aten.full_like.default + else None + ) + if any(arg is not shape_only for arg in node.all_input_nodes): + logger.info( + f"[{self.phase}] {node.target.__name__} | Skipped, value is only known at runtime" + ) + return False + supported = False if node.target.__name__ not in self.node_visitors: logger.info( diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 606fed01d72..2fb17f6f6e1 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -783,6 +783,42 @@ def forward(self, x): if "QNN" in str(e) or "qnn" in str(e): self.skipTest(f"QNN SDK not available: {e}") + def test_full_with_runtime_fill_value_stays_on_cpu(self): + """aten.full is lowered as a static tensor, so a fill_value only known at runtime + (YOLO's `torch.full(..., stride.item())`) must leave it on CPU rather than crash + the partitioner's support check. + """ + + class FullFromTensor(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("stride", torch.tensor([8.0, 16.0])) + + def forward(self, x): + return x * torch.full((4, 1), self.stride[0].item()) + + sample_input = (torch.randn(4, 3),) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + try: + edge = to_edge_transform_and_lower_to_qnn( + FullFromTensor(), sample_input, compiler_specs + ) + except RuntimeError as e: + if "QNN" in str(e) or "qnn" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + + self.assertTrue( + any( + node.op == "call_function" and "aten.full" in str(node.target) + for node in edge.exported_program().graph.nodes + ) + ) + edge.to_executorch() + def test_lowering_accepts_exported_program_without_reexport(self): """to_edge_transform_and_lower_to_qnn must lower a caller-supplied ExportedProgram as captured, never re-exporting it.