Arm backend: Remove duplicated code when inserting int32 cast node

- Extract the logic of inserting int32 cast node into a helper
  function in pass InsertCastForOpsWithInt64InputPass

Change-Id: I748d53921981d43c3a913c4e7a15f3daf33c836e
Signed-off-by: Yufeng Shi <[email protected]>

Author: YufengShi-dudu

backends/arm/_passes/insert_int64_input_cast_pass.py

@@ -55,6 +55,20 @@ def _check_aten_embedding_within_int32(self, weights, indices, node: torch.fx.No
 
         return True
 
+    def _insert_int32_cast_before_node(self, graph, node, original_input):
+        to_copy_op = self.get_decomposition(node.target)
+        with graph.inserting_before(node):
+            cast_before = create_node(
+                graph,
+                to_copy_op,
+                args=(original_input,),
+                kwargs={
+                    "dtype": torch.int32,
+                    "memory_format": torch.preserve_format,
+                },
+            )
+            node.replace_input_with(original_input, cast_before)
+
     def call(self, graph_module):
         graph = graph_module.graph
         modified_graph = False
@@ -73,24 +87,8 @@ def call(self, graph_module):
             ):
                 weights = args[0]
                 indices = args[1]
-                valid_for_insert = self._check_aten_embedding_within_int32(
-                    weights, indices, node
-                )
-
-                if valid_for_insert:
-                    to_copy_op = self.get_decomposition(node.target)
-                    with graph.inserting_before(node):
-                        cast_before = create_node(
-                            graph,
-                            to_copy_op,
-                            args=(indices,),
-                            kwargs={
-                                "dtype": torch.int32,
-                                "memory_format": torch.preserve_format,
-                            },
-                        )
-                        node.replace_input_with(indices, cast_before)
-
+                if self._check_aten_embedding_within_int32(weights, indices, node):
+                    self._insert_int32_cast_before_node(graph, node, indices)
                     modified_graph = True
 
             elif node.target in (
@@ -103,19 +101,7 @@ def call(self, graph_module):
                 if fake_tensor.dtype != torch.int64:
                     continue
 
-                to_copy_op = self.get_decomposition(node.target)
-                with graph.inserting_before(node):
-                    cast_before = create_node(
-                        graph,
-                        to_copy_op,
-                        args=(input_tensor,),
-                        kwargs={
-                            "dtype": torch.int32,
-                            "memory_format": torch.preserve_format,
-                        },
-                    )
-                    node.replace_input_with(input_tensor, cast_before)
-
+                self._insert_int32_cast_before_node(graph, node, input_tensor)
                 modified_graph = True
 
         if modified_graph: