Qualcomm AI Engine Direct - Add QNN HadamardTransform op support

backends/qualcomm/_passes/__init__.py

@@ -60,6 +60,7 @@
 from .layout_transform import LayoutTransform
 from .lift_constant_scalar_operands import LiftConstantScalarOperands
 from .lpai_partition_fallback_support import LpaiPartitionFallbackSupport
+from .recompose_hadamard import RecomposeHadamard
 from .recompose_pad_maxpool2d import RecomposePadMaxPool2d
 from .recompose_pixel_unshuffle import RecomposePixelUnshuffle
 from .recompose_rms_norm import RecomposeRmsNorm
@@ -127,6 +128,7 @@
     LayoutTransform,
     LiftConstantScalarOperands,
     LpaiPartitionFallbackSupport,
+    RecomposeHadamard,
     RecomposePadMaxPool2d,
     RecomposePixelUnshuffle,
     RecomposeRmsNorm,

backends/qualcomm/_passes/backends/htp/qnn_htp_pass_manager.py

@@ -4,7 +4,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from executorch.backends.qualcomm._passes import DecomposeReciprocal, RemoveRedundancy
+from executorch.backends.qualcomm._passes import (
+    DecomposeReciprocal,
+    RecomposeHadamard,
+    RemoveRedundancy,
+)
 from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager
 
 
@@ -33,7 +37,7 @@ def get_passes_dependency_for_capture_program(cls):
 
     @classmethod
     def get_annotation_passes(cls):
-        passes = [DecomposeReciprocal]
+        passes = [DecomposeReciprocal, RecomposeHadamard]
         passes.extend(super().get_annotation_passes())
         return passes
 

backends/qualcomm/_passes/recompose_hadamard.py

