[CINN / Fusion] Fix AnchorFusion code (#67038)

huangjiyi · web-flow · commit eabd6428e4b6 · 2024-08-07T10:16:42.000+08:00
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc b/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
@@ -252,7 +252,8 @@ ir::Expr CopyReduceBody(const FusibleOp& downstream, const ReduceOp& upstream) {
       return ir::ir_utils::IRCopy(op.GetFuncBody());
     }
     ir::Expr operator()(const TrivialOp& op) {
-      PADDLE_THROW("TrivialOp cannot be copied.");
+      PADDLE_THROW(
+          ::common::errors::Unimplemented("TrivialOp cannot be copied."));
     }
   };
   return std::visit(Visitor(), downstream);
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_util.cc b/paddle/cinn/hlir/framework/pir/trivial_op_util.cc
@@ -188,9 +188,10 @@ ExprSet ExprSetFinder::operator()(const ir::Expr& x) const { return f_(x); }
 ir::Expr ExprSetFinder::GetSingle(const ir::Expr& x) const {
   ExprSetFinder call = (*this) * ExprSetFinder::GetIdentity();
   const auto& o = call.operator()(x);
-  if (o.size() != 1) {
-    PADDLE_THROW("Try to get single result, but we get %d.", o.size());
-  }
+  PADDLE_ENFORCE_EQ(o.size(),
+                    1,
+                    ::common::errors::InvalidArgument(
+                        "Try to get single result, but we get %d.", o.size()));
   return *o.begin();
 }
 
