Qualcomm AI Engine Direct - Exact integer floor_divide and remainder

backends/qualcomm/_passes/__init__.py

@@ -29,7 +29,6 @@
 from .decompose_empty import DecomposeEmpty
 from .decompose_expm1 import DecomposeExpM1
 from .decompose_fill import DecomposeFill
-from .decompose_floor_divide import DecomposeFloorDivide
 from .decompose_glu import DecomposeGlu
 from .decompose_hardsigmoid import DecomposeHardsigmoid
 from .decompose_hyperbolic_variants import DecomposeHyperbolicVariants
@@ -102,7 +101,6 @@
     DecomposeEmpty,
     DecomposeExpM1,
     DecomposeFill,
-    DecomposeFloorDivide,
     DecomposeGlu,
     DecomposeHyperbolicVariants,
     DecomposeHardsigmoid,

backends/qualcomm/_passes/decompose_div_mode.py

@@ -44,6 +44,12 @@ def call(self, graph_module: torch.fx.GraphModule):
                 rounding_mode = node.kwargs.get("rounding_mode", None)
                 if rounding_mode is None and len(node.args) > 2:
                     rounding_mode = node.args[2]
+                if (
+                    rounding_mode == "floor"
+                    and not node.meta["val"].is_floating_point()
+                ):
+                    # integer floor division is exact in int32 (op_floor_divide)
+                    continue
 
                 div_op = (
                     exir_ops.edge.aten.div.Tensor

backends/qualcomm/_passes/decompose_remainder.py

@@ -36,7 +36,12 @@ def call(self, graph_module: torch.fx.GraphModule):
         const_cache = {}
 
         for node in list(graph.nodes):
-            if node.op == "call_function" and node.target in self.remainder_targets:
+            if (
+                node.op == "call_function"
+                and node.target in self.remainder_targets
+                # integer remainder is exact as QNN Mod (op_remainder)
+                and node.meta["val"].is_floating_point()
+            ):
                 x_arg = node.args[0]
                 y_arg = node.args[1]
                 is_edge = isinstance(node.target, EdgeOpOverload)

backends/qualcomm/_passes/layout_transform.py

@@ -82,6 +82,7 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.constant_pad_nd.default,
         exir_ops.edge.aten.cumsum.default,
         exir_ops.edge.aten.div.Tensor,
+        exir_ops.edge.aten.div.Tensor_mode,
         exir_ops.edge.aten.elu.default,
         exir_ops.edge.aten.eq.Tensor,
         exir_ops.edge.aten.exp.default,

backends/qualcomm/_passes/lift_constant_scalar_operands.py

@@ -50,12 +50,15 @@ class TensorOpInfo:
     aten.add.Tensor: TensorOpInfo(aten.add.Tensor, False, False),
     aten.div.Scalar: TensorOpInfo(aten.div.Tensor, False, False),
     aten.div.Scalar_mode: TensorOpInfo(aten.div.Tensor_mode, False, False),
+    aten.floor_divide.default: TensorOpInfo(aten.floor_divide.default, False, False),
     aten.mul.Scalar: TensorOpInfo(aten.mul.Tensor, False, False),
     aten.rsub.Scalar: TensorOpInfo(aten.rsub.Tensor, False, False),
     aten.sub.Scalar: TensorOpInfo(aten.sub.Tensor, False, False),
     aten.sub.Tensor: TensorOpInfo(aten.sub.Tensor, False, False),
     aten.pow.Tensor_Scalar: TensorOpInfo(aten.pow.Tensor_Tensor, False, False),
     aten.pow.Scalar: TensorOpInfo(aten.pow.Tensor_Tensor, False, False),
+    aten.remainder.Scalar: TensorOpInfo(aten.remainder.Tensor, False, False),
+    aten.remainder.Tensor: TensorOpInfo(aten.remainder.Tensor, False, False),
     # The scalar number arg[1] is missing when using default. Result in a corner case to deal
     aten.leaky_relu.default: TensorOpInfo(aten.prelu.default, True, False),
     aten.leaky_relu_.default: TensorOpInfo(aten.prelu.default, True, False),

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -34,7 +34,6 @@
     DecomposeEmpty,
     DecomposeExpM1,
     DecomposeFill,
-    DecomposeFloorDivide,
     DecomposeGlu,
     DecomposeHyperbolicVariants,
     DecomposeLinalgVectorNorm,
@@ -232,11 +231,6 @@ def get_export_passes(cls):
             DecomposeExpM1,
             DecomposeFill,
             DecomposeVar,
-            # DecomposeFloorDivide does not apply to the annotation pipeline,
-            # since the CPU QDQ model would reduce accuracy.
-            # We keep div and floor operations in floating-point to maintain precision.
-            # This pass is needed before to_edge pipeline to avoid mixed type for div operator with RemoveMixedTypeOperators pass.
-            DecomposeFloorDivide,
             DecomposeWrapWithAutocast,
             ConvertSquareToPow,
             LiftConstantScalarOperands,

backends/qualcomm/builders/README.md

@@ -518,7 +518,6 @@ The following PyTorch operators are supported through decomposition or annotatio
 | `aten.empty.memory_format`, `aten.empty_strided` | `DecomposeEmpty` |
 | `aten.special_expm1` | `DecomposeExpM1` |
 | `aten.fill.Scalar` | `DecomposeFill` |
-| `aten.floor_divide` | `DecomposeFloorDivide` |
 | `aten.glu` | `DecomposeGlu` |
 | `aten.linalg_vector_norm` | `DecomposeLinalgVectorNorm` |
 | `aten.log10`, `aten.log2`, `aten.log1p` | `DecomposeLogVariants` |

backends/qualcomm/builders/__init__.py

@@ -42,6 +42,7 @@
     op_expand,
     op_flip,
     op_floor,
+    op_floor_divide,
     op_full,
     op_full_like,
     op_gather,
@@ -85,6 +86,7 @@
     op_rand,
     op_randn,
     op_relu,
+    op_remainder,
     op_repeat,
     op_reshape,
     op_resize,
@@ -160,6 +162,7 @@
     op_expand,
     op_flip,
     op_floor,
+    op_floor_divide,
     op_full,
     op_full_like,
     op_gather,
@@ -203,6 +206,7 @@
     op_rand,
     op_randn,
     op_relu,
+    op_remainder,
     op_repeat,
     op_reshape,
     op_resize,

backends/qualcomm/builders/op_binary.py

@@ -14,6 +14,7 @@
 
 from .node_visitor import NodeVisitor
 from .node_visitor_manager import register_node_visitor
+from .op_floor_divide import define_int_floor_divide
 from .qnn_constants import OpElementWiseBinary, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 
@@ -35,6 +36,9 @@ def define_node(
         node: torch.fx.Node,
         nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
     ) -> PyQnnManager.PyQnnOpWrapper:
+        # QNN has no integer FloorDiv.
+        if not node.meta["val"].is_floating_point():
+            return define_int_floor_divide(self, node, nodes_to_wrappers)
         out_tensor = self.get_tensor(node, node)
         output_tensor_wrapper = self.define_tensor(
             node,

backends/qualcomm/builders/op_floor_divide.py

@@ -0,0 +1,106 @@
+# 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 torch
+
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
+from .qnn_constants import (
+    OpElementWiseDivide,
+    OpElementWiseMod,
+    OpElementWiseSubtract,
+    QNN_OP_PACKAGE_NAME_QTI_AISW,
+)
+
+
+def define_int_floor_divide(
+    visitor: NodeVisitor,
+    node: torch.fx.Node,
+    nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
+):
+    """Integer floor division, exact in int32: (x - x mod y) / y.
+
+    x - x mod y is a multiple of y, so the integer Divide's rounding mode does not
+    matter, and Mod takes the divisor's sign like torch.
+    """
+    out_tensor = visitor.get_tensor(node, node)
+    x, y = (
+        visitor.define_tensor(
+            input_node,
+            node,
+            visitor.get_tensor(input_node, node),
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+        for input_node in (visitor.get_node(arg) for arg in node.args[:2])
+    )
+    quant_encoding, quant_configs = visitor.get_quant_encoding_conf(node, node)
+
+    def intermediate(suffix):
+        return visitor.define_custom_tensor_wrapper(
+            node_name=f"{node.name}_{suffix}",
+            tensor_type=PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            dtype=visitor.get_data_type(out_tensor, quant_configs),
+            quant_encoding=quant_encoding,
+            quant_configs=quant_configs,
+            dims=out_tensor.size(),
+            tensor=out_tensor,
+            is_fake_tensor=True,
+            nodes_to_wrappers=nodes_to_wrappers,
+        )
+
+    def op(name, op_name, inputs, outputs):
+        wrapper = PyQnnManager.PyQnnOpWrapper(
+            name, QNN_OP_PACKAGE_NAME_QTI_AISW, op_name
+        )
+        wrapper.AddInputTensors(inputs)
+        wrapper.AddOutputTensors(outputs)
+        return wrapper
+
+    remainder, multiple = intermediate("mod"), intermediate("sub")
+    output = visitor.define_tensor(
+        node,
+        node,
+        out_tensor,
+        PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+        nodes_to_wrappers,
+    )
+    return [
+        op(f"{node.name}_mod", OpElementWiseMod.op_name, [x, y], [remainder]),
+        op(
+            f"{node.name}_sub",
+            OpElementWiseSubtract.op_name,
+            [x, remainder],
+            [multiple],
+        ),
+        op(node.name, OpElementWiseDivide.op_name, [multiple, y], [output]),
+    ]
+
+
+@register_node_visitor
+class IntFloorDivide(NodeVisitor):
+    """div with rounding_mode="floor" on integers. Integer floor_divide is handled in
+    op_binary; float and "trunc" division are decomposed by DecomposeDivMode."""
+
+    target = ["aten.div.Tensor_mode"]
+
+    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:
+        rounding_mode = node.kwargs.get(
+            "rounding_mode", node.args[2] if len(node.args) > 2 else None
+        )
+        if node.meta["val"].is_floating_point() or rounding_mode != "floor":
+            return None
+        return define_int_floor_divide(self, node, nodes_to_wrappers)

backends/qualcomm/builders/op_remainder.py

@@ -0,0 +1,62 @@
+# 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 torch
+
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
+from .qnn_constants import OpElementWiseMod, QNN_OP_PACKAGE_NAME_QTI_AISW
+
+
+@register_node_visitor
+class Remainder(NodeVisitor):
+    """Integer remainder as QNN Mod, whose result takes the divisor's sign like torch.
+
+    Float remainder is decomposed by DecomposeRemainder; Mod has no float support.
+    """
+
+    target = ["aten.remainder.Tensor"]
+
+    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:
+        out_tensor = self.get_tensor(node, node)
+        if out_tensor.is_floating_point():
+            return None
+
+        input_tensors = []
+        for index in range(2):
+            input_node = self.get_node(node.args[index])
+            input_tensors.append(
+                self.define_tensor(
+                    input_node,
+                    node,
+                    self.get_tensor(input_node, node),
+                    PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+                    nodes_to_wrappers,
+                )
+            )
+        output_tensor = self.define_tensor(
+            node,
+            node,
+            out_tensor,
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+        mod_op = PyQnnManager.PyQnnOpWrapper(
+            node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpElementWiseMod.op_name
+        )
+        mod_op.AddInputTensors(input_tensors)
+        mod_op.AddOutputTensors([output_tensor])
+        return mod_op

backends/qualcomm/builders/qnn_constants.py

@@ -214,6 +214,11 @@ class OpElementWiseMinimum:
     op_name: str = "ElementWiseMinimum"
 
 
+@dataclass(init=False, frozen=True)
+class OpElementWiseMod:
+    op_name: str = "ElementWiseMod"
+
+
 @dataclass(init=False, frozen=True)
 class OpElementWiseMultiply:
     op_name: str = "ElementWiseMultiply"