@@ -0,0 +1,180 @@
+# 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 operator import attrgetter
+
+import torch
+
+# Also registers torch.ops.qnn_custom.hadamard_transform.
+from executorch.backends.qualcomm.builders.custom_ops import _hadamard_matrix
+from executorch.backends.qualcomm.utils.check_qnn_version import (
+    is_qnn_sdk_version_less_than,
+)
+from executorch.exir.pass_base import ExportPass, PassResult
+from executorch.exir.passes import dead_code_elimination_pass
+
+from .utils import copy_meta
+
+
+def _is_power_of_2_sqare_matrix(weight: torch.Tensor) -> bool:
+    dim = weight.shape[0]
+    # Shape gate: non-square / non-2D / non-power-of-2 weight can never match.
+    return (
+        weight.dim() != 2 or weight.shape[0] != weight.shape[1] or dim & (dim - 1) != 0
+    )
+
+
+def _match_hadamard_weight(weight: torch.Tensor) -> bool:
+    # Returns True if `weight == _hadamard_matrix(dim) * s` for some scale s.
+    # A linear/matmul with such a weight is equivalent to a QNN HadamardTransform.
+    if _is_power_of_2_sqare_matrix(weight):
+        return False
+
+    w = weight.detach().to(torch.float64)
+    nonzero = w[w != 0]
+    if nonzero.numel() == 0:
+        return False
+    # The Hadamard weight is H * s for a single global scale s; infer s from any
+    # nonzero entry (all |H_ij| == 1). For per-channel quant this only matches
+    # when every channel's dequantized scale reconstructs the same H * s.
+    scale = float(nonzero.flatten()[0].abs())
+    hadamard = _hadamard_matrix(w.shape[0], w.device, w.dtype) * scale
+    return torch.allclose(w, hadamard, rtol=0, atol=1e-4)
+
+
+class RecomposeHadamard(ExportPass):
+    """
+    Rewrite a bias-less linear / matmul / 1x1 conv whose weight is a Hadamard
+    matrix into a single qnn_custom.hadamard_transform op, so it is annotated and
+    lowered as a first-class HadamardTransform instead of being detected late in
+    the builder and validated as FullyConnected / MatMul / Conv.
+
+    Runs in the annotation pipeline (before quantization), where the weight is a
+    real tensor and can be inspected. hadamard_transform acts on the last dim, so
+    linear / matmul rewrite directly, while conv (which mixes the channel dim) is
+    wrapped in permutes that move the channel to the last dim and back.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.hadamard_target = torch.ops.qnn_custom.hadamard_transform.default
+
+    def _is_pointwise_conv(self, node, weight: torch.Tensor) -> bool:
+        # Only a 1x1, stride-1, no-pad, dilation-1, groups-1 conv is a pure
+        # channel-mixing matmul equivalent to a Hadamard transform. conv2d args:
+        # (input, weight, bias, stride, padding, dilation, groups) with defaults.
+        stride = node.args[3] if len(node.args) > 3 else [1, 1]
+        padding = node.args[4] if len(node.args) > 4 else [0, 0]
+        dilation = node.args[5] if len(node.args) > 5 else [1, 1]
+        groups = node.args[6] if len(node.args) > 6 else 1
+        return (
+            weight.dim() == 4
+            and all(k == 1 for k in weight.shape[2:])
+            and all(s == 1 for s in stride)
+            and all(p == 0 for p in padding)
+            and all(d == 1 for d in dilation)
+            and groups == 1
+        )
+
+    def _get_hadamard_scale(self, weight: torch.Tensor) -> float:
+        # weight == H * s (all |H_ij| == 1); linear/matmul(x) = x @ H. The op
+        # applies the orthonormal H / sqrt(dim), so fold the remaining factor
+        # s * sqrt(dim) into the op's scale (== 1 for an orthonormal Hadamard).
+        dim = weight.shape[0]
+        return float(weight.detach().abs().flatten()[0]) * (dim**0.5)
+
+    def _rewrite_last_dim(self, graph, node, scale):
+        # linear / matmul already transform the last dim: replace in place.
+        with graph.inserting_before(node):
+            hadamard_node = graph.create_node(
+                "call_function",
+                self.hadamard_target,
+                (node.args[0], scale),
+            )
+        hadamard_node.meta = copy_meta(node.meta)
+        for user in node.users.copy():
+            user.replace_input_with(node, hadamard_node)
+
+    def _rewrite_channel_dim(self, graph, node, scale):
+        # conv mixes the channel dim (dim 1). Move it to the last dim, run the
+        # transform there, then move it back.
+        input_node = node.args[0]
+        input_val = input_node.meta["val"]
+        rank = input_val.dim()
+        to_last = [0, *range(2, rank), 1]
+        from_last = [0, rank - 1, *range(1, rank - 1)]
+        with graph.inserting_before(node):
+            pre = graph.create_node(
+                "call_function", torch.ops.aten.permute.default, (input_node, to_last)
+            )
+            pre.meta = copy_meta(node.meta)
+            pre.meta["val"] = input_val.permute(to_last)
+            hadamard_node = graph.create_node(
+                "call_function", self.hadamard_target, (pre, scale)
+            )
+            hadamard_node.meta = copy_meta(node.meta)
+            post = graph.create_node(
+                "call_function",
+                torch.ops.aten.permute.default,
+                (hadamard_node, from_last),
+            )
+        post.meta = copy_meta(node.meta)
+        for user in node.users.copy():
+            user.replace_input_with(node, post)
+
+    def _is_hadamard_transform(self, graph_module, node):
+        if node.op != "call_function":
+            return False
+
+        is_conv = node.target == torch.ops.aten.conv2d.default
+        is_last_dim = node.target in (
+            torch.ops.aten.linear.default,
+            torch.ops.aten.matmul.default,
+        )
+        if not (is_conv or is_last_dim):
+            return False
+
+        # linear/conv carry an optional bias in args[2]; matmul never does.
+        has_bias = len(node.args) >= 3 and node.args[2] is not None
+        if has_bias:
+            return False
+
+        weight_node = node.args[1]
+        if weight_node.op != "get_attr":
+            return False
+        weight = attrgetter(weight_node.target)(graph_module)
+        if is_conv and not self._is_pointwise_conv(node, weight):
+            return False
+        # A 1x1 conv filter is [out, in, 1, 1]; squeeze to [out, in] to match.
+        squeezed = weight.reshape(weight.shape[:2]) if is_conv else weight
+        if not _match_hadamard_weight(squeezed):
+            return False
+        return True
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        # HadamardTransform is only supported by QNN 2.47+. On older SDKs skip the
+        # rewrite so the op keeps its normal lowering path.
+        if is_qnn_sdk_version_less_than("2.47"):
+            return PassResult(graph_module, False)
+
+        graph = graph_module.graph
+        modified = False
+        for node in graph.nodes:
+            if not self._is_hadamard_transform(graph_module, node):
+                continue
+            weight_node = node.args[1]
+            weight = attrgetter(weight_node.target)(graph_module)
+            is_conv = node.target == torch.ops.aten.conv2d.default
+            squeezed = weight.reshape(weight.shape[:2]) if is_conv else weight
+            scale = self._get_hadamard_scale(squeezed)
+            if is_conv:
+                self._rewrite_channel_dim(graph, node, scale)
+            else:
+                self._rewrite_last_dim(graph, node, scale)
+            modified = True
+
+        if modified:
+            dead_code_elimination_pass(graph_module)
+        return PassResult(graph_module, modified)

backends/qualcomm/builders/README.md

@@ -436,6 +436,7 @@ Please help update following table if you are contributing new operators:
 | GetSparseValues | ✗ |
 | GridSample | ✓ |
 | GroupNorm | ✓ |
+| HadamardTransform | ✓ |
 | HardSwish | ✓ |
 | InstanceNorm | ✓ |
 | IsInf | ✓ |

backends/qualcomm/builders/__init__.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from . import (
+    custom_ops,
     node_visitor,
     op_abs,
     op_adaptive_avg_pool2d,
@@ -49,6 +50,7 @@
     op_grid_sampler_2d,
     op_group_norm,
     op_gt,
+    op_hadamard_transform,
     op_hardsigmoid,
     op_hardswish,
     op_hardtanh,
@@ -121,6 +123,7 @@
 )
 
 __all__ = [
+    custom_ops,
     node_visitor,
     op_abs,
     op_adaptive_avg_pool2d,
@@ -165,6 +168,7 @@
     op_grid_sampler_2d,
     op_group_norm,
     op_gt,
+    op_hadamard_transform,
     op_hardswish,
     op_hardtanh,
     op_hardsigmoid,

backends/qualcomm/builders/custom_ops.py

@@ -0,0 +1,37 @@
+# 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 torch.library import impl, Library, register_fake
+
+# Dedicated namespace, separate from the "qaisw" context-binary namespace.
+hadamard_op_lib = Library("qnn_custom", "DEF")
+hadamard_op_lib.define("hadamard_transform(Tensor input, float scale) -> Tensor")
+
+
+def _hadamard_matrix(dim: int, device, dtype) -> torch.Tensor:
+    # Sylvester construction of the (unnormalized, ±1) Hadamard matrix.
+    h = torch.ones((1, 1), device=device, dtype=dtype)
+    while h.shape[0] < dim:
+        h = torch.cat([torch.cat([h, h], dim=1), torch.cat([h, -h], dim=1)], dim=0)
+    return h
+
+
+@impl(hadamard_op_lib, "hadamard_transform", "CompositeExplicitAutograd")
+def hadamard_transform_impl(input: torch.Tensor, scale: float) -> torch.Tensor:
+    # Normalized Walsh-Hadamard transform along the last dim, times scale.
+    # Matches a linear/matmul whose weight is scipy.linalg.hadamard(dim) * s,
+    # where the rewrite pass sets scale = s * sqrt(dim) (scale == 1 when the
+    # weight is the orthonormal H / sqrt(dim)).
+    dim = input.shape[-1]
+    h = _hadamard_matrix(dim, input.device, input.dtype)
+    return torch.matmul(input, h) * (scale / (dim**0.5))
+
+
+@register_fake("qnn_custom::hadamard_transform")
+def hadamard_transform_fake(input: torch.Tensor, scale: float) -> torch.Tensor:
+    # Hadamard weight is square, so the transform preserves shape.
+    return torch.empty_like(input)

backends/qualcomm/builders/op_hadamard_transform.py

@@ -0,0 +1,66 @@
+# 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.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 OpHadamardTransform, QNN_OP_PACKAGE_NAME_QTI_AISW
+
+
+@register_node_visitor
+class HadamardTransformVisitor(NodeVisitor):
+    target = ["qnn_custom.hadamard_transform.default"]
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
+    ) -> PyQnnManager.PyQnnOpWrapper:
+        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,
+        )
+
+        output_tensor = self.get_tensor(node, node)
+        output_tensor_wrapper = self.define_tensor(
+            node,
+            node,
+            output_tensor,
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+
+        hadamard_op = PyQnnManager.PyQnnOpWrapper(
+            node.name,
+            QNN_OP_PACKAGE_NAME_QTI_AISW,
+            OpHadamardTransform.op_name,
+        )
+        hadamard_op.AddInputTensors([input_tensor_wrapper])
+        hadamard_op.AddOutputTensors([output_tensor_wrapper])
+
+        scale = node.args[1]
+        hadamard_op.AddScalarParam(
+            OpHadamardTransform.param_scale,
+            PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32,
+            {QCOM_DATA: np.float32(scale)},
+        )
+        return hadamard_op

backends/qualcomm/builders/qnn_constants.py

@@ -382,6 +382,12 @@ class OpGroupNorm:
     param_group = "group"
 
 
+@dataclass(init=False, frozen=True)
+class OpHadamardTransform:
+    op_name: str = "HadamardTransform"
+    param_scale: str = "scale"
+
+
 @dataclass(init=False, frozen=True)
 class OpHardSwish:
     op_name: str = "HardSwish"