Move remove_id_ops into remove_no_ops (#873)

Summary:
Pull Request resolved: #873

### This diff
Relocates `remove_id_ops` pass in `optimize_graph.py` call to `remove_no_ops.py`.

Reviewed By: muchulee8

Differential Revision: D47937117

fbshipit-source-id: cd7e5d793d34b1acf4a8a632e3a34031d136756d

python/aitemplate/compiler/transform/__init__.py

@@ -37,7 +37,6 @@
 from aitemplate.compiler.transform.optimize_graph import optimize_graph
 from aitemplate.compiler.transform.profile import profile
 from aitemplate.compiler.transform.refine_graph import refine_graph
-from aitemplate.compiler.transform.remove_id_ops import remove_id_ops
 from aitemplate.compiler.transform.remove_no_ops import remove_no_ops
 from aitemplate.compiler.transform.remove_unused_ops import remove_unused_ops
 from aitemplate.compiler.transform.split_large_concat_ops import split_large_concat_ops

python/aitemplate/compiler/transform/optimize_graph.py

@@ -41,7 +41,6 @@
 from aitemplate.compiler.transform.remove_elementwise_no_ops import (
     remove_elementwise_no_ops,
 )
-from aitemplate.compiler.transform.remove_id_ops import remove_id_ops
 from aitemplate.compiler.transform.split_large_concat_ops import split_large_concat_ops
 from aitemplate.compiler.transform.split_large_slice_scatter_ops import (
     split_large_slice_scatter_ops,
@@ -95,7 +94,6 @@ def optimize_graph(
     """
 
     funcs = [
-        remove_id_ops,
         remove_elementwise_no_ops,
         dedup_make_jagged_ops,
         fuse_permute_bmm_and_gemm,

python/aitemplate/compiler/transform/remove_no_ops.py

@@ -40,6 +40,28 @@
 from aitemplate.utils.shape_utils import is_singleton_dimension
 
 
+def _remove_id_ops(sorted_graph: List[Tensor]) -> List[Tensor]:
+    """Remove identity ops."""
+    ops = graph_utils.get_sorted_ops(sorted_graph)
+    for op in ops:
+        if op._attrs["op"] != "identity":
+            continue
+
+        inputs = op._attrs["inputs"]
+        assert len(inputs) == 1, "identity must only have 1 input"
+
+        outputs = op._attrs["outputs"]
+        identity_output = outputs[0]
+        assert len(inputs) == 1, "identity must only have 1 output"
+
+        # skip a very special case where id takes an input and produces an output
+        if identity_output._attrs["is_output"] and inputs[0]._attrs["is_input"]:
+            continue
+
+        transform_utils.remove_single_tensor_op_from_sorted_graph(op)
+    return transform_utils.sanitize_sorted_graph(sorted_graph)
+
+
 def _remove_no_op_concats(sorted_graph: List[Tensor]) -> List[Tensor]:
     """
     Remove no-op concats from the graph. A no-op concat is where the output
@@ -328,6 +350,7 @@ def remove_no_ops(sorted_graph: List[Tensor]) -> List[Tensor]:
         Graph after remove no-ops
     """
     passes = [
+        _remove_id_ops,
         _remove_no_op_concats,
         _remove_no_op_dynamic_slices,
         _remove_no_op_splits,