diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index 81dc1770856..caf311317ab 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -12,6 +12,7 @@ from .build_quant_io import BuildQuantIo from .canonicalize_conv import CanonicalizeConv from .convert_bmm_to_matmul import ConvertBmmToMatmul +from .convert_complex_stft import ConvertComplexStft from .convert_linear_to_conv2d import ConvertLinearToConv2d from .convert_mha_to_sha import ConvertMhaToSha from .convert_square_to_pow import ConvertSquareToPow @@ -85,6 +86,7 @@ BuildQuantIo, CanonicalizeConv, ConvertBmmToMatmul, + ConvertComplexStft, ConvertLinearToConv2d, ConvertMhaToSha, ConvertSquareToPow, diff --git a/backends/qualcomm/_passes/convert_complex_stft.py b/backends/qualcomm/_passes/convert_complex_stft.py new file mode 100644 index 00000000000..e04f68ce999 --- /dev/null +++ b/backends/qualcomm/_passes/convert_complex_stft.py @@ -0,0 +1,116 @@ +# 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 executorch.exir.passes import dead_code_elimination_pass + +from .utils import copy_meta + + +class ConvertComplexStft(ExportPass): + """ + Rewrite a complex-valued stft into its real-valued form so the QNN Stft + op, whose output is real with the real/imaginary parts in the last + dimension, can consume it. Only complex consumers that have a real + equivalent are handled: + + view_as_real(stft(x, return_complex=True)) -> stft(x, return_complex=False) + abs(stft(x, return_complex=True)) -> sqrt(sum(stft_real ** 2, -1)) + pow(abs(stft(...)), 2) -> sum(stft_real ** 2, -1) + + An stft with any other complex consumer is left untouched and is + decomposed (and runs on CPU) as before. + """ + + stft_op = torch.ops.aten.stft.default + return_complex_index = 7 + + def __init__(self) -> None: + super().__init__() + + @staticmethod + def _is_square(node: torch.fx.Node) -> bool: + return node.target == torch.ops.aten.pow.Tensor_Scalar and node.args[1] == 2 + + def _create_node(self, graph, target, args, meta, val): + node = graph.create_node("call_function", target, args) + node.meta = copy_meta(meta, callback=lambda m: {**m, "val": val}) + return node + + def _convert(self, graph: torch.fx.Graph, stft: torch.fx.Node) -> None: + args = list(stft.args) + args += [None] * (self.return_complex_index + 1 - len(args)) + args[self.return_complex_index] = False + real_val = torch.view_as_real(stft.meta["val"]) + with graph.inserting_after(stft): + real_stft = self._create_node( + graph, self.stft_op, tuple(args), stft.meta, real_val + ) + + for user in list(stft.users): + if user.target == torch.ops.aten.view_as_real.default: + # stft traces as a [freq, frames] view over frame-major memory, + # but the lowered op writes contiguous output; a reshape traces + # contiguous, so a delegate output keeps the matching dim order. + with graph.inserting_before(user): + contiguous = self._create_node( + graph, + torch.ops.aten.reshape.default, + (real_stft, list(real_val.shape)), + user.meta, + real_val.contiguous(), + ) + user.replace_all_uses_with(contiguous) + continue + # The remaining users are abs, as checked by the caller. + square_val = real_val * real_val + sum_val = square_val.sum(dim=-1) + with graph.inserting_before(user): + square = self._create_node( + graph, + torch.ops.aten.mul.Tensor, + (real_stft, real_stft), + user.meta, + square_val, + ) + power = self._create_node( + graph, + torch.ops.aten.sum.dim_IntList, + (square, [-1]), + user.meta, + sum_val, + ) + for abs_user in list(user.users): + if self._is_square(abs_user): + abs_user.replace_all_uses_with(power) + if user.users: + with graph.inserting_before(user): + magnitude = self._create_node( + graph, + torch.ops.aten.sqrt.default, + (power,), + user.meta, + sum_val.sqrt(), + ) + user.replace_all_uses_with(magnitude) + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + for node in list(graph.nodes): + if ( + node.target == self.stft_op + and node.meta["val"].is_complex() + and all( + user.target + in (torch.ops.aten.abs.default, torch.ops.aten.view_as_real.default) + for user in node.users + ) + ): + self._convert(graph, node) + + dead_code_elimination_pass(graph_module) + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 40af0946366..854c619061c 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -17,6 +17,7 @@ AnnotateUnbind, CanonicalizeConv, ConvertBmmToMatmul, + ConvertComplexStft, ConvertLinearToConv2d, ConvertMhaToSha, ConvertSquareToPow, @@ -217,6 +218,7 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): def get_export_passes(cls): """Return export pipeline pass classes. Override in subclasses to add backend-specific passes.""" passes = [ + ConvertComplexStft, DecomposeAsStrided, DecomposeBinaryAlpha, DecomposeCDist, diff --git a/backends/qualcomm/builders/README.md b/backends/qualcomm/builders/README.md index 164a1c6e379..1f186eff9a9 100644 --- a/backends/qualcomm/builders/README.md +++ b/backends/qualcomm/builders/README.md @@ -483,6 +483,7 @@ Please help update following table if you are contributing new operators: | SparseToDense | ✗ | | Split | ✓ | | Squeeze | ✓ | +| Stft | ✓ | | StridedSlice | ✓ | | Tanh | ✓ | | Tile | ✓ | diff --git a/backends/qualcomm/builders/__init__.py b/backends/qualcomm/builders/__init__.py index f68dac483cd..3dbb0d219b3 100644 --- a/backends/qualcomm/builders/__init__.py +++ b/backends/qualcomm/builders/__init__.py @@ -107,6 +107,7 @@ op_sqrt, op_squeeze, op_stack, + op_stft, op_sub, op_sum_int_list, op_tanh, @@ -225,6 +226,7 @@ op_squeeze, op_sqrt, op_stack, + op_stft, op_sub, op_sum_int_list, op_tanh, diff --git a/backends/qualcomm/builders/op_stft.py b/backends/qualcomm/builders/op_stft.py new file mode 100644 index 00000000000..14617704dbf --- /dev/null +++ b/backends/qualcomm/builders/op_stft.py @@ -0,0 +1,183 @@ +# 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, List + +import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_DATA + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpReshape, OpStft, OpTranspose, QNN_OP_PACKAGE_NAME_QTI_AISW +from .utils import get_parameter + + +@register_node_visitor +class Stft(NodeVisitor): + """ + Lowers the real-valued form of aten.stft (return_complex=False), see + ConvertComplexStft, as Reshape -> Stft -> Transpose [-> Reshape]. + + QNN Stft takes the signal as [batch, signal_length, 1] and produces + [batch, frames, freq, 2], while torch produces [batch, freq, frames, 2] + (without the batch dimension for a 1D signal). The window is zero-padded on + both sides to n_fft, as torch does when win_length < n_fft. + """ + + target = ["aten.stft.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def _define_native_tensor( + self, + name: str, + tensor: torch.Tensor, + source_node: torch.fx.Node, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], + ) -> PyQnnManager.TensorWrapper: + quant_encoding, quant_configs = self.get_quant_encoding_conf(source_node, node) + return self.define_custom_tensor_wrapper( + node_name=name, + tensor_type=PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + dtype=self.get_data_type(tensor, quant_configs), + quant_encoding=quant_encoding, + quant_configs=quant_configs, + dims=tensor.size(), + tensor=tensor, + is_fake_tensor=True, + nodes_to_wrappers=nodes_to_wrappers, + ) + + def _define_window( + self, + node: torch.fx.Node, + n_fft: int, + win_length: int, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], + ) -> List[PyQnnManager.TensorWrapper]: + window_node = node.args[4] if len(node.args) > 4 else None + if window_node is None: + return [] + window_node = self.get_node(window_node) + window = get_parameter(window_node, self.edge_program) + left = (n_fft - win_length) // 2 + window = torch.nn.functional.pad(window, (left, n_fft - win_length - left)) + return [ + self.define_tensor( + window_node, + node, + window, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, + nodes_to_wrappers, + ) + ] + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], + ) -> List[PyQnnManager.PyQnnOpWrapper]: + args = list(node.args) + [None] * (8 - len(node.args)) + n_fft = args[1] + hop_length = args[2] or n_fft // 4 + win_length = args[3] or n_fft + onesided = True if args[6] is None else args[6] + + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + signal = input_tensor.reshape(-1, input_tensor.shape[-1], 1) + signal_name = f"{node.name}_signal" + signal_wrapper = self._define_native_tensor( + signal_name, signal, input_node, node, nodes_to_wrappers + ) + signal_op = PyQnnManager.PyQnnOpWrapper( + signal_name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpReshape.op_name + ) + signal_op.AddInputTensors([input_tensor_wrapper]) + signal_op.AddOutputTensors([signal_wrapper]) + + output_tensor = self.get_tensor(node, node) + # torch layout with the batch dimension: [batch, freq, frames, 2] + spectrum = output_tensor.reshape(-1, *output_tensor.shape[-3:]) + frames_name = f"{node.name}_frames" + frames_wrapper = self._define_native_tensor( + frames_name, spectrum.transpose(1, 2), node, node, nodes_to_wrappers + ) + stft_op = PyQnnManager.PyQnnOpWrapper( + frames_name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpStft.op_name + ) + stft_op.AddInputTensors( + [signal_wrapper] + + self._define_window(node, n_fft, win_length, nodes_to_wrappers) + ) + stft_op.AddOutputTensors([frames_wrapper]) + stft_op.AddScalarParam( + OpStft.param_frame_step, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(hop_length)}, + ) + stft_op.AddScalarParam( + OpStft.param_frame_length, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(n_fft)}, + ) + stft_op.AddScalarParam( + OpStft.param_onesided, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_BOOL_8, + {QCOM_DATA: onesided}, + ) + + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + is_batched = output_tensor.dim() == 4 + transpose_name = node.name if is_batched else f"{node.name}_transpose" + transpose_wrapper = ( + output_tensor_wrapper + if is_batched + else self._define_native_tensor( + transpose_name, spectrum, node, node, nodes_to_wrappers + ) + ) + perm = [0, 2, 1, 3] + transpose_op = PyQnnManager.PyQnnOpWrapper( + transpose_name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpTranspose.op_name + ) + transpose_op.AddInputTensors([frames_wrapper]) + transpose_op.AddOutputTensors([transpose_wrapper]) + transpose_op.AddTensorParam( + OpTranspose.param_perm, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + 1, + [len(perm)], + np.array(perm, dtype=np.uint32), + True, + ) + op_wrapper_list = [signal_op, stft_op, transpose_op] + + if not is_batched: + unbatch_op = PyQnnManager.PyQnnOpWrapper( + node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpReshape.op_name + ) + unbatch_op.AddInputTensors([transpose_wrapper]) + unbatch_op.AddOutputTensors([output_tensor_wrapper]) + op_wrapper_list.append(unbatch_op) + return op_wrapper_list diff --git a/backends/qualcomm/builders/qnn_constants.py b/backends/qualcomm/builders/qnn_constants.py index 0df2e5a4731..0c0c4ce63de 100644 --- a/backends/qualcomm/builders/qnn_constants.py +++ b/backends/qualcomm/builders/qnn_constants.py @@ -654,6 +654,14 @@ class OpSqueeze: op_name: str = "Squeeze" +@dataclass(init=False, frozen=True) +class OpStft: + op_name: str = "Stft" + param_frame_step: str = "frame_step" + param_frame_length: str = "frame_length" + param_onesided: str = "onesided" + + @dataclass(init=False, frozen=True) class OpStridedSlice: op_name: str = "StridedSlice" diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index 93885b4fdf1..5fa5633ac70 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -148,6 +148,19 @@ def filter_fn(node: torch.fx.Node) -> bool: torch.int64, ]: return False + # QNN Stft only produces the real-valued spectrum and has no normalization. + # ConvertComplexStft rewrites the supported complex cases into this form. + # Without a window, QNN's rectangular default spans n_fft rather than + # win_length. + if node.target == torch.ops.aten.stft.default: + args = list(node.args) + [None] * (8 - len(node.args)) + n_fft, win_length, window, normalized = args[1], args[3], args[4], args[5] + return ( + not node.meta["val"].is_complex() + and node.args[0].meta["val"].is_floating_point() + and not normalized + and (window is not None or win_length in (None, n_fft)) + ) return True @@ -176,6 +189,7 @@ def get_skip_decomp_table() -> List[torch._ops.OperatorBase]: torch.ops.aten.scatter_add.default, torch.ops.aten.scatter_reduce.two, torch.ops.aten.stack.default, + torch.ops.aten.stft.default, torch.ops.aten.upsample_bicubic2d.vec, # This request is ignored because it is in a blocklist. Refer to exir/program/_program.py torch.ops.aten.unbind.int, diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 14aff8e7fbb..94d875617c7 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -2864,6 +2864,42 @@ def forward(self, x, y, z): return torch.stack((x, y, z)) +class Stft(torch.nn.Module): + def __init__( + self, + n_fft, + hop_length=None, + win_length=None, + use_window=True, + onesided=True, + output="power", + ): + super().__init__() + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + self.onesided = onesided + self.output = output + self.window = torch.hann_window(win_length or n_fft) if use_window else None + + def forward(self, x): + spectrum = torch.stft( + x, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + window=self.window, + center=False, + onesided=self.onesided, + return_complex=True, + ) + if self.output == "power": + return spectrum.abs().pow(2) + if self.output == "magnitude": + return spectrum.abs() + return torch.view_as_real(spectrum) + + class Sub(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index e9f69a34547..221b7b4dbac 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -2616,6 +2616,39 @@ def test_qnn_backend_softmax(self): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_stft(self): + test_comb = [ + { + # Pre-framed speech front end: one 512-point frame per row, + # with a 400-sample window zero-padded to n_fft. + QCOM_MODULE: Stft( # noqa: F405 + n_fft=512, hop_length=160, win_length=400 + ), + QCOM_SAMPLE_INPUTS: (torch.randn(4, 512) * 0.1,), + }, + { + QCOM_MODULE: Stft( # noqa: F405 + n_fft=256, hop_length=128, output="magnitude" + ), + QCOM_SAMPLE_INPUTS: (torch.randn(1664) * 0.1,), + }, + { + QCOM_MODULE: Stft( # noqa: F405 + n_fft=64, + hop_length=16, + use_window=False, + onesided=False, + output="real", + ), + QCOM_SAMPLE_INPUTS: (torch.randn(2, 256) * 0.1,), + }, + ] + for i, test in enumerate(test_comb): + with self.subTest(i=i): + self.lower_module_and_test_output( + test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS] + ) + def test_qnn_backend_sort(self): modules = [ Conv2dSort(descending=True), # noqa: F405