Eliminate nop cast (onnx#10)

* Add pass to eliminate no-op casts

Signed-off-by: crisp-snakey <[email protected]>

* Add basic test for no-op cast elimination

Signed-off-by: crisp-snakey <[email protected]>

* Formatting

Signed-off-by: crisp-snakey <[email protected]>

* Propagate sizes

Signed-off-by: crisp-snakey <[email protected]>

* Propagate output name

Signed-off-by: crisp-snakey <[email protected]>

Author: crisp-snakey

onnxoptimizer/pass_registry.h

@@ -10,6 +10,7 @@
 
 #include "onnxoptimizer/passes/eliminate_deadend.h"
 #include "onnxoptimizer/passes/eliminate_identity.h"
+#include "onnxoptimizer/passes/eliminate_nop_cast.h"
 #include "onnxoptimizer/passes/eliminate_nop_dropout.h"
 #include "onnxoptimizer/passes/eliminate_nop_monotone_argmax.h"
 #include "onnxoptimizer/passes/eliminate_nop_pad.h"
@@ -44,6 +45,7 @@ struct GlobalPassRegistry {
     // Register the optimization passes to the optimizer.
     registerPass<NopEmptyPass>();
     registerPass<EliminateDeadEnd>();
+    registerPass<EliminateNopCast>();
     registerPass<EliminateNopDropout>();
     registerPass<EliminateIdentity>();
     registerPass<EliminateNopMonotoneArgmax>();

onnxoptimizer/passes/eliminate_nop_cast.h

@@ -0,0 +1,43 @@
+// ATTENTION: The code in this file is highly EXPERIMENTAL.
+// Adventurous users should note that the APIs will probably change.
+
+#pragma once
+
+#include "onnxoptimizer/pass.h"
+
+namespace ONNX_NAMESPACE {
+namespace optimization {
+
+struct EliminateNopCast final : public PredicateBasedPass {
+  explicit EliminateNopCast()
+      : PredicateBasedPass(
+            PassType::Nop,
+            PassEfficiency::Complete,
+            PassOptimizationType::Compute) {}
+
+  std::string getPassName() const override {
+    return "eliminate_nop_cast";
+  }
+
+  bool patternMatchPredicate(Node* node) override {
+    return (node->kind() == kCast && node->hasAttribute(kto) &&
+      node->input()->elemType() == node->i(kto));
+  }
+
+  bool runTransform(Node* node, Graph& graph, NodeDestroyType& destroy_current)
+      override {
+    if (node->output()->has_sizes()) {
+        node->input()->setSizes(node->output()->sizes());
+    }
+    if (std::find(graph.outputs().rbegin(), graph.outputs().rend(),
+                  node->output()) != graph.outputs().rend()) {
+      node->input()->setUniqueName(node->output()->uniqueName());
+    }
+    node->output()->replaceAllUsesWith(node->input());
+    destroy_current = NodeDestroyType::DestroyOne;
+    return true;
+  }
+};
+
+} // namespace optimization
+} // namespace ONNX_NAMESPACE

onnxoptimizer/test/optimizer_test.py

@@ -161,6 +161,18 @@ def test_eliminate_identity_multiple_uses(self):  # type: () -> None
             assert node.op_type != "Identity"
         assert len(optimized_model.graph.node) == 2
 
+    def test_nop_cast(self):
+        cast = helper.make_node("Cast", ["A"], ["B"], to=TensorProto.FLOAT)
+        graph = helper.make_graph(
+            [cast],
+            "test",
+            [helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3))],
+            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (2, 3))])
+
+        optimized_model = self._optimized(graph, ["eliminate_nop_cast"])
+
+        assert len(optimized_model.graph.node) == 0
+
     def test_nop_transpose_graph_output(self):  # type: () -> None
         add = helper.make_node("Add", ["X", "Y"], ["A"])
         trans = helper.make_node("Transpose", ["A"], ["B"], perm=[0, 1])