From 59df9ef546faaf2c70622a202b9bf5830a1f76c3 Mon Sep 17 00:00:00 2001 From: winskuo-quic Date: Tue, 18 Mar 2025 17:28:47 +0800 Subject: [PATCH 1/2] Qualcomm AI Engine Direct - Mimi Enablement Stage 1 --- backends/qualcomm/_passes/__init__.py | 16 +- .../_passes/convert_conv1d_to_conv2d.py | 99 ++++++++++ backends/qualcomm/_passes/decompose_expm1.py | 46 +++++ backends/qualcomm/_passes/decompose_silu.py | 13 +- backends/qualcomm/_passes/layout_transform.py | 6 + .../_passes/lift_constant_scalar_operands.py | 34 ++-- .../qualcomm/_passes/replace_arange_args.py | 48 +++++ ...ce_inf_buffer.py => replace_inf_values.py} | 14 +- .../qualcomm/_passes/tensor_i64_to_i32.py | 11 +- backends/qualcomm/_passes/utils.py | 14 ++ backends/qualcomm/builders/__init__.py | 12 ++ backends/qualcomm/builders/op_and.py | 59 ++++++ backends/qualcomm/builders/op_conv2d.py | 169 +----------------- backends/qualcomm/builders/op_elu.py | 68 +++++++ backends/qualcomm/builders/op_exp.py | 59 ++++++ backends/qualcomm/builders/op_pad.py | 4 +- .../qualcomm/builders/op_scalar_tensor.py | 50 ++++++ backends/qualcomm/builders/op_sqrt.py | 4 +- backends/qualcomm/builders/op_stack.py | 71 ++++++++ backends/qualcomm/builders/op_unbind.py | 73 ++++++++ backends/qualcomm/builders/qnn_constants.py | 32 +++- backends/qualcomm/partition/common_defs.py | 3 + backends/qualcomm/quantizer/annotators.py | 108 ++++++++--- backends/qualcomm/quantizer/quantizer.py | 8 +- backends/qualcomm/tests/models.py | 71 +++++++- backends/qualcomm/tests/test_qnn_delegate.py | 80 ++++++++- backends/qualcomm/utils/utils.py | 10 +- 27 files changed, 947 insertions(+), 235 deletions(-) create mode 100644 backends/qualcomm/_passes/convert_conv1d_to_conv2d.py create mode 100644 backends/qualcomm/_passes/decompose_expm1.py create mode 100644 backends/qualcomm/_passes/replace_arange_args.py rename backends/qualcomm/_passes/{replace_inf_buffer.py => replace_inf_values.py} (58%) create mode 100644 backends/qualcomm/builders/op_and.py create mode 100644 backends/qualcomm/builders/op_elu.py create mode 100644 backends/qualcomm/builders/op_exp.py create mode 100644 backends/qualcomm/builders/op_scalar_tensor.py create mode 100644 backends/qualcomm/builders/op_stack.py create mode 100644 backends/qualcomm/builders/op_unbind.py diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index fb65e6b5f75..fb1f985edb9 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -1,10 +1,18 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + from .annotate_decomposed import AnnotateDecomposed from .annotate_quant_attrs import AnnotateQuantAttrs from .constant_i64_to_i32 import ConstantI64toI32 from .convert_bmm_to_matmul import ConvertBmmToMatmul +from .convert_conv1d_to_conv2d import ConvertConv1dToConv2d from .convert_to_linear import ConvertToLinear from .decompose_any import DecomposeAny from .decompose_einsum import DecomposeEinsum +from .decompose_expm1 import DecomposeExpM1 from .decompose_linalg_vector_norm import DecomposeLinalgVectorNorm from .decompose_silu import DecomposeSilu from .expand_broadcast_tensor_shape import ExpandBroadcastTensorShape @@ -19,8 +27,9 @@ from .recompose_rms_norm import RecomposeRmsNorm from .reduce_dynamic_range import ReduceDynamicRange from .remove_redundancy import RemoveRedundancy +from .replace_arange_args import ReplaceArangeArgs from .replace_index_put_input import ReplaceIndexPutInput -from .replace_inf_buffer import ReplaceInfBuffer +from .replace_inf_values import ReplaceInfValues from .tensor_i64_to_i32 import TensorI64toI32 @@ -29,10 +38,12 @@ AnnotateQuantAttrs, ConstantI64toI32, ConvertBmmToMatmul, + ConvertConv1dToConv2d, RecomposePReLU, ConvertToLinear, DecomposeAny, DecomposeEinsum, + DecomposeExpM1, DecomposeLinalgVectorNorm, DecomposeSilu, ExpandBroadcastTensorShape, @@ -46,7 +57,8 @@ RecomposeRmsNorm, ReduceDynamicRange, RemoveRedundancy, + ReplaceArangeArgs, ReplaceIndexPutInput, - ReplaceInfBuffer, + ReplaceInfValues, TensorI64toI32, ] diff --git a/backends/qualcomm/_passes/convert_conv1d_to_conv2d.py b/backends/qualcomm/_passes/convert_conv1d_to_conv2d.py new file mode 100644 index 00000000000..947b631dbbf --- /dev/null +++ b/backends/qualcomm/_passes/convert_conv1d_to_conv2d.py @@ -0,0 +1,99 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# 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 +import torch.nn as nn +from executorch.backends.qualcomm.builders.utils import get_parameter, set_parameter +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + +from .utils import copy_meta + + +class ConvertConv1dToConv2d(ExportPass): + """ + Conv1d is not supported by QNN. + Change it to input -> unsqueeze -> conv2d -> squeeze -> output + """ + + def __init__(self, edge_program: torch.export.ExportedProgram): + super(ConvertConv1dToConv2d, self).__init__() + self.edge_program = edge_program + + def call(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + conv_op = exir_ops.edge.aten.convolution.default + for node in graph.nodes: + if node.target == conv_op and node.meta["val"].dim() == 3: + + input_node = node.args[0] + with graph_module.graph.inserting_after(input_node): + unsqueeze_op = exir_ops.edge.aten.unsqueeze_copy.default + unsqueeze_node = graph.create_node( + "call_function", + unsqueeze_op, + ( + input_node, + 2, + ), + ) + unsqueeze_node.meta = copy_meta( + input_node.meta, lambda m: {**m, "val": m["val"].unsqueeze(2)} + ) + with graph_module.graph.inserting_after(unsqueeze_node): + + filter_node = node.args[1] + filter_node.meta["val"] = ( + filter_node.meta["val"].unsqueeze(2).contiguous() + ) + filter_tensor = get_parameter(filter_node, self.edge_program) + # Ensure tensor is nn.Parameter type, so program does not fail during edge_program._validate() + filter_tensor = nn.Parameter(filter_tensor.unsqueeze(2)) + set_parameter(filter_tensor, filter_node, self.edge_program) + + bias_node = node.args[2] + stride = [1] + node.args[3] + padding = [0] + node.args[4] + dilation = [1] + node.args[5] + transpose = node.args[6] + output_padding = [0] + node.args[7] + groups = node.args[8] + + conv2d_node = graph.create_node( + "call_function", + conv_op, + ( + unsqueeze_node, + filter_node, + bias_node, + stride, + padding, + dilation, + transpose, + output_padding, + groups, + ), + ) + conv2d_node.meta = copy_meta( + node.meta, lambda m: {**m, "val": m["val"].unsqueeze(2)} + ) + + with graph_module.graph.inserting_after(conv2d_node): + squeeze_op = exir_ops.edge.aten.squeeze_copy.dims + squeeze_node = graph.create_node( + "call_function", + squeeze_op, + ( + conv2d_node, + [2], + ), + ) + squeeze_node.meta = copy_meta(node.meta) + for user in node.users.copy(): + user.replace_input_with(node, squeeze_node) + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/decompose_expm1.py b/backends/qualcomm/_passes/decompose_expm1.py new file mode 100644 index 00000000000..8fe6ebdec5b --- /dev/null +++ b/backends/qualcomm/_passes/decompose_expm1.py @@ -0,0 +1,46 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# 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.exir.pass_base import ExportPass, PassResult + +from .utils import copy_meta + + +class DecomposeExpM1(ExportPass): + """ + Decompose for expm1 to exponential and minus 1. + """ + + def __init__(self, quantization_capture=False) -> None: + super().__init__() + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + for node in graph.nodes: + if node.target == torch.ops.aten.special_expm1.default: + input_node = node.args[0] + with graph_module.graph.inserting_after(input_node): + exp_op = torch.ops.aten.exp.default + exp_node = graph.create_node("call_function", exp_op, (input_node,)) + exp_node.meta = copy_meta(node.meta) + with graph_module.graph.inserting_after(exp_node): + sub_op = torch.ops.aten.sub.Tensor + sub_node = graph.create_node( + "call_function", + sub_op, + ( + exp_node, + 1, + ), + ) + sub_node.meta = copy_meta(node.meta) + for user in node.users.copy(): + user.replace_input_with(node, sub_node) + + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/decompose_silu.py b/backends/qualcomm/_passes/decompose_silu.py index 96c48920419..c3ac45a8d9d 100644 --- a/backends/qualcomm/_passes/decompose_silu.py +++ b/backends/qualcomm/_passes/decompose_silu.py @@ -3,22 +3,17 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Dict import torch from executorch.exir.pass_base import ExportPass, PassResult +from .utils import copy_meta + class DecomposeSilu(ExportPass): def __init__(self): super(DecomposeSilu, self).__init__() - def _copy_meta(self, meta: Dict): - copied = {} - for k, v in meta.items(): - copied[k] = v - return copied - def call(self, graph_module: torch.fx.GraphModule): graph = graph_module.graph for node in graph.nodes: @@ -34,14 +29,14 @@ def call(self, graph_module: torch.fx.GraphModule): torch.ops.aten.sigmoid.default, (silu_node_input,), ) - sigmoid_node.meta = self._copy_meta(silu_node.meta) + sigmoid_node.meta = copy_meta(silu_node.meta) with graph_module.graph.inserting_after(sigmoid_node): mul_node = graph.create_node( "call_function", torch.ops.aten.mul.Tensor, (silu_node_input, sigmoid_node), ) - mul_node.meta = self._copy_meta(silu_node.meta) + mul_node.meta = copy_meta(silu_node.meta) for user in silu_node.users.copy(): user.replace_input_with(silu_node, mul_node) diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 31bb936f3c4..64fdcb2bb88 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -49,12 +49,15 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.bitwise_or.Tensor, exir_ops.edge.aten.bmm.default, + exir_ops.edge.aten.bitwise_and.Tensor, exir_ops.edge.aten.cat.default, exir_ops.edge.aten.ceil.default, exir_ops.edge.aten.clamp.default, exir_ops.edge.aten.constant_pad_nd.default, exir_ops.edge.aten.div.Tensor, + exir_ops.edge.aten.elu.default, exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.exp.default, exir_ops.edge.aten.full.default, exir_ops.edge.aten.full_like.default, exir_ops.edge.aten.ge.Tensor, @@ -87,10 +90,13 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.sqrt.default, exir_ops.edge.aten.sub.Tensor, exir_ops.edge.aten.sum.dim_IntList, + exir_ops.edge.aten.stack.default, exir_ops.edge.aten.topk.default, exir_ops.edge.aten._to_copy.default, + exir_ops.edge.aten.unbind.int, exir_ops.edge.aten.where.self, _operator.getitem, + torch.ops.aten.scalar_tensor.default, } layout_type = { diff --git a/backends/qualcomm/_passes/lift_constant_scalar_operands.py b/backends/qualcomm/_passes/lift_constant_scalar_operands.py index 749d30f3564..cef28988520 100644 --- a/backends/qualcomm/_passes/lift_constant_scalar_operands.py +++ b/backends/qualcomm/_passes/lift_constant_scalar_operands.py @@ -28,24 +28,27 @@ class TensorConstant: class TensorOpInfo: target: torch._ops.OpOverload use_schema_args: bool + use_self_dtype: bool SCALAR_OPS = { - aten.eq.Scalar: TensorOpInfo(aten.eq.Tensor, False), - aten.ge.Scalar: TensorOpInfo(aten.ge.Tensor, False), - aten.gt.Scalar: TensorOpInfo(aten.gt.Tensor, False), - aten.le.Scalar: TensorOpInfo(aten.le.Tensor, False), - aten.lt.Scalar: TensorOpInfo(aten.lt.Tensor, False), - aten.ne.Scalar: TensorOpInfo(aten.ne.Tensor, False), - aten.add.Scalar: TensorOpInfo(aten.add.Tensor, False), - aten.add_.Scalar: TensorOpInfo(aten.add_.Tensor, False), - aten.div.Scalar: TensorOpInfo(aten.div.Tensor, False), - aten.mul.Scalar: TensorOpInfo(aten.mul.Tensor, False), - aten.rsub.Scalar: TensorOpInfo(aten.rsub.Tensor, False), - aten.sub.Scalar: TensorOpInfo(aten.sub.Tensor, False), - aten.pow.Tensor_Scalar: TensorOpInfo(aten.pow.Tensor_Tensor, False), + aten.eq.Scalar: TensorOpInfo(aten.eq.Tensor, False, False), + aten.ge.Scalar: TensorOpInfo(aten.ge.Tensor, False, False), + aten.gt.Scalar: TensorOpInfo(aten.gt.Tensor, False, False), + aten.le.Scalar: TensorOpInfo(aten.le.Tensor, False, False), + aten.lt.Scalar: TensorOpInfo(aten.lt.Tensor, False, False), + aten.ne.Scalar: TensorOpInfo(aten.ne.Tensor, False, False), + aten.add.Scalar: TensorOpInfo(aten.add.Tensor, False, False), + aten.add_.Scalar: TensorOpInfo(aten.add_.Tensor, False, False), + aten.div.Scalar: TensorOpInfo(aten.div.Tensor, False, False), + aten.mul.Scalar: TensorOpInfo(aten.mul.Tensor, False, False), + aten.rsub.Scalar: TensorOpInfo(aten.rsub.Tensor, False, False), + aten.sub.Scalar: TensorOpInfo(aten.sub.Tensor, False, False), + aten.pow.Tensor_Scalar: TensorOpInfo(aten.pow.Tensor_Tensor, False, False), # The scalar number arg[1] is missing when using default. Result in a corner case to deal - aten.leaky_relu.default: TensorOpInfo(aten.prelu.default, True), + aten.leaky_relu.default: TensorOpInfo(aten.prelu.default, True, False), + aten.where.ScalarOther: TensorOpInfo(aten.where.self, False, True), + aten.where.Scalar: TensorOpInfo(aten.where.self, False, True), } @@ -63,11 +66,14 @@ def __init__(self): def _build_tensor_constant( self, gm: torch.fx.GraphModule, node: fx.Node, const_val ) -> TensorConstant: + # For dtype, in some cases, we cannot use node.args[0] as scalar dtype. + # Ex: Where op args[0] can be bool, however, we probably want args[1] and args[2] to be dtype same as node.meta["val"] instead of bool type tensor = torch.tensor( [const_val], dtype=( node.args[0].meta["val"].dtype if not is_float_tensor(node) + and not SCALAR_OPS.get(node.target).use_self_dtype else node.meta["val"].dtype ), device=node.meta["val"].device, diff --git a/backends/qualcomm/_passes/replace_arange_args.py b/backends/qualcomm/_passes/replace_arange_args.py new file mode 100644 index 00000000000..19ebc60227f --- /dev/null +++ b/backends/qualcomm/_passes/replace_arange_args.py @@ -0,0 +1,48 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# 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.exir.pass_base import ExportPass, PassResult + +from .utils import copy_meta + + +class ReplaceArangeArgs(ExportPass): + """ + During annotation, kwargs for arange will be removed due to restrictions by quantizer. + This causes arange to have no dtype, which means FP nodes might be inferred as INT nodes during calibration. + This can cause calibration to fail since QDQ can only be applied on FP nodes but not INT nodes. + To hint the dtype, we provide step size as 1.0 instead of 1, which makes the node a FP node. + """ + + def __init__(self, quantization_capture=False) -> None: + super().__init__() + self.quantization_capture = quantization_capture + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + for node in graph.nodes: + if node.target == torch.ops.aten.arange.default: + if torch.is_floating_point(node.meta["val"]) and len(node.args) == 1: + with graph_module.graph.inserting_after(node): + step_arange_op = torch.torch.ops.aten.arange.start_step + step_arange_node = graph.create_node( + "call_function", + step_arange_op, + ( + 0, + node.args[0], + 1.0, + ), + ) + step_arange_node.meta = copy_meta(node.meta) + + for user in node.users.copy(): + user.replace_input_with(node, step_arange_node) + + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/replace_inf_buffer.py b/backends/qualcomm/_passes/replace_inf_values.py similarity index 58% rename from backends/qualcomm/_passes/replace_inf_buffer.py rename to backends/qualcomm/_passes/replace_inf_values.py index 776bc9beeba..5f7fb9bd768 100644 --- a/backends/qualcomm/_passes/replace_inf_buffer.py +++ b/backends/qualcomm/_passes/replace_inf_values.py @@ -7,20 +7,30 @@ from executorch.exir.pass_base import ExportPass, PassResult -class ReplaceInfBuffer(ExportPass): +class ReplaceInfValues(ExportPass): """ Due to limitation in Qnn, we need to change inf or -inf to arbitrary value in quantization. """ def __init__(self): - super(ReplaceInfBuffer, self).__init__() + super(ReplaceInfValues, self).__init__() def call(self, graph_module: torch.fx.GraphModule): for buf_name, tensor in graph_module.named_buffers(): if tensor.is_floating_point(): + # 255 here is mainly for attention_mask in Llama for reasonable quant scale tensor[tensor == float("inf")] = 255 tensor[tensor == float("-inf")] = -255 setattr(graph_module, buf_name, tensor) + for node in graph_module.graph.nodes: + arg_list = list(node.args) + for index, arg in enumerate(arg_list): + if arg == float("-inf"): + arg_list[index] = torch.finfo(torch.float32).min + elif arg == float("inf"): + arg_list[index] = torch.finfo(torch.float32).max + node.args = tuple(arg_list) + graph_module.recompile() return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/tensor_i64_to_i32.py b/backends/qualcomm/_passes/tensor_i64_to_i32.py index b590e30884c..baddd747f99 100644 --- a/backends/qualcomm/_passes/tensor_i64_to_i32.py +++ b/backends/qualcomm/_passes/tensor_i64_to_i32.py @@ -24,6 +24,9 @@ class TensorI64toI32(ExportPass): cast_ops = { torch.ops.aten.argmin.default, + torch.ops.aten.arange.start_step, + torch.ops.aten.full.default, + torch.ops.aten.scalar_tensor.default, } def __init__(self, edge_program): @@ -61,7 +64,13 @@ def _cast_to_int32(self, core_ep: ExirExportedProgram): cast_node.args = args for user in users: - user.replace_input_with(n, cast_node) + # _assert_tensor_metadata is used to check dtype, which will cause lowering to fail since we are changing int64 to int32 + # We also skip if the next op is already a cast op, which prevents redundant casting. + if user.target not in { + torch.ops.aten._assert_tensor_metadata.default, + torch.ops.aten._to_copy.default, + }: + user.replace_input_with(n, cast_node) core_ep.exported_program._graph_signature = _get_updated_graph_signature( core_ep.exported_program._graph_signature, diff --git a/backends/qualcomm/_passes/utils.py b/backends/qualcomm/_passes/utils.py index 23dfb569a8f..0c838e9a676 100755 --- a/backends/qualcomm/_passes/utils.py +++ b/backends/qualcomm/_passes/utils.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. +from typing import Dict + import torch from executorch.backends.qualcomm.builders.utils import get_parameter from executorch.backends.qualcomm.utils.constants import QCOM_DTYPE, QCOM_ENCODING @@ -24,6 +26,15 @@ } +def copy_meta(meta: Dict, callback=None): + copied = {} + for k, v in meta.items(): + copied[k] = v + if callback: + copied = callback(copied) + return copied + + def get_quant_attrs( edge_program: torch.export.ExportedProgram, quant_node: torch.fx.Node ): @@ -66,6 +77,7 @@ def get_passes_dependency_for_capture_program(): AnnotateQuantAttrs, ConstantI64toI32, ConvertBmmToMatmul, + ConvertConv1dToConv2d, ConvertToLinear, DecomposeAny, DecomposeLinalgVectorNorm, @@ -91,6 +103,7 @@ def get_passes_dependency_for_capture_program(): ], ConstantI64toI32: [RemoveRedundancy], ConvertBmmToMatmul: [ConvertToLinear], + ConvertConv1dToConv2d: [FoldQDQ], ConvertToLinear: [RecomposePixelUnshuffle], DecomposeAny: [RemoveRedundancy], DecomposeLinalgVectorNorm: [RemoveRedundancy], @@ -98,6 +111,7 @@ def get_passes_dependency_for_capture_program(): FoldQDQ: [AnnotateQuantAttrs, AnnotateDecomposed], LayoutTransform: [ AnnotateQuantAttrs, + ConvertConv1dToConv2d, ExpandBroadcastTensorShape, ], RecomposePixelUnshuffle: [RemoveRedundancy], diff --git a/backends/qualcomm/builders/__init__.py b/backends/qualcomm/builders/__init__.py index c5352a7fbee..cc85333f26b 100644 --- a/backends/qualcomm/builders/__init__.py +++ b/backends/qualcomm/builders/__init__.py @@ -9,6 +9,7 @@ op_abs, op_adaptive_avg_pool2d, op_add, + op_and, op_arange, op_argmin, op_avg_pool2d, @@ -22,8 +23,10 @@ op_depth_to_space, op_dequantize, op_div, + op_elu, op_embedding, op_eq, + op_exp, op_expand, op_full, op_full_like, @@ -62,6 +65,7 @@ op_reshape, op_rms_norm, op_rsqrt, + op_scalar_tensor, op_select_copy, op_sigmoid, op_sin, @@ -72,12 +76,14 @@ op_split_with_sizes, op_sqrt, op_squeeze, + op_stack, op_sub, op_sum_int_list, op_tanh, op_to, op_topk, op_transpose, + op_unbind, op_unsqueeze, op_upsample_bilinear2d, op_upsample_nearest2d, @@ -89,6 +95,7 @@ op_abs, op_adaptive_avg_pool2d, op_add, + op_and, op_arange, op_argmin, op_avg_pool2d, @@ -102,8 +109,10 @@ op_depth_to_space, op_dequantize, op_div, + op_elu, op_embedding, op_eq, + op_exp, op_expand, op_full, op_full_like, @@ -142,6 +151,7 @@ op_reshape, op_rms_norm, op_rsqrt, + op_scalar_tensor, op_select_copy, op_sigmoid, op_sin, @@ -152,12 +162,14 @@ op_split_with_sizes, op_squeeze, op_sqrt, + op_stack, op_sub, op_sum_int_list, op_tanh, op_topk, op_to, op_transpose, + op_unbind, op_unsqueeze, op_upsample_bilinear2d, op_upsample_nearest2d, diff --git a/backends/qualcomm/builders/op_and.py b/backends/qualcomm/builders/op_and.py new file mode 100644 index 00000000000..44e6f2893f5 --- /dev/null +++ b/backends/qualcomm/builders/op_and.py @@ -0,0 +1,59 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import torch + +from .node_visitor import NodeVisitor, register_node_visitor +from .qnn_constants import OpElementWiseAnd, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class OpAnd(NodeVisitor): + target = ["aten.bitwise_and.Tensor"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + out_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + out_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + and_output_tensors = [output_tensor_wrapper] + + and_input_tensors = [] + for index in range(2): + input_node = node.args[index] + input_tensor = self.get_tensor(input_node, node) + tensor_type = PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE + + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + tensor_type, + nodes_to_wrappers, + ) + and_input_tensors.append(input_tensor_wrapper) + and_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpElementWiseAnd.op_name, + ) + and_op.AddInputTensors(and_input_tensors) + and_op.AddOutputTensors(and_output_tensors) + return and_op diff --git a/backends/qualcomm/builders/op_conv2d.py b/backends/qualcomm/builders/op_conv2d.py index a6051636d3e..c019a835223 100644 --- a/backends/qualcomm/builders/op_conv2d.py +++ b/backends/qualcomm/builders/op_conv2d.py @@ -4,7 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import warnings from typing import cast, Dict, List import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper @@ -17,8 +16,6 @@ from .qnn_constants import ( OpConv2d, OpDepthWiseConv2d, - OpExpandDims, - OpReshape, OpTransposeConv2d, QNN_OP_PACKAGE_NAME_QTI_AISW, ) @@ -102,176 +99,16 @@ def _add_conv_op_parameter( return conv_op - def _define_conv1d( - self, - node: torch.fx.Node, - nodes_to_wrappers: Dict[str, PyQnnWrapper.TensorWrapper], - ) -> PyQnnWrapper.PyQnnOpWrapper: - """ - Conv1D is a special case for convolutional operation. QNN does not support Conv1D, therefore, - we need to cast from input -> Conv1d -> output to input -> unsqueeze -> Conv2d -> squeeze -> output. - """ - transpose_conv = cast(bool, node.args[6]) - if transpose_conv: - print("ConvTranspose1d is not yet supported") - return - - op_wrapper_list = [] # op_wrapper to return - unsqueeze_input_node = node.args[0] - input_quant_encoding, input_quant_configs = self.get_quant_encoding_conf( - unsqueeze_input_node, node - ) - - unsqueeze_input_tensor = self.get_tensor(unsqueeze_input_node, node) - unsqueeze_input_tensor_wrapper = self.define_tensor( - unsqueeze_input_node, - node, - unsqueeze_input_tensor, - PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, - nodes_to_wrappers, - ) - unsqueeze_output_tensor = unsqueeze_input_tensor.unsqueeze(1).contiguous() - dtype = self.get_data_type(unsqueeze_output_tensor, input_quant_configs) - unsqueeze_output_tensor_wrapper = self.define_custom_tensor_wrapper( - node_name=node.name + "_unsqueeze", - tensor_type=PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, - dtype=dtype, - quant_encoding=input_quant_encoding, - quant_configs=input_quant_configs, - dims=unsqueeze_output_tensor.size(), - tensor=unsqueeze_output_tensor, - is_fake_tensor=True, - nodes_to_wrappers=nodes_to_wrappers, - ) - unsqueeze_op = PyQnnWrapper.PyQnnOpWrapper( - node.name + "_unsqueeze", - QNN_OP_PACKAGE_NAME_QTI_AISW, - OpExpandDims.op_name, - ) - unsqueeze_op.AddInputTensors([unsqueeze_input_tensor_wrapper]) - unsqueeze_op.AddOutputTensors([unsqueeze_output_tensor_wrapper]) - unsqueeze_op.AddScalarParam( - OpExpandDims.param_axis, - PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32, - {QCOM_DATA: np.uint32(1)}, - ) - op_wrapper_list.append(unsqueeze_op) - - filter_node = node.args[1] - filter_tensor = ( - get_parameter(filter_node, self.edge_program).unsqueeze(2).contiguous() - ) - filter_axis_order = (2, 3, 1, 0) - filter_tensor = filter_tensor.permute(dims=filter_axis_order).contiguous() - filter_tensor_wrapper = self.define_tensor( - filter_node, - node, - filter_tensor, - PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, - nodes_to_wrappers, - ) - conv_input_tensors = [unsqueeze_output_tensor_wrapper, filter_tensor_wrapper] - if node.args[2] is not None: - bias_node = node.args[2] - bias_tensor = get_parameter(bias_node, self.edge_program) - bias_tensor_wrapper = self.define_tensor( - bias_node, - node, - bias_tensor, - PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, - nodes_to_wrappers, - ) - conv_input_tensors.append(bias_tensor_wrapper) - - stride = [1] + cast(List[int], node.args[3]) - padding = [0] + cast(List[int], node.args[4]) - dilation = [1] + cast(List[int], node.args[5]) - groups = cast(int, node.args[8]) - - # args[6] = transposed - if cast(bool, node.args[6]): - warnings.warn( - "[QNN Delegate Op Builder]: Currently, No support for transposed convolution.", - stacklevel=1, - ) - return - - # args[7] = output padding - if not all(out_pad == 0 for out_pad in cast(List[int], node.args[7])): - warnings.warn( - "[QNN Delegate Op Builder]: QNN does not support output padding.", - stacklevel=1, - ) - return - - stride_shape = [len(stride)] - padding_shape = [2, 2] - dilation_shape = [len(dilation)] - - conv_op = PyQnnWrapper.PyQnnOpWrapper( - node.name + "_squeeze", - QNN_OP_PACKAGE_NAME_QTI_AISW, - OpConv2d.op_name, - ) - conv_output_tensor = self.get_tensor(node, node) - conv_output_tensor = conv_output_tensor.unsqueeze(1).contiguous() - dtype = self.get_data_type(conv_output_tensor, input_quant_configs) - conv_output_tensor_wrapper = self.define_custom_tensor_wrapper( - node_name=node.name + "_squeeze", - tensor_type=PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, - dtype=dtype, - quant_encoding=input_quant_encoding, - quant_configs=input_quant_configs, - dims=conv_output_tensor.size(), - tensor=conv_output_tensor, - is_fake_tensor=True, - nodes_to_wrappers=nodes_to_wrappers, - ) - conv_op = self._add_conv_op_parameter( - OpConv2d, - conv_op, - conv_input_tensors, - [conv_output_tensor_wrapper], - stride, - stride_shape, - padding, - padding_shape, - dilation, - dilation_shape, - groups=groups, - ) - op_wrapper_list.append(conv_op) - - squeeze_op = PyQnnWrapper.PyQnnOpWrapper( - node.name, - QNN_OP_PACKAGE_NAME_QTI_AISW, - OpReshape.op_name, - ) - squeeze_output_tensor = self.get_tensor(node, node) - squeeze_output_tensor_wrapper = self.define_tensor( - node, - node, - squeeze_output_tensor, - PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, - nodes_to_wrappers, - node_name=node.name, - ) - squeeze_op.AddInputTensors([conv_output_tensor_wrapper]) - squeeze_op.AddOutputTensors([squeeze_output_tensor_wrapper]) - op_wrapper_list.append(squeeze_op) - - return op_wrapper_list - def define_node( self, node: torch.fx.Node, nodes_to_wrappers: Dict[str, PyQnnWrapper.TensorWrapper], ) -> PyQnnWrapper.PyQnnOpWrapper: - if get_parameter(node.args[1], self.edge_program).dim() == 3: - return self._define_conv1d(node, nodes_to_wrappers) - input_node = node.args[0] input_tensor = self.get_tensor(input_node, node) + assert ( + input_tensor.dim() == 4 + ), "All Conv should be converted to Conv2D in ConvertConv1dToConv2d" input_tensor_wrapper = self.define_tensor( input_node, node, diff --git a/backends/qualcomm/builders/op_elu.py b/backends/qualcomm/builders/op_elu.py new file mode 100644 index 00000000000..f9cc089c7bb --- /dev/null +++ b/backends/qualcomm/builders/op_elu.py @@ -0,0 +1,68 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_DATA + +from .node_visitor import NodeVisitor, register_node_visitor +from .qnn_constants import OpElu, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Elu(NodeVisitor): + target = ["aten.elu.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + # tensor input + input_node = node.args[0] + input_tensor = self.get_tensor(input_node, node) + + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + elu_input_tensors = [input_tensor_wrapper] + + out_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + out_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + elu_output_tensors = [output_tensor_wrapper] + + elu_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpElu.op_name, + ) + elu_op.AddInputTensors(elu_input_tensors) + elu_op.AddOutputTensors(elu_output_tensors) + + if len(node.args) == 2: + elu_op.AddScalarParam( + OpElu.param_alpha, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32, + {QCOM_DATA: np.uint32(node.args[1])}, + ) + + return elu_op diff --git a/backends/qualcomm/builders/op_exp.py b/backends/qualcomm/builders/op_exp.py new file mode 100644 index 00000000000..8c4794c9725 --- /dev/null +++ b/backends/qualcomm/builders/op_exp.py @@ -0,0 +1,59 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import torch + +from .node_visitor import NodeVisitor, register_node_visitor +from .qnn_constants import OpElementWiseExp, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Exp(NodeVisitor): + target = ["aten.exp.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + # tensor input + input_node = node.args[0] + input_tensor = self.get_tensor(input_node, node) + + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + exp_input_tensors = [input_tensor_wrapper] + + out_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + out_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + exp_output_tensors = [output_tensor_wrapper] + + exp_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpElementWiseExp.op_name, + ) + exp_op.AddInputTensors(exp_input_tensors) + exp_op.AddOutputTensors(exp_output_tensors) + + return exp_op diff --git a/backends/qualcomm/builders/op_pad.py b/backends/qualcomm/builders/op_pad.py index 10948859be9..5ec34065f8b 100644 --- a/backends/qualcomm/builders/op_pad.py +++ b/backends/qualcomm/builders/op_pad.py @@ -53,14 +53,14 @@ def define_node( pad_amount = np.reshape(cast(List[int], node.args[1]), (-1, 2))[::-1].astype( np.uint32 ) - # fullfill the pad amount for each idex of tensor + # fulfill the pad amount for each idex of tensor if zero_amounts := pad_amount_shape[0] - pad_amount.shape[0]: pad_amount = np.concatenate( (np.array([(0, 0)] * zero_amounts), pad_amount) ).astype(np.uint32) if QCOM_AXIS_ORDER in node.meta: - pad_amount = np.transpose(pad_amount, node.meta[QCOM_AXIS_ORDER]) + pad_amount = pad_amount[list(node.meta[QCOM_AXIS_ORDER])] pad_amount_val = node.args[2] pad_op = PyQnnWrapper.PyQnnOpWrapper( diff --git a/backends/qualcomm/builders/op_scalar_tensor.py b/backends/qualcomm/builders/op_scalar_tensor.py new file mode 100644 index 00000000000..d236f6674df --- /dev/null +++ b/backends/qualcomm/builders/op_scalar_tensor.py @@ -0,0 +1,50 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import torch + +from .node_visitor import NodeVisitor, register_node_visitor + + +@register_node_visitor +class Arange(NodeVisitor): + target = ["scalar_tensor.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + val = node.args[0] + out_tensor = torch.tensor([val], dtype=node.meta["val"].dtype) + + # The following clamping will only occur in FP mode. Clamping for quantized mode will happen in the pass ReplaceInfValues. + # negative infinite + if torch.isinf(out_tensor)[0] and (out_tensor < 0): + out_tensor = torch.tensor( + [torch.finfo(torch.float32).min], dtype=node.meta["val"].dtype + ) + # positive infinite + elif torch.isinf(out_tensor)[0] and (out_tensor > 0): + out_tensor = torch.tensor( + [torch.finfo(torch.float32).max], dtype=node.meta["val"].dtype + ) + # since we can derive the constant value of current op in AoT stage + # we only build static tensor here for consumers of current node + # to correctly reference the data + self.define_tensor( + node, + node, + out_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, + nodes_to_wrappers, + ) diff --git a/backends/qualcomm/builders/op_sqrt.py b/backends/qualcomm/builders/op_sqrt.py index dc6691460ca..030e6c3e10a 100644 --- a/backends/qualcomm/builders/op_sqrt.py +++ b/backends/qualcomm/builders/op_sqrt.py @@ -10,7 +10,7 @@ import torch from .node_visitor import NodeVisitor, register_node_visitor -from .qnn_constants import OpSqrt, QNN_OP_PACKAGE_NAME_QTI_AISW +from .qnn_constants import OpElementWiseSqrt, QNN_OP_PACKAGE_NAME_QTI_AISW @register_node_visitor @@ -51,7 +51,7 @@ def define_node( sqrt_op = PyQnnWrapper.PyQnnOpWrapper( node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, - OpSqrt.op_name, + OpElementWiseSqrt.op_name, ) sqrt_op.AddInputTensors(sqrt_input_tensors) sqrt_op.AddOutputTensors(sqrt_output_tensors) diff --git a/backends/qualcomm/builders/op_stack.py b/backends/qualcomm/builders/op_stack.py new file mode 100644 index 00000000000..616d0ee0ccc --- /dev/null +++ b/backends/qualcomm/builders/op_stack.py @@ -0,0 +1,71 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import cast, Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA + +from .node_visitor import NodeVisitor, register_node_visitor +from .qnn_constants import OpPack, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Stack(NodeVisitor): + target = ["aten.stack.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node_list = node.args[0] + stack_input_tensors = [] + for input_node in input_node_list: + input_tensor = self.get_tensor(input_node, node) + stack_inp_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + stack_input_tensors.append(stack_inp_tensor_wrapper) + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + stack_output_tensors = [output_tensor_wrapper] + + dim = 0 if len(node.args) == 1 else cast(int, node.args[1]) + if dim < 0: + dim = dim % len(input_tensor.shape) + if QCOM_AXIS_ORDER in node.meta: + dim = node.meta[QCOM_AXIS_ORDER].index(dim) + stack_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpPack.op_name, + ) + stack_op.AddInputTensors(stack_input_tensors) + stack_op.AddOutputTensors(stack_output_tensors) + + stack_op.AddScalarParam( + OpPack.param_axis, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(dim)}, + ) + + return stack_op diff --git a/backends/qualcomm/builders/op_unbind.py b/backends/qualcomm/builders/op_unbind.py new file mode 100644 index 00000000000..8ca62e2a07b --- /dev/null +++ b/backends/qualcomm/builders/op_unbind.py @@ -0,0 +1,73 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import cast, Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA + +from .node_visitor import NodeVisitor, register_node_visitor +from .qnn_constants import OpUnpack, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Unbind(NodeVisitor): + target = ["aten.unbind.int"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = node.args[0] + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, + nodes_to_wrappers, + ) + unbind_input_tensors = [input_tensor_wrapper] + + unbind_output_tensors = [] + for i in range(len(node.meta["val"])): + output_tensor = self.get_tensor(node, node, i) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + wrapper_idx=i, + ) + unbind_output_tensors.append(output_tensor_wrapper) + + dim = 0 if len(node.args) == 1 else cast(int, node.args[1]) + if dim < 0: + dim = dim % len(input_tensor.shape) + if QCOM_AXIS_ORDER in node.meta: + dim = node.meta[QCOM_AXIS_ORDER].index(dim) + unbind_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpUnpack.op_name, + ) + unbind_op.AddInputTensors(unbind_input_tensors) + unbind_op.AddOutputTensors(unbind_output_tensors) + + unbind_op.AddScalarParam( + OpUnpack.param_axis, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(dim)}, + ) + + return unbind_op diff --git a/backends/qualcomm/builders/qnn_constants.py b/backends/qualcomm/builders/qnn_constants.py index 1d55d56de0f..9613c755c7c 100644 --- a/backends/qualcomm/builders/qnn_constants.py +++ b/backends/qualcomm/builders/qnn_constants.py @@ -85,6 +85,11 @@ class OpElementWiseAdd: op_name: str = "ElementWiseAdd" +@dataclass(init=False, frozen=True) +class OpElementWiseAnd: + op_name: str = "ElementWiseAnd" + + @dataclass(init=False, frozen=True) class OpElementWiseCeil: op_name = "ElementWiseCeil" @@ -100,6 +105,11 @@ class OpElementWiseDivide: op_name: str = "ElementWiseDivide" +@dataclass(init=False, frozen=True) +class OpElementWiseExp: + op_name: str = "ElementWiseExp" + + @dataclass(init=False, frozen=True) class OpElementWiseEqual: op_name: str = "ElementWiseEqual" @@ -193,11 +203,22 @@ class OpElementWiseSelect: op_name = "ElementWiseSelect" +@dataclass(init=False, frozen=True) +class OpElementWiseSqrt: + op_name = "ElementWiseSquareRoot" + + @dataclass(init=False, frozen=True) class OpElementWiseSubtract: op_name = "ElementWiseSubtract" +@dataclass(init=False, frozen=True) +class OpElu: + op_name: str = "Elu" + param_alpha: str = "alpha" + + @dataclass(init=False, frozen=True) class OpExpandDims: op_name: str = "ExpandDims" @@ -423,11 +444,6 @@ class OpSplit: param_split_index: str = "split_index" -@dataclass(init=False, frozen=True) -class OpSqrt: - op_name: str = "ElementWiseSquareRoot" - - @dataclass(init=False, frozen=True) class OpSqueeze: op_name: str = "Squeeze" @@ -474,3 +490,9 @@ class OpTransposeConv2d: param_pad_amount: str = "pad_amount" param_group: str = "group" param_output_padding: str = "output_padding" + + +@dataclass(init=False, frozen=True) +class OpUnpack: + op_name: str = "UnPack" + param_axis: str = "axis" diff --git a/backends/qualcomm/partition/common_defs.py b/backends/qualcomm/partition/common_defs.py index 8254bb64db0..b427c59ce07 100644 --- a/backends/qualcomm/partition/common_defs.py +++ b/backends/qualcomm/partition/common_defs.py @@ -5,6 +5,8 @@ # LICENSE file in the root directory of this source tree. import _operator +import torch + from executorch.exir.dialects._ops import ops as exir_ops not_supported_operator = [ @@ -20,6 +22,7 @@ exir_ops.edge.aten.arange.start_step, exir_ops.edge.aten.full.default, exir_ops.edge.aten.full_like.default, + torch.ops.aten.scalar_tensor.default, ] allow_list_operator = [ diff --git a/backends/qualcomm/quantizer/annotators.py b/backends/qualcomm/quantizer/annotators.py index c1e1aa25b08..93af5e86c97 100644 --- a/backends/qualcomm/quantizer/annotators.py +++ b/backends/qualcomm/quantizer/annotators.py @@ -378,6 +378,20 @@ def annotate_sin(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) +@register_annotator([torch.ops.aten.scalar_tensor.default]) +def annotate_scalar_tensor(node: Node, quantization_config: QuantizationConfig) -> None: + if _is_annotated([node]): + return + if _is_float_tensor(node): + # workaround for node with kwargs could not be correctly annotated + node.kwargs = {} + node.meta[QUANT_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map={}, + output_qspec=quantization_config.output_activation, + _annotated=True, + ) + + @register_annotator([torch.ops.aten.tanh.default]) def annotate_tanh(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) @@ -680,6 +694,11 @@ def annotate_sigmoid(node: Node, quantization_config: QuantizationConfig) -> Non ) +@register_annotator([torch.ops.aten.__and__.Tensor]) +def annotate_and(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_binary(node, quantization_config) + + @register_annotator([torch.ops.aten.bitwise_or.Tensor, torch.ops.aten.__or__.Tensor]) def annotate_bitwise_or(node: Node, quantization_config: QuantizationConfig) -> None: annotate_binary(node, quantization_config) @@ -717,6 +736,11 @@ def annotate_transpose(node: Node, quantization_config: QuantizationConfig) -> N annotate_single_in_single_out(node, quantization_config) +@register_annotator([torch.ops.aten.elu.default]) +def annotate_elu(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_single_in_single_out(node, quantization_config) + + @register_annotator([torch.ops.aten.embedding.default]) def annotate_embedding(node: Node, quantization_config: QuantizationConfig) -> None: weight = node.args[0] @@ -763,6 +787,11 @@ def annotate_index_put(node: Node, quantization_config: QuantizationConfig) -> N ) +@register_annotator([torch.ops.aten.exp.default]) +def annotate_exp(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_single_in_single_out(node, quantization_config) + + @register_annotator([torch.ops.aten.expand.default, torch.ops.aten.expand_as.default]) def annotate_expand(node: Node, quantization_config: QuantizationConfig) -> None: annotate_in_out_obs_sharing_op(node, quantization_config) @@ -812,18 +841,28 @@ def annotate_flatten(node: Node, quantization_config: QuantizationConfig) -> Non @register_annotator([torch.ops.aten.stack.default]) def annotate_stack(node: Node, quantization_config: QuantizationConfig) -> None: + input_nodes = node.args[0] + if _is_annotated([node]) or not _is_float_tensor(node): + return + + assert isinstance(input_nodes, Sequence) + + first_input_node = input_nodes[0] input_qspec_map = {} - for input_act in node.args[0]: - assert isinstance(input_act, Node) - input_qspec_map[input_act] = quantization_config.input_activation + assert isinstance(first_input_node, Node) + input_qspec_map[first_input_node] = quantization_config.input_activation + share_qparams_with_input_act0_qspec = SharedQuantizationSpec( + (first_input_node, node) + ) - node_tensor = node.meta.get("val") - if torch.is_tensor(node_tensor) and node_tensor.dtype == torch.int64: - continue + for input_node in input_nodes[1:]: + if input_node not in input_qspec_map: + assert isinstance(input_node, Node) + input_qspec_map[input_node] = share_qparams_with_input_act0_qspec node.meta[QUANT_ANNOTATION_KEY] = QuantizationAnnotation( input_qspec_map=input_qspec_map, - output_qspec=quantization_config.output_activation, + output_qspec=share_qparams_with_input_act0_qspec, _annotated=True, ) @@ -894,6 +933,7 @@ def annotate_bmm(node: Node, quantization_config: QuantizationConfig) -> None: torch.ops.aten.conv2d.default, torch.ops.aten.conv1d.default, torch.ops.aten.conv_transpose2d.input, + torch.ops.aten.conv_transpose1d.default, ] ) def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None: @@ -1059,7 +1099,7 @@ def annotate_layer_norm(node: Node, quantization_config: QuantizationConfig) -> @register_annotator([torch.ops.aten.cat.default, torch.ops.aten.concat.default]) def annotate_cat(node: Node, quantization_config: QuantizationConfig) -> None: input_nodes = node.args[0] - if _is_annotated([node]): + if _is_annotated([node]) or not _is_float_tensor(node): return assert isinstance(input_nodes, Sequence) @@ -1087,23 +1127,28 @@ def annotate_cat(node: Node, quantization_config: QuantizationConfig) -> None: @register_annotator([torch.ops.aten.unbind.int]) def annotate_unbind(node: Node, quantization_config: QuantizationConfig) -> None: - if _is_annotated([node]): + # Seems like unbind.int can be either float or int. Only quant when input is float. + if _is_annotated([node]) or not _is_float_tensor(node.args[0]): return input_qspec_map = {} input_act = node.args[0] assert isinstance(input_act, Node) + share_qparams_with_out_node0_qspec = SharedQuantizationSpec((node.args[0], node)) input_qspec_map[input_act] = quantization_config.input_activation - node_tensor = node.meta.get("val") - if torch.is_tensor(node_tensor) and node_tensor.dtype == torch.int64: - return - node.meta[QUANT_ANNOTATION_KEY] = QuantizationAnnotation( input_qspec_map=input_qspec_map, + output_qspec=share_qparams_with_out_node0_qspec, _annotated=True, ) + for user in node.users: + user.meta[QUANT_ANNOTATION_KEY] = QuantizationAnnotation( + output_qspec=share_qparams_with_out_node0_qspec, + _annotated=True, + ) + @register_annotator([torch.ops.aten.split.Tensor, torch.ops.aten.chunk.default]) def annotate_chunk(node: Node, quantization_config: QuantizationConfig) -> None: @@ -1129,22 +1174,33 @@ def annotate_chunk(node: Node, quantization_config: QuantizationConfig) -> None: @register_annotator([torch.ops.aten.where.self]) def annotate_where(node: Node, quantization_config: QuantizationConfig) -> None: - true_input_act = node.args[1] - false_input_act = node.args[2] if _is_annotated([node]): return - _annotate_input_qspec_map( - node, - true_input_act, - quantization_config.input_activation, - ) + input_qspec_map = {} + for input_node in node.args: + assert isinstance(input_node, Node) + if _is_float_tensor(input_node): + input_qspec_map[input_node] = quantization_config.input_activation - _annotate_input_qspec_map( - node, - false_input_act, - quantization_config.input_activation, + node.meta[QUANT_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=( + quantization_config.output_activation if _is_float_tensor(node) else None + ), + _annotated=True, ) - _annotate_output_qspec(node, quantization_config.output_activation) - _mark_nodes_as_annotated([node]) + +@register_annotator([torch.ops.aten.zeros.default]) +def annotate_zeros(node: Node, quantization_config: QuantizationConfig) -> None: + if _is_annotated([node]) or not _is_float_tensor(node): + return + + # workaround for node with kwargs could not be correctly annotated + node.kwargs = {} + node.meta[QUANT_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map={}, + output_qspec=quantization_config.output_activation, + _annotated=True, + ) diff --git a/backends/qualcomm/quantizer/quantizer.py b/backends/qualcomm/quantizer/quantizer.py index 38570835bea..028ffb69f1d 100644 --- a/backends/qualcomm/quantizer/quantizer.py +++ b/backends/qualcomm/quantizer/quantizer.py @@ -10,12 +10,14 @@ import torch from executorch.backends.qualcomm._passes import ( DecomposeEinsum, + DecomposeExpM1, DecomposeLinalgVectorNorm, DecomposeSilu, LiftConstantScalarOperands, RecomposePixelUnshuffle, ReduceDynamicRange, - ReplaceInfBuffer, + ReplaceArangeArgs, + ReplaceInfValues, ) from executorch.backends.transforms.decompose_sdpa import ( DecomposeScaledDotProductAttention, @@ -273,11 +275,13 @@ def set_per_channel_linear_quant(self, enable: bool) -> None: def transform_for_annotation(self, model: GraphModule) -> GraphModule: model = ReduceDynamicRange()(model).graph_module model = RecomposePixelUnshuffle(quantization_capture=True)(model).graph_module + model = ReplaceArangeArgs()(model).graph_module model = DecomposeScaledDotProductAttention()(model).graph_module model = DecomposeSilu()(model).graph_module model = DecomposeEinsum()(model).graph_module + model = DecomposeExpM1()(model).graph_module model = DecomposeLinalgVectorNorm(aten_dialect_capture=True)(model).graph_module - model = ReplaceInfBuffer()(model).graph_module + model = ReplaceInfValues()(model).graph_module model = LiftConstantScalarOperands()(model).graph_module return model diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index e5a9be8e75b..c3c439261d2 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -8,6 +8,19 @@ # module with related operator only + + +class And(torch.nn.Module): + def __init__(self, pos, neg): + super().__init__() + self.pos = pos + self.neg = neg + + def forward(self, x, y): + bitwise_and = torch.bitwise_and(x, y).bool() + return torch.where(bitwise_and, self.pos, self.neg) + + class Abs(torch.nn.Module): def __init__(self): super().__init__() @@ -462,6 +475,17 @@ def forward(self, x): return self.conv(x) +class ConvTranspose1dSingle(torch.nn.Module): + def __init__(self, bias=True): + super().__init__() + self.conv_transpose = torch.nn.ConvTranspose1d( + in_channels=1, out_channels=3, kernel_size=3, stride=2, padding=1, bias=bias + ) + + def forward(self, x): + return self.conv_transpose(x) + + class ConvTranspose2dSingle(torch.nn.Module): def __init__(self, bias=True): super().__init__() @@ -601,6 +625,15 @@ def forward(self, i, j): return torch.relu(torch.einsum("i,j->ij", i, j)) +class Elu(torch.nn.Module): + def __init__(self): + super().__init__() + self.elu = torch.nn.ELU(alpha=0.5) + + def forward(self, i): + return self.elu(i) + + class Embedding(torch.nn.Module): def __init__(self): super().__init__() @@ -645,6 +678,14 @@ def forward(self, x): return y.expand_as(x) +class ExpM1(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.special.expm1(x) + + class Full(torch.nn.Module): def __init__(self, fill, shape): super().__init__() @@ -1383,8 +1424,8 @@ class Stack(torch.nn.Module): def __init__(self): super().__init__() - def forward(self, x, y): - return torch.stack((x, y)) + def forward(self, x, y, z): + return torch.stack((x, y, z)) class Sub(torch.nn.Module): @@ -1493,3 +1534,29 @@ def __init__(self, pos, neg): def forward(self, x): return torch.where(x >= torch.zeros(x.shape), self.pos, self.neg) + + +class WhereConstantOther(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.where(x >= 0, torch.ones(x.shape), 0) + + +class WhereConstantAll(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.where(x >= 0, 1, 0) + + +class WhereConstantInf(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.nn.functional.softmax( + torch.where(x >= 0, 0.1, float("-inf")), dim=-1 + ) diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 936b9c3efe4..05e368f372e 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -195,6 +195,16 @@ def test_qnn_backend_conv2d_channel_last(self): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv_transpose1d(self): + modules = [ + ConvTranspose1dSingle(), # noqa: F405 + ConvTranspose1dSingle(bias=False), # noqa: F405 + ] + sample_input = (torch.randn([1, 1, 3]),) + for i, module in enumerate(modules): + with self.subTest(i=i): + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv_transpose2d(self): modules = [ ConvTranspose2dSingle(), # noqa: F405 @@ -255,6 +265,14 @@ def test_qnn_backend_element_wise_add(self): self.lower_module_and_test_output(module, sample_input) index += 1 + def test_qnn_backend_element_wise_and(self): + module = And(torch.tensor(1.7), torch.tensor(0.2)) # noqa: F405 + sample_input = ( + torch.tensor([1, 0, 1, 0], dtype=torch.bool), + torch.tensor([1, 1, 0, 0], dtype=torch.bool), + ) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_element_wise_ceil(self): module = Ceil() # noqa: F405 sample_input = (torch.randn([2, 5, 1, 3]),) @@ -369,6 +387,12 @@ def test_qnn_backend_element_wise_sub(self): self.lower_module_and_test_output(module, sample_input) index += 1 + @unittest.expectedFailure + def test_qnn_backend_elu(self): + module = Elu() # noqa: F405 + sample_input = (torch.randn(2, 5, 1, 3),) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_embedding(self): module = Embedding() # noqa: F405 sample_input = (torch.Tensor([[1, 2, 4, 5], [4, 3, 2, 9]]).to(torch.int32),) @@ -398,6 +422,11 @@ def test_qnn_backend_expand(self): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_expm1(self): + sample_input = (torch.randn(3, 4, 5),) + module = ExpM1() # noqa: F405 + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_full(self): shape = (1, 2, 3, 4) module = Full(0.5, shape) # noqa: F405 @@ -758,7 +787,11 @@ def test_qnn_backend_slice_copy(self): def test_qnn_backend_stack(self): module = Stack() # noqa: F405 - sample_input = (torch.randn([1, 2, 3, 4]), torch.randn([1, 2, 3, 4])) + sample_input = ( + torch.randn([1, 2, 3, 4]), + torch.randn([1, 2, 3, 4]), + torch.randn([1, 2, 3, 4]), + ) self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_softmax(self): @@ -800,10 +833,16 @@ def test_qnn_backend_where(self): modules = [ Where(), # noqa: F405 WhereConstant(torch.randn(3, 2), torch.randn(3, 2)), # noqa: F405 + WhereConstantOther(), # noqa: F405 + # WhereConstantAll(), # noqa: F405 TODO: constant dtype does not propogate when doing const i64->32, causing where to fail since where does not support int64 output + WhereConstantInf(), # noqa: F405 ] sample_inputs = [ (torch.randn(3, 2), torch.randn(3, 2), torch.randn(3, 2)), (torch.randn(3, 2),), + (torch.randn(3, 2),), + # (torch.randn(3, 2),), + (torch.randn(30, 20),), ] for i, module in enumerate(modules): self.lower_module_and_test_output(module, sample_inputs[i]) @@ -1206,6 +1245,17 @@ def test_qnn_backend_conv2d_channel_last(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv_transpose1d(self): + modules = [ + ConvTranspose1dSingle(), # noqa: F405 + ConvTranspose1dSingle(bias=False), # noqa: F405 + ] + sample_input = (torch.randn([1, 1, 3]),) + for i, module in enumerate(modules): + with self.subTest(i=i): + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv_transpose2d(self): modules = [ ConvTranspose2dSingle(), # noqa: F405 @@ -1271,6 +1321,15 @@ def test_qnn_backend_element_wise_add(self): self.lower_module_and_test_output(module, sample_input) index += 1 + def test_qnn_backend_element_wise_and(self): + module = And(torch.tensor(1.7), torch.tensor(0.2)) # noqa: F405 + sample_input = ( + torch.tensor([1, 0, 1, 0], dtype=torch.bool), + torch.tensor([1, 1, 0, 0], dtype=torch.bool), + ) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_element_wise_ceil(self): module = Ceil() # noqa: F405 sample_input = (torch.randn([2, 5, 1, 3]),) @@ -1391,6 +1450,12 @@ def test_qnn_backend_element_wise_sub(self): self.lower_module_and_test_output(module, sample_input) index += 1 + def test_qnn_backend_elu(self): + module = Elu() # noqa: F405 + sample_input = (torch.randn(2, 5, 1, 3),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_embedding(self): module = Embedding() # noqa: F405 sample_input = (torch.Tensor([[1, 2, 4, 5], [4, 3, 2, 9]]).to(torch.int32),) @@ -1423,6 +1488,12 @@ def test_qnn_backend_expand(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_expm1(self): + sample_input = (torch.randn(3, 4, 5),) + module = ExpM1() # noqa: F405 + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_full(self): shape = (1, 2, 3, 4) module = Full(0.5, shape) # noqa: F405 @@ -1856,6 +1927,7 @@ def test_qnn_backend_stack(self): sample_input = ( torch.randn([1, 2, 3, 4]), torch.randn([1, 2, 3, 4]), + torch.randn([1, 2, 3, 4]), ) module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) @@ -1894,10 +1966,16 @@ def test_qnn_backend_where(self): modules = [ Where(), # noqa: F405 WhereConstant(torch.randn(3, 2), torch.randn(3, 2)), # noqa: F405 + WhereConstantOther(), # noqa: F405 + # WhereConstantAll(), # noqa: F405, TODO: constant dtype does not propogate when doing const i64->32, causing where to fail since where does not support int64 output + WhereConstantInf(), # noqa: F405 ] sample_inputs = [ (torch.randn(3, 2), torch.randn(3, 2), torch.randn(3, 2)), (torch.randn(3, 2),), + (torch.randn(3, 2),), + # (torch.randn(3, 2),), + (torch.randn(30, 20),), ] for i, module in enumerate(modules): module = self.get_qdq_module(module, sample_inputs[i]) diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 8045e9e6443..67b08034439 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -21,8 +21,10 @@ AnnotateQuantAttrs, ConstantI64toI32, ConvertBmmToMatmul, + ConvertConv1dToConv2d, ConvertToLinear, DecomposeAny, + DecomposeExpM1, DecomposeLinalgVectorNorm, ExpandBroadcastTensorShape, FoldQDQ, @@ -326,6 +328,7 @@ def get_decomp_table() -> Dict[torch._ops.OperatorBase, Callable]: # The below super ops are supported by QNN skip_decompositions = [ torch.ops.aten.adaptive_avg_pool2d.default, + torch.ops.aten.elu.default, torch.ops.aten.instance_norm.default, torch.ops.aten.pixel_shuffle.default, torch.ops.aten.pixel_unshuffle.default, @@ -334,6 +337,8 @@ def get_decomp_table() -> Dict[torch._ops.OperatorBase, Callable]: torch.ops.pt2e_quant.quantize_affine.default, torch.ops.pt2e_quant.dequantize_affine.default, torch.ops.aten._safe_softmax.default, + torch.ops.aten.stack.default, # TODO: Might need to remove this later due to Mimi. QNN does not support int io for stack op. + torch.ops.aten.unbind.int, ] remove_decompositions(source_decompositions, skip_decompositions) @@ -353,10 +358,11 @@ def get_capture_program_passes(): # The second value in each tuple in `default_passes_and_setting` indicates whether the corresponding pass is activated by default. # If a pass is activated, it will be executed by default. default_passes_and_setting = [ - (AnnotateDecomposed, True), + (AnnotateDecomposed, False), (AnnotateQuantAttrs, True), (ConstantI64toI32, True), (ConvertBmmToMatmul, True), + (ConvertConv1dToConv2d, True), (ConvertToLinear, True), (DecomposeAny, True), (DecomposeLinalgVectorNorm, True), @@ -448,6 +454,7 @@ def _preprocess_module(module: torch.nn.Module, inputs: Tuple[torch.Tensor]): module = torch.export.export(module, inputs, strict=True).module() module = DecomposeScaledDotProductAttention()(module).graph_module module = DecomposeLinalgVectorNorm(True)(module).graph_module + module = DecomposeExpM1()(module).graph_module module = LiftConstantScalarOperands()(module).graph_module return module @@ -460,6 +467,7 @@ def capture_program( ) -> exir.ExirExportedProgram: module = _preprocess_module(module, inputs) ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True) + # TODO: Handle stack op. If we want to run annotate_decomposed pass for stack op, we need to make stack op decompose, which means we need to find a method to remove it from skip_decomp table decomposed_ep = ep.run_decompositions(get_decomp_table()) core_ep = ExirExportedProgram(decomposed_ep, False) core_ep.transform(TensorI64toI32(edge_program=core_ep)) From f277d442b0d5a496ec4b8d597f1992d6c8a300ca Mon Sep 17 00:00:00 2001 From: winskuo-quic Date: Tue, 25 Mar 2025 17:34:38 +0800 Subject: [PATCH 2/2] Skip decompose op --- backends/qualcomm/_passes/annotate_decomposed.py | 4 +++- backends/qualcomm/utils/utils.py | 13 ++++++++++--- examples/models/llama/export_llama_lib.py | 14 ++++++++++++-- 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/backends/qualcomm/_passes/annotate_decomposed.py b/backends/qualcomm/_passes/annotate_decomposed.py index a8a757ce9bf..918b705e5e9 100644 --- a/backends/qualcomm/_passes/annotate_decomposed.py +++ b/backends/qualcomm/_passes/annotate_decomposed.py @@ -17,6 +17,8 @@ class AnnotateDecomposed(ExportPass): generated after quantization process. """ + decomp_ops = [torch.ops.aten.stack.default, torch.ops.aten.unbind.int] + def __init__(self, edge_program: torch.export.ExportedProgram): super(AnnotateDecomposed, self).__init__() self.edge_program = edge_program @@ -32,7 +34,7 @@ def _annotate_unbind(self, graph_module: torch.fx.GraphModule): n.meta[QCOM_QUANT_ATTRS] = quant_attrs.copy() def _annotate_stack(self, graph_module: torch.fx.GraphModule): - partitions = get_source_partitions(graph_module.graph, [torch.stack]) + partitions = get_source_partitions(graph_module.graph, [torch.stack, "stack"]) for _, src_partitions in partitions.items(): for src_partition in src_partitions: output = src_partition.output_nodes[0] diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 67b08034439..7033f30997a 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -323,7 +323,7 @@ def canonicalize_program(obj): update_spill_fill_size(obj) -def get_decomp_table() -> Dict[torch._ops.OperatorBase, Callable]: +def get_decomp_table(passes_job) -> Dict[torch._ops.OperatorBase, Callable]: source_decompositions = core_aten_decompositions() # The below super ops are supported by QNN skip_decompositions = [ @@ -337,10 +337,17 @@ def get_decomp_table() -> Dict[torch._ops.OperatorBase, Callable]: torch.ops.pt2e_quant.quantize_affine.default, torch.ops.pt2e_quant.dequantize_affine.default, torch.ops.aten._safe_softmax.default, - torch.ops.aten.stack.default, # TODO: Might need to remove this later due to Mimi. QNN does not support int io for stack op. + torch.ops.aten.stack.default, torch.ops.aten.unbind.int, ] + # If we want to annotate the decomposed ops, then we should decompose the operation. + if passes_job and passes_job.get(AnnotateDecomposed, False): + skip_decompositions = [ + skip_decomp_op + for skip_decomp_op in skip_decompositions + if skip_decomp_op not in AnnotateDecomposed.decomp_ops + ] remove_decompositions(source_decompositions, skip_decompositions) return source_decompositions @@ -468,7 +475,7 @@ def capture_program( module = _preprocess_module(module, inputs) ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True) # TODO: Handle stack op. If we want to run annotate_decomposed pass for stack op, we need to make stack op decompose, which means we need to find a method to remove it from skip_decomp table - decomposed_ep = ep.run_decompositions(get_decomp_table()) + decomposed_ep = ep.run_decompositions(get_decomp_table(passes_job)) core_ep = ExirExportedProgram(decomposed_ep, False) core_ep.transform(TensorI64toI32(edge_program=core_ep)) edge_ep = core_ep.to_edge(qnn_edge_config()) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 37a4e6952d8..cfcc68874b1 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -794,9 +794,19 @@ def _to_edge_and_lower_llama( # noqa: C901 ) ) # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.utils.utils` - from executorch.backends.qualcomm.utils.utils import _transform, tag_quant_io + from executorch.backends.qualcomm._passes.annotate_decomposed import ( + AnnotateDecomposed, + ) + from executorch.backends.qualcomm.utils.constants import QCOM_PASS_ACTIVATE_KEY + from executorch.backends.qualcomm.utils.utils import ( + _transform, + get_capture_program_passes, + tag_quant_io, + ) - _transform(builder_exported_to_edge.edge_manager.exported_program()) + passes_job = get_capture_program_passes() + passes_job[AnnotateDecomposed][QCOM_PASS_ACTIVATE_KEY] = True + _transform(builder_exported_to_edge.edge_manager.exported_program(), passes_job) if args.num_sharding > 0: model_sharding.split_graph(