@@ -544,9 +545,10 @@ ExprTransformer SubstitudeByScheduleBlockRealize(const ir::Expr& realize) {
 ExprTransformer WrapScheduleRealizer(const std::vector<ir::Var>& block_vars,
                                      const std::string& tensor_name) {
   const auto& f = [=](const ir::Expr& e) -> ir::Expr {
-    if (e.As<ir::ScheduleBlock>()) {
-      PADDLE_THROW("please input a non-schedule block expr.");
-    }
+    PADDLE_ENFORCE_EQ(e.As<ir::ScheduleBlock>(),
+                      nullptr,
+                      ::common::errors::InvalidArgument(
+                          "please input a non-schedule block expr."));
     const auto& inner_block_var = CreateInnerBlockVars(block_vars);
     const auto& replaced_e =
         ChangeVarTransformer(block_vars, inner_block_var)(e);
diff --git a/paddle/cinn/operator_fusion/fusion_tracker/expr_utils.h b/paddle/cinn/operator_fusion/fusion_tracker/expr_utils.h
@@ -43,14 +43,17 @@ struct GetSplitedExprFromFusionOp {
 struct ApplyTransform {
   explicit ApplyTransform(const ir::Expr& expr) : expr_(expr) {}
   ir::Expr operator()(const UnsupportTransformPtr& transform) {
-    PADDLE_THROW("Can not do UnsupportTransform");
+    PADDLE_THROW(
+        ::common::errors::Unavailable("Can not do UnsupportTransform"));
   }
   ir::Expr operator()(const IdentityTransformPtr& transform) { return expr_; }
   ir::Expr operator()(const AppendDimTransformPtr& transform) {
-    PADDLE_THROW("AppendDimTransform not implemented");
+    PADDLE_THROW(
+        ::common::errors::Unimplemented("AppendDimTransform not implemented"));
   }
   ir::Expr operator()(const DeleteDimTransformPtr& transform) {
-    PADDLE_THROW("DeleteDimTransform not implemented");
+    PADDLE_THROW(
+        ::common::errors::Unimplemented("DeleteDimTransform not implemented"));
   }
 
  private:
diff --git a/paddle/cinn/operator_fusion/fusion_tracker/interpreter.cc b/paddle/cinn/operator_fusion/fusion_tracker/interpreter.cc
@@ -46,7 +46,10 @@ void RunInitPatternInstr(const std::shared_ptr<InitPatternInstr>& instr,
 
 void RunTrivialInlineInstr(const std::shared_ptr<TrivialInlineInstr>& instr,
                            FusionInterpreter* interpreter) {
-  PADDLE_ENFORCE_EQ(interpreter->scope[instr->upstream_]->fusion_ops.size(), 1);
+  PADDLE_ENFORCE_EQ(interpreter->scope[instr->upstream_]->fusion_ops.size(),
+                    1,
+                    ::common::errors::InvalidArgument(
+                        "Upstream op must have only one fusion_op."));
   auto upstream_op = std::get<TrivialOp>(
       interpreter->scope[instr->upstream_]->fusion_ops.front());
   ScopeElementPtr new_pattern = std::make_shared<ScopeElement>();
@@ -69,7 +72,8 @@ void RunTmpTransformInstr(const std::shared_ptr<TmpTransformInstr>& instr,
   VLOG(4) << interpreter->scope[instr->upstream_]->fusion_ops.size();
   PADDLE_ENFORCE_EQ(interpreter->scope[instr->downstream_]->fusion_ops.size(),
                     1,
-                    "Downstream op must have only one fusion_op.");
+                    ::common::errors::InvalidArgument(
+                        "Downstream op must have only one fusion_op."));
   auto upstream_op = std::get<ReduceOp>(
       interpreter->scope[instr->upstream_]->fusion_ops.front());
   auto downstream_op =
@@ -90,7 +94,9 @@ void RunTrivialLoopAlignInstr(
     const std::shared_ptr<TrivialLoopAlignInstr>& instr,
     FusionInterpreter* interpreter) {
   PADDLE_ENFORCE_EQ(interpreter->scope[instr->downstream_]->fusion_ops.size(),
-                    1);
+                    1,
+                    ::common::errors::InvalidArgument(
+                        "Downstream op must have only one fusion_op."));
   auto upstream_op = std::get<ReduceOp>(
       interpreter->scope[instr->upstream_]->fusion_ops.front());
   auto downstream_op = std::get<TrivialOp>(
@@ -104,7 +110,10 @@ void RunTrivialLoopAlignInstr(
 
 void RunAnchorTransformInstr(const std::shared_ptr<AnchorTransformInstr>& instr,
                              FusionInterpreter* interpreter) {
-  PADDLE_ENFORCE_EQ(interpreter->scope[instr->target_]->fusion_ops.size(), 1);
+  PADDLE_ENFORCE_EQ(interpreter->scope[instr->target_]->fusion_ops.size(),
+                    1,
+                    ::common::errors::InvalidArgument(
+                        "Target op must have only one fusion_op."));
   ScopeElementPtr new_pattern = std::make_shared<ScopeElement>();
 
   std::function<ir::Expr(ir::Expr)> do_transform =
@@ -187,7 +196,8 @@ std::vector<ir::Expr> FusionInterpreter::Run() {
             dynamic_cast_instr_with_err<TrivialLoopAlignInstr>(instr), this);
         break;
       default:
-        PADDLE_THROW("Unsupported Fusion Instrution");
+        PADDLE_THROW(
+            ::common::errors::Unavailable("Unsupported Fusion Instrution"));
     }
   }
 
diff --git a/paddle/cinn/operator_fusion/fusion_tracker/tracker.h b/paddle/cinn/operator_fusion/fusion_tracker/tracker.h
@@ -103,9 +103,10 @@ struct InitPatternInstr : public FusionInstruction {
   }
   void set_idx(int idx) { idx_in_fusion_op = idx; }
   int get_idx() const {
-    if (idx_in_fusion_op == -1) {
-      PADDLE_THROW("Not initialized.");
-    }
+    PADDLE_ENFORCE_NE(
+        idx_in_fusion_op,
+        -1,
+        ::common::errors::InvalidArgument("FusionPattern is not initialized."));
     return idx_in_fusion_op;
   }
   virtual InstructionType type() const { return T_InitPattern; }
@@ -232,7 +233,9 @@ struct PaddingInstr : public FusionInstruction {
 template <typename T>
 std::shared_ptr<T> dynamic_cast_instr_with_err(FusionInstrPtr instr) {
   auto chile_instr = std::dynamic_pointer_cast<T>(instr);
-  if (!chile_instr) PADDLE_THROW("Cast Fusion Instr Failed.");
+  PADDLE_ENFORCE_NOT_NULL(
+      chile_instr,
+      ::common::errors::InvalidArgument("Cast Fusion Instr Failed."));
   return chile_instr;
 }
 
diff --git a/paddle/cinn/operator_fusion/graph_transformer/matcher.h b/paddle/cinn/operator_fusion/graph_transformer/matcher.h
@@ -164,7 +164,7 @@ struct LEOneElementWiseDownstreamMatcher {
         count += is_elementwise;
       }
     }
-    return (count <= 1);
+    return (count < node->downstream().size());
   }
 };
 
@@ -349,9 +349,11 @@ struct HorizontalCheckMiddleOutputVar {
     }
 
     // compare non_one value dims of
-    PADDLE_ENFORCE_EQ(lhs_squeeze_value_dim.size(),
-                      rhs_squeeze_value_dim.size(),
-                      "lhs squeezed dims is not equal to rhs squeezed dims");
+    PADDLE_ENFORCE_EQ(
+        lhs_squeeze_value_dim.size(),
+        rhs_squeeze_value_dim.size(),
+        ::common::errors::PreconditionNotMet(
+            "lhs squeezed dims is not equal to rhs squeezed dims"));
     for (int i = 0; i < lhs_squeeze_value_dim.size(); ++i) {
       if (get_axes_from_valuedim(lhs_squeeze_value_dim[i]) !=
           get_axes_from_valuedim(rhs_squeeze_value_dim[i]))
diff --git a/paddle/cinn/operator_fusion/pattern_fuser.h b/paddle/cinn/operator_fusion/pattern_fuser.h
@@ -188,7 +188,8 @@ static std::vector<ExprPromise> InitExprPromiseImpl(
 template <typename PATTERN>
 std::vector<ExprPromise> InitExprPromiseImpl(const PATTERN& pattern,
                                              pir::Value anchor) {
-  PADDLE_THROW("Can not Init ExprPromise");
+  PADDLE_THROW(::common::errors::Unimplemented(
+      "Can not Init ExprPromise with Unsupport Pattern."));
 }
 
 static std::vector<ExprPromise> InitExprPromise(const StmtPattern& pattern,
@@ -247,10 +248,12 @@ static std::vector<pir::Operation*> GetOutputOpsInPattern(
       return {pattern.sink_op()};
     }
     std::vector<pir::Operation*> operator()(const UnsupportPattern& pattern) {
-      PADDLE_THROW("not implement!");
+      PADDLE_THROW(::common::errors::Unimplemented(
+          "Get output ops in UnsupportPattern is not implement!"));
     }
     std::vector<pir::Operation*> operator()(const AnchorPattern& pattern) {
-      PADDLE_THROW("not implement!");
+      PADDLE_THROW(::common::errors::Unimplemented(
+          "Can't get output ops in AnchorPattern Currently."));
     }
     std::vector<pir::Operation*> operator()(const ReduceTreePattern& pattern) {
       return this->operator()(pattern.GetRootPattern());
diff --git a/paddle/cinn/operator_fusion/policy/policy_manager.h b/paddle/cinn/operator_fusion/policy/policy_manager.h
@@ -35,9 +35,10 @@ class PolicyManager {
   template <typename POLICY>
   std::shared_ptr<POLICY> GetPolicy() const {
     auto key = POLICY::Kind;
-    PADDLE_ENFORCE_NE(policies.find(key),
-                      policies.end(),
-                      phi::errors::NotFound("Policy %d Not Found", key));
+    PADDLE_ENFORCE_NE(
+        policies.find(key),
+        policies.end(),
+        phi::errors::NotFound("Can not find Policy %d, please set it.", key));
     return std::static_pointer_cast<POLICY>(policies.at(key));
   }
 
diff --git a/test/ir/pir/cinn/test_anchor_fusion.py b/test/ir/pir/cinn/test_anchor_fusion.py
@@ -75,7 +75,6 @@ def init():
             x = paddle.rand((32, 32, 128))
             return (x,)
 
-        print("feifeifei")
         self.compare_result(func, None, init)
 
     def test_anchor_fusion_2(self):

Original file line number	Diff line number	Diff line change
`@@ -252,7 +252,8 @@ ir::Expr CopyReduceBody(const FusibleOp& downstream, const ReduceOp& upstream) {`
`252`	`252`	`return ir::ir_utils::IRCopy(op.GetFuncBody());`
`253`	`253`	`}`
`254`	`254`	`ir::Expr operator()(const TrivialOp& op) {`
`255`		`- PADDLE_THROW("TrivialOp cannot be copied.");`
	`255`	`+ PADDLE_THROW(`
	`256`	`+ ::common::errors::Unimplemented("TrivialOp cannot be copied."));`
`256`	`257`	`}`
`257`	`258`	`};`
`258`	`259`	`return std::visit(Visitor(), downstream);`
Original file line number	Diff line number	Diff line change
`@@ -164,7 +164,7 @@ struct LEOneElementWiseDownstreamMatcher {`
`164`	`164`	`count += is_elementwise;`
`165`	`165`	`}`
`166`	`166`	`}`
`167`		`- return (count <= 1);`
	`167`	`+ return (count < node->downstream().size());`
`168`	`168`	`}`
`169`	`169`	`};`
`170`	`170`
`@@ -349,9 +349,11 @@ struct HorizontalCheckMiddleOutputVar {`
`349`	`349`	`}`
`350`	`350`
`351`	`351`	`// compare non_one value dims of`
`352`		`- PADDLE_ENFORCE_EQ(lhs_squeeze_value_dim.size(),`
`353`		`- rhs_squeeze_value_dim.size(),`
`354`		`- "lhs squeezed dims is not equal to rhs squeezed dims");`
	`352`	`+ PADDLE_ENFORCE_EQ(`
	`353`	`+ lhs_squeeze_value_dim.size(),`
	`354`	`+ rhs_squeeze_value_dim.size(),`
	`355`	`+ ::common::errors::PreconditionNotMet(`
	`356`	`+ "lhs squeezed dims is not equal to rhs squeezed dims"));`
`355`	`357`	`for (int i = 0; i < lhs_squeeze_value_dim.size(); ++i) {`
`356`	`358`	`if (get_axes_from_valuedim(lhs_squeeze_value_dim[i]) !=`
`357`	`359`	`get_axes_from_valuedim(rhs_squeeze_value_dim[i]))`
Original file line number	Diff line number	Diff line change
`@@ -188,7 +188,8 @@ static std::vector<ExprPromise> InitExprPromiseImpl(`
`188`	`188`	`template <typename PATTERN>`
`189`	`189`	`std::vector<ExprPromise> InitExprPromiseImpl(const PATTERN& pattern,`
`190`	`190`	`pir::Value anchor) {`
`191`		`- PADDLE_THROW("Can not Init ExprPromise");`
	`191`	`+ PADDLE_THROW(::common::errors::Unimplemented(`
	`192`	`+ "Can not Init ExprPromise with Unsupport Pattern."));`
`192`	`193`	`}`
`193`	`194`
`194`	`195`	`static std::vector<ExprPromise> InitExprPromise(const StmtPattern& pattern,`
`@@ -247,10 +248,12 @@ static std::vector<pir::Operation*> GetOutputOpsInPattern(`
`247`	`248`	`return {pattern.sink_op()};`
`248`	`249`	`}`
`249`	`250`	`std::vector<pir::Operation*> operator()(const UnsupportPattern& pattern) {`
`250`		`- PADDLE_THROW("not implement!");`
	`251`	`+ PADDLE_THROW(::common::errors::Unimplemented(`
	`252`	`+ "Get output ops in UnsupportPattern is not implement!"));`
`251`	`253`	`}`
`252`	`254`	`std::vector<pir::Operation*> operator()(const AnchorPattern& pattern) {`
`253`		`- PADDLE_THROW("not implement!");`
	`255`	`+ PADDLE_THROW(::common::errors::Unimplemented(`
	`256`	`+ "Can't get output ops in AnchorPattern Currently."));`
`254`	`257`	`}`
`255`	`258`	`std::vector<pir::Operation*> operator()(const ReduceTreePattern& pattern) {`
`256`	`259`	`return this->operator()(pattern.GetRootPattern());`