Refactor permute code

Author: keyprocedure

backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py

@@ -282,12 +282,13 @@ def input_to_nhwc(
                 ChannelsLastTaggedReshapePass.PARTNER_NODE
             ]
         else:
-            # trace back to permute
+            # Need to create NHWC node
             origin = 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]
 
-            # at x choose_qparams and quantize insert permute
             with graph_module.graph.inserting_after(origin):
                 input_node_nhwc = self.create_call_function_node(
                     graph_module=graph_module,
@@ -296,24 +297,17 @@ def input_to_nhwc(
                     memory_format=torch.channels_last,
                 )
 
-            for user in list(origin.users):
-                if user != input_node_nhwc:
-                    user.replace_input_with(origin, input_node_nhwc)
+            # If input_node was not source
+            if origin != input_node:
+                print("Permuted\n\n")
+                # Replace downstream source node with NHWC node
+                for user in list(origin.users):
+                    if user != input_node_nhwc:
+                        user.replace_input_with(origin, input_node_nhwc)
+                graph_module.recompile()
 
-            graph_module.recompile()
             self.mark_as_nhwc_node(input_node_nhwc)
 
-        # TODO: uncomment, use case when permute not needed
-        #     # Need to create NHWC node                     ----------------------------- CONVERSION HAPPENING ----->>
-        #     with graph_module.graph.inserting_after(input_node):
-        #         input_node_nhwc = self.create_call_function_node(
-        #             graph_module=graph_module,
-        #             target=exir_ops.edge.aten._to_copy.default,
-        #             args=(input_node,),
-        #             memory_format=torch.channels_last,
-        #         )
-        #     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,15 +249,14 @@ def _test_dq_conv2d(
 
         tester = Tester(m, inputs, dynamic_shapes=dynamic_shapes)
         tester = tester.quantize(Quantize(quantization_config=quant_config))
-
-        tester.stages["quantize"] = tester.stages[tester.cur]
-
         exported = tester.export()
 
-        tester.stages["export"] = exported.stages[exported.cur]
-
         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])
         )
@@ -266,7 +265,6 @@ def _test_dq_conv2d(
         tester.check_not(["executorch_exir_dialects_edge__ops_aten_conv2d_default"])
 
         tester.to_executorch()
-
         #tester.serialize()
         tester.serialize().dump_artifact("conv2d.pte")
 

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)