diff --git a/backends/qualcomm/_passes/recompose_hadamard.py b/backends/qualcomm/_passes/recompose_hadamard.py index 6f5d5339d4c..31a2c1e0283 100644 --- a/backends/qualcomm/_passes/recompose_hadamard.py +++ b/backends/qualcomm/_passes/recompose_hadamard.py @@ -116,6 +116,7 @@ def _rewrite_channel_dim(self, graph, node, scale): "call_function", self.hadamard_target, (pre, scale) ) hadamard_node.meta = copy_meta(node.meta) + hadamard_node.meta["val"] = pre.meta["val"] post = graph.create_node( "call_function", torch.ops.aten.permute.default, @@ -155,9 +156,12 @@ def _is_hadamard_transform(self, graph_module, node): 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 not os.environ.get("QNN_SDK_ROOT") or is_qnn_sdk_version_less_than("2.47"): + try: + if not os.environ.get("QNN_SDK_ROOT") or is_qnn_sdk_version_less_than( + "2.47" + ): + return PassResult(graph_module, False) + except Exception: return PassResult(graph_module, False) graph = graph_module.graph diff --git a/backends/qualcomm/builders/custom_ops.py b/backends/qualcomm/builders/custom_ops.py index 016e7bacd73..16ead2de5d7 100644 --- a/backends/qualcomm/builders/custom_ops.py +++ b/backends/qualcomm/builders/custom_ops.py @@ -7,10 +7,6 @@ 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. @@ -20,18 +16,26 @@ def _hadamard_matrix(dim: int, device, dtype) -> torch.Tensor: 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) +if not hasattr(torch.ops, "qnn_custom") or not hasattr( + torch.ops.qnn_custom, "hadamard_transform" +): + hadamard_op_lib = Library("qnn_custom", "DEF") + hadamard_op_lib.define("hadamard_transform(Tensor input, float scale) -> Tensor") + + @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) + +else: + hadamard_op_lib = Library("qnn_custom", "FRAGMENT") diff --git a/backends/qualcomm/utils/check_qnn_version.py b/backends/qualcomm/utils/check_qnn_version.py index b2342d312bf..d38629a78b1 100644 --- a/backends/qualcomm/utils/check_qnn_version.py +++ b/backends/qualcomm/utils/check_qnn_version.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import functools import os import platform import re @@ -25,25 +26,31 @@ def _get_qnn_host_lib_dir_name() -> str: return "x86_64-linux-clang" -def get_sdk_build_id(): - qnn_sdk_root = os.environ.get("QNN_SDK_ROOT") - if not qnn_sdk_root: - raise EnvironmentError( - "QNN_SDK_ROOT must be set to query the QNN SDK build id." - ) +@functools.cache +def _get_sdk_build_id(qnn_sdk_root: str): htp_library_path = os.path.join( qnn_sdk_root, "lib", _get_qnn_host_lib_dir_name(), get_qnn_lib_name("QnnHtp"), ) - # The GetQnnSdkBuildId API can be used without needing to create a backend first, so it works regardless of which backend is used. - sdk_build_id = PyQnnManagerAdaptor.GetQnnSdkBuildId(htp_library_path) - return sdk_build_id + return PyQnnManagerAdaptor.GetQnnSdkBuildId(htp_library_path) + + +def get_sdk_build_id(): + qnn_sdk_root = os.environ.get("QNN_SDK_ROOT") + if not qnn_sdk_root: + raise EnvironmentError( + "QNN_SDK_ROOT must be set to query the QNN SDK build id." + ) + return _get_sdk_build_id(qnn_sdk_root) def is_qnn_sdk_version_less_than(target_version): - current_version = get_sdk_build_id() + try: + current_version = get_sdk_build_id() + except Exception: + return True match = re.search(r"v(\d+)\.(\d+)", current_version) if match: @@ -59,7 +66,10 @@ def is_qnn_sdk_version_less_than(target_version): def is_qnn_sdk_version_greater_than(target_version): - current_version = get_sdk_build_id() + try: + current_version = get_sdk_build_id() + except Exception: + return False match = re.search(r"v(\d+)\.(\d+)", current_version) if match: