Qualcomm AI Engine Direct - Adding QNN backend support for isInf core ATen op (#18249)

Author: qti-horodnic

backends/qualcomm/builders/README.md

@@ -439,6 +439,7 @@ Please help update following table if you are contributing new operators:
 | GroupNorm | ✓ |
 | HardSwish | ✓ |
 | InstanceNorm | ✓ |
+| IsInf | ✓ |
 | L2Norm | ✗ |
 | LayerNorm | ✓ |
 | LogSoftmax | ✓ |

backends/qualcomm/builders/__init__.py

@@ -55,6 +55,7 @@
     op_index_put,
     op_index_select,
     op_instance_norm,
+    op_is_inf,
     op_layer_norm,
     op_le,
     op_linear,
@@ -164,6 +165,7 @@
     op_index_put,
     op_index_select,
     op_instance_norm,
+    op_is_inf,
     op_layer_norm,
     op_le,
     op_linear,

backends/qualcomm/builders/node_visitor.py

@@ -59,6 +59,7 @@
 }
 QNN_TENSOR_TYPE_MAP = {
     torch.bool: PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_BOOL_8,
+    torch.float16: PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_16,
     torch.float32: PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32,
     # Note that there is no float64 tensor data type in Qnn.
     torch.float64: PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32,

backends/qualcomm/builders/op_is_inf.py

@@ -0,0 +1,80 @@
+# 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 warnings
+from typing import Dict
+
+import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager
+
+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 OpIsInf, QNN_OP_PACKAGE_NAME_QTI_AISW
+
+
+@register_node_visitor
+class IsInf(NodeVisitor):
+    target = ["aten.isinf.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)
+
+        if input_tensor.dtype != torch.float16:
+            warnings.warn(
+                "[QNN Delegate Op Builder]: QNN IsInf only supports FP16 inputs.",
+                stacklevel=1,
+            )
+            return None
+
+        input_tensor_wrapper = self.define_tensor(
+            input_node,
+            node,
+            self.get_tensor(input_node, node),
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+
+        input_tensors = [input_tensor_wrapper]
+
+        out_tensor = self.get_tensor(node, node)
+        output_tensor_wrapper = self.define_tensor(
+            node,
+            node,
+            out_tensor,
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+        output_tensors = [output_tensor_wrapper]
+
+        isinf_op = PyQnnManager.PyQnnOpWrapper(
+            node.name,
+            QNN_OP_PACKAGE_NAME_QTI_AISW,
+            OpIsInf.op_name,
+        )
+        isinf_op.AddInputTensors(input_tensors)
+        isinf_op.AddOutputTensors(output_tensors)
+
+        isinf_op.AddScalarParam(
+            OpIsInf.param_detect_negative,
+            PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_BOOL_8,
+            {QCOM_DATA: True},
+        )
+        isinf_op.AddScalarParam(
+            OpIsInf.param_detect_positive,
+            PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_BOOL_8,
+            {QCOM_DATA: True},
+        )
+
+        return isinf_op

backends/qualcomm/builders/qnn_constants.py

@@ -389,6 +389,13 @@ class OpInstanceNorm:
     param_region = "region"
 
 
+@dataclass(init=False, frozen=True)
+class OpIsInf:
+    op_name: str = "IsInf"
+    param_detect_negative = "detect_negative"
+    param_detect_positive = "detect_positive"
+
+
 @dataclass(init=False, frozen=True)
 class OpLayerNorm:
     op_name: str = "LayerNorm"

backends/qualcomm/tests/models.py

@@ -1310,6 +1310,14 @@ def forward(self, x):
         return self.instance_norm(x)
 
 
+class IsInf(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.isinf(x)
+
+
 class LargeTensorLinear(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1251,6 +1251,25 @@ def test_qnn_backend_instance_norm_2d(self):
             with self.subTest(i=i):
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_is_inf(self):
+        module = IsInf()  # noqa: F405
+        sample_input = (
+            torch.tensor(
+                [
+                    1.1,
+                    float("inf"),
+                    -float("inf"),
+                    0.0,
+                    float("nan"),
+                    0.6,
+                    float("nan"),
+                    -5.0,
+                ],
+                dtype=torch.float16,
+            ),
+        )
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_interpolate_bicubic(self):
         modules = [
             ResizeBicubic([2, 2], None, False),  # noqa: F405