Add Stft op support to the OSS QNN backend (#24031)

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,

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)

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,

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 | ✓ |

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,

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

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"

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,