try fix

Author: hgt312

src/relay/transforms/simplify_expr.cc

@@ -75,57 +75,31 @@ class SimplifyReshape : public DFPatternRewrite {
   DFPattern x_;
 };
 
-/*!
- * \brief SimplifyCastLike matches the pattern of cast data to the same dtype.
- */
-class SimplifyCastLike : public DFPatternRewrite {
- public:
-  SimplifyCastLike() {
-    data_pat_ = IsWildcard();
-    like_pat_ = IsWildcard();
-    pattern_ = IsOp("cast_like")({data_pat_, like_pat_});
-  }
-
-  Expr Callback(const Expr& pre, const Expr& post,
-                const Map<DFPattern, Array<Expr>>& node_map) const override {
-    auto data = node_map[data_pat_][0];
-    const TensorTypeNode* data_ty = data->checked_type().as<TensorTypeNode>();
-    const TensorTypeNode* like_ty = pre->checked_type().as<TensorTypeNode>();
-    if (like_ty->dtype == data_ty->dtype) {
-      return data;
-    }
-    return post;
-  }
-
- protected:
-  DFPattern data_pat_;
-  DFPattern like_pat_;
-};
-
 /*!
  * \brief SimplifyCast matches the pattern of cast data to the same dtype.
  */
 class SimplifyCast : public DFPatternRewrite {
  public:
   SimplifyCast() {
     data_pat_ = IsWildcard();
-    pattern_ = IsOp("cast")({data_pat_});
+    like_pat_ = IsWildcard();
+    pattern_ = IsOp("cast_like")({data_pat_, like_pat_}) || IsOp("cast")({data_pat_});
   }
 
   Expr Callback(const Expr& pre, const Expr& post,
                 const Map<DFPattern, Array<Expr>>& node_map) const override {
     const CallNode* call = pre.as<CallNode>();
-    auto attrs = call->attrs.as<CastAttrs>();
-    auto data = node_map[data_pat_][0];
-    const TensorTypeNode* data_ty = data->checked_type().as<TensorTypeNode>();
-    if (attrs->dtype == data_ty->dtype) {
-      return data;
+    const TensorTypeNode* data_ty = call->args[0]->checked_type().as<TensorTypeNode>();
+    const TensorTypeNode* like_ty = pre->checked_type().as<TensorTypeNode>();
+    if (like_ty->dtype == data_ty->dtype) {
+      return node_map[data_pat_][0];
     }
     return post;
   }
 
  protected:
   DFPattern data_pat_;
+  DFPattern like_pat_;
 };
 
 /*!
@@ -510,7 +484,6 @@ Expr SimplifyExpr(const Expr& expr, const IRModule& mod) {
   composer.AddRewrite<EliminateIdentityRewrite>();
   composer.AddRewrite<SimplifyReshape>();
   composer.AddRewrite<SimplifyTranspose>();
-  composer.AddRewrite<SimplifyCastLike>();
   composer.AddRewrite<SimplifyCast>();
   composer.AddRewrite<FullElementwise>();
   return RewritePatterns(composer.MakeCallbacks(), expr, mod);

tests/python/relay/test_pass_simplify_expr.py

@@ -236,25 +236,18 @@ def check(x, y=None, do_nothing=False):
         check(id_op(const, x), id_op(op_like(x), x))
 
 
-def test_simplify_cast_like():
-    dtype = "int32"
-    data = relay.var("data", shape=(3, 4, 5), dtype=dtype)
-    dtype_like = relay.var("dtype_like", shape=(2, 2, 2), dtype=dtype)
-    expr = relay.cast_like(data, dtype_like)
-
-    expected = run_infer_type(data)
-    actual = run_opt_pass(expr, relay.transform.SimplifyExpr())
-    assert tvm.ir.structural_equal(actual, expected)
-
-
 def test_simplify_cast():
     dtype = "int32"
     data = relay.var("data", shape=(3, 4, 5), dtype=dtype)
-    expr = relay.cast(data, dtype)
+    expr1 = relay.cast(data, dtype)
+    dtype_like = relay.var("dtype_like", shape=(2, 2, 2), dtype=dtype)
+    expr2 = relay.cast_like(data, dtype_like)
 
     expected = run_infer_type(data)
-    actual = run_opt_pass(expr, relay.transform.SimplifyExpr())
-    assert tvm.ir.structural_equal(actual, expected)
+    actual1 = run_opt_pass(expr1, relay.transform.SimplifyExpr())
+    assert tvm.ir.structural_equal(actual1, expected)
+    actual2 = run_opt_pass(expr2, relay.transform.SimplifyExpr())
+    assert tvm.ir.structural_equal(actual2, expected)
 
 
 def test_concretize_reshape_like():