Qualcomm AI Engine Direct - Add composite lowering for aten.convolution_backward

backends/qualcomm/_passes/i64_to_i32.py

@@ -73,7 +73,10 @@ def _record_original_output_dtype(self, graph_module: torch.fx.GraphModule):
             # Keep track of original output dtype so we ensure the dtype of the graph is consistent with nn.Module
             if is_graph_output(n):
                 if isinstance(n.meta["val"], (tuple, list)):
-                    dtype_list = [tensor.dtype for tensor in n.meta["val"]]
+                    dtype_list = [
+                        tensor.dtype if tensor is not None else None
+                        for tensor in n.meta["val"]
+                    ]
                     n.meta[QCOM_ORIG_DTYPE] = dtype_list
                 else:
                     n.meta[QCOM_ORIG_DTYPE] = n.meta["val"].dtype
@@ -84,6 +87,8 @@ def _preserve_output_dtype(self, graph_module: torch.fx.GraphModule):
                 if isinstance(n.meta["val"], (tuple, list)):
                     for i, dtype in enumerate(n.meta[QCOM_ORIG_DTYPE]):
                         # TODO: Enable this in future to support OP such as topK
+                        if n.meta["val"][i] is None or dtype is None:
+                            continue
                         if n.meta["val"][i].dtype != dtype:
                             raise AssertionError(
                                 "Multi output nodes currently don't support casting dtype back."
@@ -115,10 +120,10 @@ def _preserve_output_dtype(self, graph_module: torch.fx.GraphModule):
     def _update_meta(self, node: torch.fx.node) -> None:
         meta_val = node.meta["val"]
         if isinstance(meta_val, tuple):
-            node.meta["val"] = (
+            node.meta["val"] = tuple(
                 (
                     fake_tensor.to(torch.int32)
-                    if fake_tensor.dtype == torch.int64
+                    if fake_tensor is not None and fake_tensor.dtype == torch.int64
                     else fake_tensor
                 )
                 for fake_tensor in meta_val

backends/qualcomm/_passes/layout_transform.py

@@ -47,6 +47,7 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.avg_pool2d.default,
         exir_ops.edge.aten.avg_pool3d.default,
         exir_ops.edge.aten.convolution.default,
+        exir_ops.edge.aten.convolution_backward.default,
         exir_ops.edge.aten.grid_sampler_2d.default,
         exir_ops.edge.aten.grid_sampler_3d.default,
         exir_ops.edge.aten.instance_norm.default,
@@ -169,8 +170,8 @@ def mark_as_transformed(self, node: torch.fx.Node) -> None:
             getitem_node = list(node.users.keys())[0]
             if getitem_node.target.__name__ != "getitem":
                 raise AssertionError(
-                    "Expected node's user to be getitem, "
-                    f"got {getitem_node.target.__name__}"
+                    f"Expected tuple-output node {node.name} ({node.target}) to have a "
+                    f"getitem user, got {getitem_node.target.__name__}"
                 )
             index = getitem_node.args[1]
             node.meta[self.transformed_tag] = self.get_axis_order(
@@ -226,32 +227,45 @@ def is_edge_condition(self, node):
 
         return False
 
-    def insert_node(self, graph_module, node, revert_layout: bool) -> None:
+    def insert_node(
+        self, graph_module, node, revert_layout: bool, layout_node=None
+    ) -> None:
         if not self.insert_permute:
             return
-        with graph_module.graph.inserting_after(node):
-            users = node.users.copy()
-            if isinstance(node.meta["val"], tuple):
-                getitem_node = list(node.users.keys())[0]
+        if layout_node is None:
+            layout_node = node
+        if isinstance(node.meta["val"], tuple):
+            for getitem_node in node.users:
                 if getitem_node.target.__name__ != "getitem":
                     raise AssertionError(
-                        f"Expected bn node's user to be getitem, got {getitem_node.target.__name__}"
+                        f"Expected tuple-output node {node.name} ({node.target}) to have a "
+                        f"getitem user, got {getitem_node.target.__name__}"
                     )
-                index = getitem_node.args[1]
-                tensor = node.meta["val"][index]
-            else:
-                tensor = node.meta["val"]
+                tensor = node.meta["val"][getitem_node.args[1]]
+                if tensor is not None:
+                    self.insert_node(
+                        graph_module,
+                        getitem_node,
+                        revert_layout,
+                        layout_node,
+                    )
+            return
+
+        tensor = node.meta["val"]
+        permute_input = node
+        users = node.users.copy()
 
+        with graph_module.graph.inserting_after(permute_input):
             permute = self.create_call_function_node(
                 graph_module,
                 exir_ops.edge.aten.permute_copy.default,
                 (
-                    node,
+                    permute_input,
                     self.get_axis_order(eval_shape(tensor.shape), revert_layout),
                 ),
             )
             permute.meta["val"] = tensor
-            permute.meta[QCOM_QUANT_ATTRS] = node.meta.get(QCOM_QUANT_ATTRS)
+            permute.meta[QCOM_QUANT_ATTRS] = layout_node.meta.get(QCOM_QUANT_ATTRS)
             # we need this to check the annotation boundary
             permute.meta[QCOM_INSERTED_PERMUTE] = True
 
@@ -264,14 +278,14 @@ def insert_node(self, graph_module, node, revert_layout: bool) -> None:
             #               └--------------------------------------> qnn_premute -┙
             # i.e. insert permute by condition between user and current node
             #      if there are multiple users included
-            is_node_transformed = self.is_transformed_node(node)
+            is_node_transformed = self.is_transformed_node(layout_node)
             for user in users:
                 is_user_transformed = (
                     self.is_transformed_node(user) or QCOM_LAYOUT_CHANGE in user.meta
                 )
                 # insert permute only in exclusive condition
                 if is_node_transformed != is_user_transformed:
-                    user.replace_input_with(node, permute)
+                    user.replace_input_with(permute_input, permute)
 
     def create_call_function_node(
         self,
@@ -286,7 +300,12 @@ def create_call_function_node(
         )
 
     def traverse(self, node: torch.fx.Node, graph_module: torch.fx.GraphModule) -> None:
-        for arg in node.args:
+        for index, arg in enumerate(node.args):
+            if (
+                node.target == exir_ops.edge.aten.convolution_backward.default
+                and index == 1
+            ):
+                continue
             if isinstance(arg, list):
                 for arg_node in arg:
                     self.annotate_layout(arg_node, graph_module, revert_layout=False)

backends/qualcomm/builders/README.md

@@ -13,6 +13,7 @@ Thank you for contributing to Qualcomm AI Engine Direct delegate for ExecuTorch.
     - [Implementation](#implementation)
     - [Quantizer Annotation](#quantizer-annotation)
   - [Operator Support Status](#operator-support-status)
+  - [Additional PyTorch Operators Supported via Composite QNN Lowering](#additional-pytorch-operators-supported-via-composite-qnn-lowering)
   - [Issues](#issues)
   - [Pull Requests](#pull-requests)
 
@@ -540,6 +541,19 @@ The following PyTorch operators are supported through decomposition or annotatio
 | `aten.trunc` | `DecomposeTrunc` |
 | `aten.var.correction`, `aten.var.dim` | `DecomposeVar` |
 
+## Additional PyTorch Operators Supported via Composite QNN Lowering
+
+The following PyTorch operators are lowered by a node visitor into multiple QNN operations.
+
+ConvBackward uses this path to reuse the forward Conv2d's QNN weight wrapper and derive HWOI at runtime; an ATen decomposition would otherwise require another static weight or changes to the forward Conv2d visitor. It currently supports floating-point inputs and weights only; quantized ConvBackward support is planned for a future change.
+
+| PyTorch Op | Builder | QNN Lowering |
+|---|---|---|
+| `aten.convolution_backward.default` | `ConvBackward` | `Transpose` → `TransposeConv2d` |
+| `aten.copy.default` | `Copy` | `Reshape` → `Tile` |
+| `aten.embedding.default` | `Embedding` | `Convert` → `Gather` or `Gather` → `Convert` |
+| `aten.index_put.default` | `IndexPutVisitor` | tensor preparation → `ScatterNd` |
+
 ## Issues
 Please refer to the [issue section](../README.md#issues) for more information.
 

backends/qualcomm/builders/__init__.py

@@ -29,6 +29,7 @@
     op_channel_shuffle,
     op_clamp,
     op_conv,
+    op_conv_backward,
     op_copy,
     op_cos,
     op_cum_sum,
@@ -148,6 +149,7 @@
     op_channel_shuffle,
     op_clamp,
     op_conv,
+    op_conv_backward,
     op_copy,
     op_cos,
     op_cum_sum,