diff --git a/backends/qualcomm/_passes/annotate_unbind.py b/backends/qualcomm/_passes/annotate_unbind.py index ec03dc517..c67319cfe 100644 --- a/backends/qualcomm/_passes/annotate_unbind.py +++ b/backends/qualcomm/_passes/annotate_unbind.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import torch -from executorch.backends.qualcomm.builders.node_visitor import dq_ops +from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_ops from executorch.backends.qualcomm.utils.constants import QCOM_QUANT_ATTRS from executorch.exir.pass_base import ExportPass, PassResult from torch.fx.passes.utils.source_matcher_utils import get_source_partitions @@ -30,8 +30,12 @@ def _annotate_unbind(self, graph_module: torch.fx.GraphModule): ) for src_partitions in partitions.values(): for src_partition in src_partitions: - if src_partition.input_nodes[0].target in dq_ops: - q_node = src_partition.input_nodes[0].args[0] + dq_node = src_partition.input_nodes[0] + if dq_node.target in dq_ops: + # a folded quantized constant feeds its dequantize directly + q_node = dq_node.args[0] + if q_node.target not in q_ops: + q_node = dq_node quant_attrs = get_quant_attrs(self.edge_program, q_node) for n in src_partition.nodes: n.meta[QCOM_QUANT_ATTRS] = quant_attrs.copy() diff --git a/backends/qualcomm/_passes/remove_0d_tensor.py b/backends/qualcomm/_passes/remove_0d_tensor.py index 4056910f6..dc99dd598 100644 --- a/backends/qualcomm/_passes/remove_0d_tensor.py +++ b/backends/qualcomm/_passes/remove_0d_tensor.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import copy + import torch from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult @@ -28,8 +30,23 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: graph = graph_module.graph for node in graph.nodes: if node.target in self.remove_ops and len(node.meta["val"].shape) == 0: + source, dim, index = node.args + replacement = source + size = source.meta["val"].shape[dim] + if size != 1: + # Keep the selected element as a 1-element slice; the source + # itself is only equivalent when it holds a single element. + index %= size + with graph.inserting_before(node): + replacement = graph.create_node( + "call_function", + exir_ops.edge.aten.slice_copy.Tensor, + (source, dim, index, index + 1), + ) + replacement.meta = copy.copy(node.meta) + replacement.meta["val"] = source.meta["val"].narrow(dim, index, 1) for user_n in list(node.users.keys()): - user_n.replace_input_with(node, node.args[0]) + user_n.replace_input_with(node, replacement) graph.erase_node(node) dead_code_elimination_pass(graph_module) diff --git a/backends/qualcomm/builders/op_to.py b/backends/qualcomm/builders/op_to.py index 9b0243079..1c77cffc4 100644 --- a/backends/qualcomm/builders/op_to.py +++ b/backends/qualcomm/builders/op_to.py @@ -12,7 +12,7 @@ from .node_visitor import NodeVisitor, QNN_TENSOR_TYPE_MAP from .node_visitor_manager import register_node_visitor -from .qnn_constants import OpCast, OpConvert, QNN_OP_PACKAGE_NAME_QTI_AISW +from .qnn_constants import OpCast, OpConvert, OpQuantize, QNN_OP_PACKAGE_NAME_QTI_AISW @register_node_visitor @@ -138,6 +138,36 @@ def define_node( nodes_to_wrappers, ) + # Cast copies an integer's raw value into a quantized tensor, and the HTP + # rejects Convert there: cast to float, then quantize. + if ( + node.meta.get(QCOM_QUANT_ATTRS) + and not input_node.meta.get(QCOM_QUANT_ATTRS) + and not input_node.meta["val"].is_floating_point() + ): + float_tensor_wrapper = self.define_custom_tensor_wrapper( + node_name=node.name + "_float", + tensor_type=PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + dtype=PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32, + quant_encoding=PyQnnManager.Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_UNDEFINED, + quant_configs={}, + dims=output_tensor.size(), + tensor=output_tensor, + is_fake_tensor=True, + nodes_to_wrappers=nodes_to_wrappers, + ) + cast_op = PyQnnManager.PyQnnOpWrapper( + f"{node.name}_float", QNN_OP_PACKAGE_NAME_QTI_AISW, OpCast.op_name + ) + cast_op.AddInputTensors(node_input_tensors) + cast_op.AddOutputTensors([float_tensor_wrapper]) + quantize_op = PyQnnManager.PyQnnOpWrapper( + node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpQuantize.op_name + ) + quantize_op.AddInputTensors([float_tensor_wrapper]) + quantize_op.AddOutputTensors([output_tensor_wrapper]) + return ops + [cast_op, quantize_op] + qnn_op = OpCast if self.is_cast_node(node) else OpConvert op = PyQnnManager.PyQnnOpWrapper( node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, qnn_op.op_name diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 7d424101f..32d3502ab 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -378,6 +378,14 @@ def forward(self, x): return x.type(torch.IntTensor) +class CastIntToFloat(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, y): + return y + (x + 1).float() + + class CastMultiUsers(torch.nn.Module): def __init__(self): super().__init__() @@ -2580,6 +2588,15 @@ def forward(self, x, y): ) +class SelectElement(torch.nn.Module): + def __init__(self, index): + super().__init__() + self.index = index + + def forward(self, x, y): + return y + x[self.index] + + class SelectScatter(torch.nn.Module): def __init__(self, dim, index): super().__init__() diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index e8da5fbd8..88a26f7a4 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -709,6 +709,41 @@ def forward(self, x): raise edge.to_executorch() + def test_unbind_of_quantized_constant_lowers(self): + """AnnotateUnbind read the quant attrs from the quantize node feeding unbind's + dequantize, but a folded quantized constant feeds the dequantize directly. + """ + + class UnbindConstant(torch.nn.Module): + def __init__(self): + super().__init__() + self.table = torch.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]) + + def forward(self, x): + a, b = torch.unbind(self.table, dim=0) + return x * a + b + + sample_input = (torch.randn(4, 4),) + quantizer = QnnQuantizer() + quantizer.set_default_quant_config(quant_dtype=QuantDtype.use_16a8w) + prepared = prepare_pt2e( + torch.export.export(UnbindConstant(), sample_input).module(), quantizer + ) + prepared(*sample_input) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=False), + ) + try: + edge = to_edge_transform_and_lower_to_qnn( + convert_pt2e(prepared), 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 + edge.to_executorch() + def test_constant_only_partition_stays_on_cpu(self): """An int64 buffer gets an int32 cast from I64toI32. When the buffer's consumer is not delegated, that cast alone would form a partition with no graph input, which diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index b260597c2..57b7e252d 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -445,10 +445,11 @@ def test_qnn_backend_bmm(self): self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_cast(self): - modules = [Cast(), CastMultiUsers()] # noqa: F405 + modules = [Cast(), CastMultiUsers(), CastIntToFloat()] # noqa: F405 sample_inputs = [ (10 * torch.rand((9, 4, 5, 3)),), (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3)), + (torch.randint(0, 80, size=(3, 3)), torch.randn(3, 3)), ] for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)): with self.subTest(i=i): @@ -2246,6 +2247,13 @@ def test_qnn_backend_select_copy(self): sample_input = (torch.randn([1, 3, 3, 3]),) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_select_element(self): + sample_input = (torch.randn(4), torch.randn(3, 4)) + for i, index in enumerate([1, -1]): + with self.subTest(i=i): + module = SelectElement(index) # noqa: F405 + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_select_scatter(self): test_comb = [ { @@ -3350,6 +3358,14 @@ def test_qnn_backend_cast(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_cast_int_to_float(self): + module = CastIntToFloat() # noqa: F405 + sample_input = (torch.randint(0, 80, size=(3, 3)), torch.randn(3, 3)) + module = self.get_qdq_module(module, sample_input) + # absolute: the class rtol of 1 would accept an all-zero output + self.atol, self.rtol = 1, 0 + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_cat(self): modules = [Cat2(), Cat3(), Cat4(), Cat5()] # noqa: F405 sample_input = (torch.randn(1, 1, 2, 2), torch.randn(1, 1, 4, 2)) @@ -5385,6 +5401,14 @@ def test_qnn_backend_select_copy(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_select_element(self): + sample_input = (torch.randn(4), torch.randn(3, 4)) + for i, index in enumerate([1, -1]): + with self.subTest(i=i): + module = SelectElement(index) # noqa: F405 + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_select_scatter(self): test_comb = [ {