Qualcomm AI Engine Direct - Adding QNN backend support for log2, log10, log1p core ATen ops (pytorch#18542)

Author: qti-horodnic

backends/qualcomm/_passes/__init__.py

@@ -22,6 +22,7 @@
 from .decompose_floor_divide import DecomposeFloorDivide
 from .decompose_glu import DecomposeGlu
 from .decompose_linalg_vector_norm import DecomposeLinalgVectorNorm
+from .decompose_log_variants import DecomposeLogVariants
 from .decompose_maxpool3d import DecomposeMaxPool3d
 from .decompose_minmaxdim import DecomposeMinMaxDim
 from .decompose_reciprocal import DecomposeReciprocal
@@ -72,6 +73,7 @@
     DecomposeFloorDivide,
     DecomposeGlu,
     DecomposeLinalgVectorNorm,
+    DecomposeLogVariants,
     DecomposeMaxPool3d,
     DecomposeMinMaxDim,
     DecomposeReciprocal,

backends/qualcomm/_passes/decompose_log_variants.py

@@ -0,0 +1,117 @@
+# 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 math
+from functools import partial
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+
+from .utils import copy_meta, get_const_node
+
+
+class DecomposeLogVariants(ExportPass):
+    """
+    Decompose log variants [log10, log2, log1p] operations using the identities:
+        log10(x) = log(x) / log(10)
+        log2(x)  = log(x) / log(2)
+        log1p(x) = log(1 + x)
+    """
+
+    _EDGE_OPS = {
+        exir_ops.edge.aten.log10.default,
+        exir_ops.edge.aten.log2.default,
+        exir_ops.edge.aten.log1p.default,
+    }
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._dispatcher = {
+            # Edge dialect (post-to_edge) - FP
+            exir_ops.edge.aten.log10.default: partial(self._decompose_log_n, n=10),
+            exir_ops.edge.aten.log2.default: partial(self._decompose_log_n, n=2),
+            exir_ops.edge.aten.log1p.default: partial(self._decompose_log_p, p=1),
+            # ATen dialect (pre-to_edge) - Quantized
+            torch.ops.aten.log10.default: partial(self._decompose_log_n, n=10),
+            torch.ops.aten.log2.default: partial(self._decompose_log_n, n=2),
+            torch.ops.aten.log1p.default: partial(self._decompose_log_p, p=1),
+        }
+
+    def _decompose_log_n(self, node, graph, graph_module, n):
+        input_node = node.args[0]
+        is_edge = node.target in self._EDGE_OPS
+
+        if is_edge:
+            log_op = exir_ops.edge.aten.log.default
+            div_op = exir_ops.edge.aten.div.Tensor
+            div_arg = get_const_node(
+                graph,
+                graph_module,
+                f"_log_base_{n}_constant",
+                math.log(n),
+                node,
+            )
+
+        else:
+            log_op = torch.ops.aten.log.default
+            div_op = torch.ops.aten.div.Tensor
+            div_arg = math.log(n)
+
+        with graph.inserting_after(input_node):
+            log_node = graph.create_node("call_function", log_op, (input_node,))
+            log_node.meta = copy_meta(node.meta)
+
+            with graph.inserting_after(log_node):
+                div_node = graph.create_node(
+                    "call_function", div_op, (log_node, div_arg)
+                )
+                div_node.meta = copy_meta(node.meta)
+
+        for user in node.users.copy():
+            user.replace_input_with(node, div_node)
+
+    def _decompose_log_p(self, node, graph, graph_module, p):
+        input_node = node.args[0]
+        is_edge = node.target in self._EDGE_OPS
+
+        if is_edge:
+            add_op = exir_ops.edge.aten.add.Tensor
+            log_op = exir_ops.edge.aten.log.default
+            add_arg = get_const_node(
+                graph,
+                graph_module,
+                f"_log1p_addend_{p}_constant",
+                p,
+                node,
+            )
+
+        else:
+            add_op = torch.ops.aten.add.Tensor
+            log_op = torch.ops.aten.log.default
+            add_arg = p
+
+        with graph.inserting_after(input_node):
+            add_node = graph.create_node("call_function", add_op, (input_node, add_arg))
+            add_node.meta = copy_meta(node.meta)
+
+            with graph.inserting_after(add_node):
+                log_node = graph.create_node("call_function", log_op, (add_node,))
+                log_node.meta = copy_meta(node.meta)
+
+        for user in node.users.copy():
+            user.replace_input_with(node, log_node)
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        graph = graph_module.graph
+
+        for node in list(graph.nodes):
+            if node.target in self._dispatcher:
+                self._dispatcher[node.target](node, graph, graph_module)
+
+        graph.eliminate_dead_code()
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -27,6 +27,7 @@
     DecomposeFloorDivide,
     DecomposeGlu,
     DecomposeLinalgVectorNorm,
+    DecomposeLogVariants,
     DecomposeMaxPool3d,
     DecomposeMinMaxDim,
     DecomposeReciprocal,
@@ -96,6 +97,7 @@ def get_capture_program_passes():
         (ConvertBmmToMatmul, False),
         (DecomposeAny, True),
         (DecomposeColIm, True),
+        (DecomposeLogVariants, True),
         (DecomposeMaxPool3d, True),
         (DecomposeMinMaxDim, True),
         (ExpandBroadcastTensorShape, True),
@@ -226,6 +228,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         # TODO: Skip this pass for CPU backend (Dependency: Backend-aware passes manager)
         self.add_pass(DecomposeReciprocal())
         self.add_pass(DecomposeLinalgVectorNorm(quantization_capture=True))
+        self.add_pass(DecomposeLogVariants())
         self.add_pass(ReplaceInfValues())
         self.add_pass(LiftConstantScalarOperands())
         self.add_pass(InsertReshapeForReduceOps())

backends/qualcomm/_passes/utils.py

@@ -69,6 +69,7 @@ def get_passes_dependency_for_capture_program():
         DecomposeAny,
         DecomposeColIm,
         DecomposeLinalgVectorNorm,
+        DecomposeLogVariants,
         DecomposeMaxPool3d,
         ExpandBroadcastTensorShape,
         FixedLinearKeepDim,
@@ -96,6 +97,7 @@ def get_passes_dependency_for_capture_program():
         DecomposeAny: [RemoveRedundancy],
         DecomposeColIm: [FoldQDQ],
         DecomposeLinalgVectorNorm: [RemoveRedundancy],
+        DecomposeLogVariants: [RemoveRedundancy],
         DecomposeMaxPool3d: [RemoveRedundancy],
         ExpandBroadcastTensorShape: [FoldQDQ],
         FixedLinearKeepDim: [FoldQDQ],
@@ -285,3 +287,25 @@ def append_qdq(
             dq_node = graph_module.graph.create_node("call_function", dq_op, dq_args)
             dq_node.meta = copy_meta(node.meta)
     return dq_node
+
+
+def get_const_node(
+    graph: torch.fx.Graph,
+    graph_module: torch.fx.GraphModule,
+    attr_name: str,
+    value,
+    source_node: torch.fx.Node,
+) -> torch.fx.Node:
+    """
+    Register a scalar constant as a named buffer on the graph module and return a get_attr node referencing it.
+    Used in edge dialect op decomposition passes where raw scalar arguments are not accepted by QNN op builders which need the inputs to be graph nodes.
+    """
+    dtype = source_node.meta["val"].dtype
+    tensor = torch.tensor(value, dtype=dtype)
+    graph_module.register_buffer(attr_name, tensor)
+
+    fake_mode = source_node.meta["val"].fake_mode
+    with graph.inserting_before(next(iter(graph.nodes))):
+        const_node = graph.get_attr(attr_name)
+        const_node.meta["val"] = fake_mode.from_tensor(tensor)
+    return const_node

backends/qualcomm/partition/common_defs.py

@@ -21,9 +21,6 @@
 to_be_implemented_operator = [
     exir_ops.edge.aten.adaptive_max_pool3d.default,
     exir_ops.edge.aten.div.Tensor_mode,
-    exir_ops.edge.aten.log10.default,
-    exir_ops.edge.aten.log1p.default,
-    exir_ops.edge.aten.log2.default,
     exir_ops.edge.aten.max_pool3d_with_indices.default,
     exir_ops.edge.aten.median.default,
     exir_ops.edge.aten.median.dim,

backends/qualcomm/tests/models.py

@@ -1498,6 +1498,30 @@ def forward(self, x):
         return torch.nn.functional.log_softmax(x, dim=-1)
 
 
+class Log10(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.log10(x)
+
+
+class Log1p(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.log1p(x)
+
+
+class Log2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.log2(x)
+
+
 class MaxPool2d(torch.nn.Module):
     def __init__(self, kernel_size=3, stride=1, padding=1, ceil_mode=True):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1455,6 +1455,21 @@ def test_qnn_backend_log_softmax(self):
         sample_input = (torch.randn([1, 4, 8, 8]),)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_log10(self):
+        module = Log10()  # noqa: F405
+        sample_input = (torch.abs(torch.rand(2, 5, 1, 3) + 0.1),)
+        self.lower_module_and_test_output(module, sample_input)
+
+    def test_qnn_backend_log1p(self):
+        module = Log1p()  # noqa: F405
+        sample_input = (torch.abs(torch.rand(2, 5, 1, 3) + 0.1),)
+        self.lower_module_and_test_output(module, sample_input)
+
+    def test_qnn_backend_log2(self):
+        module = Log2()  # noqa: F405
+        sample_input = (torch.abs(torch.rand(2, 5, 1, 3) + 0.1),)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_maximum(self):
         module = Maximum()  # noqa: F405
         sample_input = (torch.randn(1, 2, 3, 4), torch.randn(2, 3, 4))
@@ -3780,6 +3795,24 @@ def test_qnn_backend_log_softmax(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_log10(self):
+        module = Log10()  # noqa: F405
+        sample_input = (torch.abs(torch.rand(2, 5, 1, 3) + 0.1),)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
+
+    def test_qnn_backend_log1p(self):
+        module = Log1p()  # noqa: F405
+        sample_input = (torch.abs(torch.rand(2, 5, 1, 3) + 0.1),)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
+
+    def test_qnn_backend_log2(self):
+        module = Log2()  # noqa: F405
+        sample_input = (torch.abs(torch.rand(2, 5, 1, 3) + 0.1),)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_maximum(self):
         module = Maximum()  # noqa: F405
         sample_input = (torch.randn(1, 2, 3, 4), torch.randn(2, 3, 4))