Corrects input to conv

Author: keyprocedure

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -283,31 +283,34 @@ def input_to_nhwc(
             ]
         else:
             # Need to create NHWC node
-            origin = input_node
+            source_node = input_node
+
             # TODO: safe/correct to always trace back?
-            # Trace back to source node
-            while hasattr(origin, "args") and isinstance(origin.args, tuple) and len(origin.args) > 0:
-                origin = origin.args[0]
+            # Trace back to find original source node
+            while (
+                hasattr(source_node, "args")
+                and isinstance(source_node.args, tuple)
+                and len(source_node.args) > 0
+            ):
+                source_node = source_node.args[0]
 
-            with graph_module.graph.inserting_after(origin):
+            with graph_module.graph.inserting_after(source_node):
                 input_node_nhwc = self.create_call_function_node(
                     graph_module=graph_module,
                     target=exir_ops.edge.aten._to_copy.default,
-                    args=(origin,),
+                    args=(source_node,),
                     memory_format=torch.channels_last,
                 )
 
-            # If input_node was not source
-            if origin != input_node:
-                print("Permuted\n\n")
+            # If input_node was not the original source node
+            if source_node != input_node:
+                input_node = source_node
                 # Replace downstream source node with NHWC node
-                for user in list(origin.users):
+                for user in list(input_node.users):
                     if user != input_node_nhwc:
-                        user.replace_input_with(origin, input_node_nhwc)
+                        user.replace_input_with(input_node, input_node_nhwc)
                 graph_module.recompile()
 
-            self.mark_as_nhwc_node(input_node_nhwc)
-
         self.insert_copy_and_assign_partner_nodes_quantization_sensitive(
             graph_module=graph_module,
             original_input=input_node,

backends/xnnpack/test/ops/test_conv2d.py

@@ -249,14 +249,10 @@ def _test_dq_conv2d(
 
         tester = Tester(m, inputs, dynamic_shapes=dynamic_shapes)
         tester = tester.quantize(Quantize(quantization_config=quant_config))
-        exported = tester.export()
+        tester.export()
 
         tester.check(["torch.ops.quantized_decomposed.choose_qparams"])
 
-        tester.stages["export"] = exported.stages[exported.cur]
-        print("\n----------Exported Graph:")
-        print(tester.stages["export"].graph_module.code)
-
         tester.to_edge_transform_and_lower(
             ToEdgeTransformAndLower([DynamicallyQuantizedPartitioner])
         )
@@ -265,7 +261,7 @@ def _test_dq_conv2d(
         tester.check_not(["executorch_exir_dialects_edge__ops_aten_conv2d_default"])
 
         tester.to_executorch()
-        #tester.serialize()
+        # tester.serialize()
         tester.serialize().dump_artifact("conv2d.pte")
 
         tester.run_method_and_compare_outputs(atol=atol)
@@ -751,7 +747,11 @@ def test_dq_conv2d(self) -> None:
         class SimpleConv2d(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.conv = torch.nn.Conv2d(3, 10, 3, )
+                self.conv = torch.nn.Conv2d(
+                    3,
+                    10,
+                    3,
+                )
                 self.conv.weight.requires_grad = False
                 self.conv.bias.requires_grad = False
 

backends/xnnpack/xnnpack_preprocess.py

@@ -144,7 +144,7 @@ def preprocess(
         graph_module = ep.graph_module
 
         node_to_external_map = generate_node_to_external_map(ep, graph_module)
-        print("\n----------XNNPack Preprocess Graph:", graph_module)
+
         # Make sure all inputs are contiguous_format or NCHW or default dim order
         assert_default_dim_order(graph_module)