try fix

Author: hgt312

src/relay/backend/build_module.cc

@@ -300,6 +300,7 @@ class RelayBuildModule : public runtime::ModuleNode {
       }
     });
     pass_seqs.push_back(transform::EliminateCommonSubexpr(fskip));
+    pass_seqs.push_back(transform::InferType());
     pass_seqs.push_back(transform::SimplifyExpr());
     pass_seqs.push_back(transform::CombineParallelConv2D(3));
     pass_seqs.push_back(transform::CombineParallelDense(3));

src/relay/transforms/simplify_expr.cc

@@ -76,14 +76,14 @@ class SimplifyReshape : public DFPatternRewrite {
 };
 
 /*!
- * \brief SimplifyCastLike matches the pattern of cast data to the same dtype.
+ * \brief SimplifyCast matches the pattern of cast data to the same dtype.
  */
-class SimplifyCastLike : public DFPatternRewrite {
+class SimplifyCast : public DFPatternRewrite {
  public:
-  SimplifyCastLike() {
+  SimplifyCast() {
     data_pat_ = IsWildcard();
     like_pat_ = IsWildcard();
-    pattern_ = IsOp("cast_like")({data_pat_, like_pat_});
+    pattern_ = IsOp("cast_like")({data_pat_, like_pat_}) || IsOp("cast")({data_pat_});
   }
 
   Expr Callback(const Expr& pre, const Expr& post,
@@ -102,32 +102,6 @@ class SimplifyCastLike : public DFPatternRewrite {
   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_});
-  }
-
-  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;
-    }
-    return post;
-  }
-
- protected:
-  DFPattern data_pat_;
-};
-
 /*!
  * \brief SimplifyTranspose matches the pattern of consecutive transpose op,
  *   and merges or cancels them.
@@ -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():