fix (#65258)

enkilee · web-flow · commit 9b97101edae9 · 2024-06-24T18:36:28.000+08:00
diff --git a/paddle/cinn/poly/ast_gen.cc b/paddle/cinn/poly/ast_gen.cc
@@ -23,6 +23,7 @@
 #include "paddle/cinn/ir/ir_printer.h"
 #include "paddle/cinn/poly/domain_add_unit_loop_mutator.h"
 #include "paddle/cinn/poly/isl_utils.h"
+#include "paddle/common/enforce.h"
 
 namespace cinn {
 namespace poly {
@@ -371,7 +372,10 @@ void EatBlock(const isl::ast_node& node, ir::Expr* expr) {
   VLOG(2) << "get isl ast body node";
   CHECK(!node.is_null());
   CHECK(expr);
-  CHECK_EQ(isl_ast_node_get_type(node.get()), isl_ast_node_block);
+  PADDLE_ENFORCE_EQ(isl_ast_node_get_type(node.get()),
+                    isl_ast_node_block,
+                    phi::errors::InvalidArgument(
+                        "The node type should be isl_ast_node_block"));
   isl::ast_node_list list =
       isl::manage(isl_ast_node_block_get_children(node.get()));
   std::vector<ir::Expr> exprs;
@@ -387,13 +391,19 @@ void EatBlock(const isl::ast_node& node, ir::Expr* expr) {
 }
 // Eat an isl user node.
 void EatUser(const isl::ast_node& node, ir::Expr* expr) {
-  CHECK_EQ(isl_ast_node_get_type(node.get()), isl_ast_node_user);
+  PADDLE_ENFORCE_EQ(isl_ast_node_get_type(node.get()),
+                    isl_ast_node_user,
+                    phi::errors::InvalidArgument(
+                        "The node type should be isl_ast_node_user"));
   isl::ast_expr isl_expr = isl::manage(isl_ast_node_user_get_expr(node.get()));
   IslAstExprToCinnExpr(isl_expr, expr);
 }
 // Eat an isl `for` node.
 void EatFor(const isl::ast_node& node, ir::Expr* expr) {
-  CHECK_EQ(isl_ast_node_get_type(node.get()), isl_ast_node_for);
+  PADDLE_ENFORCE_EQ(
+      isl_ast_node_get_type(node.get()),
+      isl_ast_node_for,
+      phi::errors::InvalidArgument("The node type should be isl_ast_node_for"));
 
   // iter name
   isl::ast_expr iter = isl::manage(isl_ast_node_for_get_iterator(node.get()));
@@ -436,7 +446,10 @@ void EatFor(const isl::ast_node& node, ir::Expr* expr) {
 }
 
 void EatIf(const isl::ast_node& node, ir::Expr* expr) {
-  CHECK_EQ(isl_ast_node_get_type(node.get()), isl_ast_node_if);
+  PADDLE_ENFORCE_EQ(
+      isl_ast_node_get_type(node.get()),
+      isl_ast_node_if,
+      phi::errors::InvalidArgument("The node type should be isl_ast_node_if."));
   isl::ast_node then_body = isl::manage(isl_ast_node_if_get_then(node.get()));
   isl::ast_expr condition = isl::manage(isl_ast_node_if_get_cond(node.get()));
 
@@ -561,8 +574,10 @@ void IslAstExprToCinnExpr(const isl::ast_expr& node, ir::Expr* expr) {
           *expr = ir::Div::Make(ops[0], ops[1]);
           break;
         case isl_ast_op_select:
-          CHECK_EQ(ops.size(), 3UL)
-              << "In ir::Select, the ops size should be 3";
+          PADDLE_ENFORCE_EQ(ops.size(),
+                            3UL,
+                            phi::errors::InvalidArgument(
+                                "In ir::Select, the ops size should be 3"));
           ops[0]->set_type(Bool());
           *expr = ir::Select::Make(ops[0], ops[1], ops[2]);
           break;
diff --git a/paddle/cinn/poly/isl_utils.cc b/paddle/cinn/poly/isl_utils.cc
@@ -22,7 +22,7 @@
 
 #include "paddle/cinn/common/common.h"
 #include "paddle/cinn/utils/string.h"
-
+#include "paddle/common/enforce.h"
 namespace cinn {
 namespace poly {
 using utils::Join;
@@ -57,7 +57,11 @@ void isl_set_dim_names(isl::map *map,
                        isl_dim_type dim_type,
                        const std::vector<std::string> &names) {
   const int dim = isl_map_dim(map->get(), dim_type);
-  CHECK_EQ(dim, names.size());
+  PADDLE_ENFORCE_EQ(
+      dim,
+      names.size(),
+      phi::errors::InvalidArgument(
+          "The size of names should be equal to the dimension of the map."));
 
   for (int i = 0; i < dim; i++) {
     *map = isl::manage(
@@ -67,7 +71,11 @@ void isl_set_dim_names(isl::map *map,
 
 void isl_set_dim_names(isl::set *set, const std::vector<std::string> &names) {
   int dim = isl_set_dim(set->get(), isl_dim_set);
-  CHECK_EQ(dim, names.size());
+  PADDLE_ENFORCE_EQ(
+      dim,
+      names.size(),
+      phi::errors::InvalidArgument(
+          "The size of names should be equal to the dimension of the set."));
 
   for (int i = 0; i < dim; i++) {
     *set = isl::manage(
@@ -122,7 +130,11 @@ isl::set SetGetDims(isl::set set, const std::vector<int> &dims) {
   auto dim_names = isl_get_dim_names(set);
   std::vector<std::string> selected_dim_names;
   for (int v : dims) {
-    CHECK_LT(v, dim_names.size());
+    PADDLE_ENFORCE_LT(
+        v,
+        dim_names.size(),
+        phi::errors::InvalidArgument(
+            "The dim index should be less than the size of dim names."));
     selected_dim_names.push_back(dim_names[v]);
   }
 
@@ -138,7 +150,11 @@ isl::set SetGetDims(isl::set set, const std::vector<int> &dims) {
 
 isl_set *isl_get_preceding_axis(isl_set *set, int level, bool with_tuple_name) {
   int n = isl_set_dim(set, isl_dim_set);
-  CHECK_LT(level, n);
+  PADDLE_ENFORCE_LT(
+      level,
+      n,
+      phi::errors::InvalidArgument(
+          "The level should be less than the dimension of the set."));
 
   std::vector<std::string> domain_iterators;
   std::vector<std::string> range_iterators;
@@ -225,8 +241,16 @@ bool isl_is_removed_axis(isl_set __isl_keep *a, int level) {
 int isl_max_level_compatible(isl_set *a, isl_set *b) {
   int an = isl_set_dim(a, isl_dim_set);
   int bn = isl_set_dim(b, isl_dim_set);
-  CHECK_GE(an, 0);
-  CHECK_GE(bn, 0);
+  PADDLE_ENFORCE_GE(
+      an,
+      0,
+      phi::errors::InvalidArgument(
+          "The dimension of the set should be greater than or equal to 0."));
+  PADDLE_ENFORCE_GE(
+      bn,
+      0,
+      phi::errors::InvalidArgument(
+          "The dimension of the set should be greater than or equal to 0."));
 
   int compatible_level = -1;
   for (int i = 0; i < std::min(an, bn); i++) {
diff --git a/paddle/cinn/poly/naive_scheduler.cc b/paddle/cinn/poly/naive_scheduler.cc
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 #include "paddle/cinn/poly/naive_scheduler.h"
-
 #include <vector>
+#include "paddle/common/enforce.h"
 
 namespace cinn {
 namespace poly {
@@ -25,7 +25,11 @@ std::unique_ptr<Schedule> NaiveScheduler::BuildSchedule() {
 
   for (auto &group : groups_) {
     std::vector<Stage *> status;
-    CHECK_EQ(group.nodes.size(), 1UL);
+    PADDLE_ENFORCE_EQ(
+        group.nodes.size(),
+        1UL,
+        phi::errors::InvalidArgument(
+            "group.nodes.size() should be 1, but got %d", group.nodes.size()));
     NaiveGroupScheduler scheduler(
         const_cast<Stage *>(group.nodes.front()->stage));
     scheduler.Build();
diff --git a/paddle/cinn/poly/poly_scheduler.cc b/paddle/cinn/poly/poly_scheduler.cc
@@ -22,8 +22,8 @@
 #include <set>
 #include <stack>
 #include <unordered_set>
-
 #include "paddle/cinn/poly/isl_utils.h"
+#include "paddle/common/enforce.h"
 
 namespace cinn {
 namespace poly {
@@ -129,7 +129,11 @@ std::vector<Group> PartitionGraphByIterationDomain(cinn::common::Graph* graph) {
   // @{
   int num_node_in_groups = 0;
   for (auto& group : groups) num_node_in_groups += group.nodes.size();
-  CHECK_EQ(num_node_in_groups, graph->num_nodes());
+  PADDLE_ENFORCE_EQ(num_node_in_groups,
+                    graph->num_nodes(),
+                    phi::errors::InvalidArgument(
+                        "The number of nodes in groups should be the same as "
+                        "the number of nodes in the graph"));
   // @}
 
   return groups;
@@ -222,7 +226,11 @@ std::vector<Group> TopoSortGroups(std::vector<Group>& groups) {  // NOLINT
     }
     for (auto& node : top_group->nodes) {
       for (auto& edge : node->outlinks()) {
-        CHECK_EQ(edge->source()->id(), node->id());
+        PADDLE_ENFORCE_EQ(
+            edge->source()->id(),
+            node->id(),
+            phi::errors::InvalidArgument(
+                "The edge source should be the same as the node"));
         auto* sink = edge->sink();
         if (all_nodes.count(sink->id()) == 0 &&
             (--group_indegree[node2group[sink->id()]]) == 0) {
@@ -314,8 +322,11 @@ std::vector<Group> NaivePartitionGraph(cinn::common::Graph* graph) {
     graph_node_count += group.nodes.size();
   }
   // check the groups contains all the nodes in graph.
-  CHECK_EQ(graph_node_count, graph->nodes().size())
-      << "the groups should contain all the nodes in the graph";
+  PADDLE_ENFORCE_EQ(
+      graph_node_count,
+      graph->nodes().size(),
+      phi::errors::InvalidArgument(
+          "the groups should contain all the nodes in the graph"));
 #endif
 
   return group_order;
@@ -332,8 +343,10 @@ std::unique_ptr<Schedule> PolyScheduler::BuildSchedule() {
 
   // transform the DFG groups to schedule groups.
   CHECK(!schedule_graph_.nodes().empty());
-  CHECK_EQ(schedule_graph_.nodes().size(), dfg_->nodes().size())
-      << "DFG graph is not match schedule graph";
+  PADDLE_ENFORCE_EQ(
+      schedule_graph_.nodes().size(),
+      dfg_->nodes().size(),
+      phi::errors::InvalidArgument("DFG graph is not match schedule graph"));
   schedule_groups_.clear();
   for (auto& dfg_group : dfg_groups) {
     ScheduleGroup group;
@@ -345,7 +358,11 @@ std::unique_ptr<Schedule> PolyScheduler::BuildSchedule() {
     }
     schedule_groups_.emplace_back(std::move(group));
   }
-  CHECK_EQ(schedule_groups_.size(), dfg_groups.size());
+  PADDLE_ENFORCE_EQ(schedule_groups_.size(),
+                    dfg_groups.size(),
+                    phi::errors::InvalidArgument(
+                        "The number of groups should be the same as the DFG "
+                        "groups"));
 
   // Schedule each group
   ScheduleGroups();
@@ -424,10 +441,14 @@ std::vector<Shared<ScheduleGraphNode>> PolyGroupScheduler::Build() {
   }
   std::map<std::string, int> stage_level;
   for (auto& link : compute_at_links) {
-    CHECK_NE(stage_map.count(link.first), 0)
-        << link.first << " not found in stage_map";
-    CHECK_NE(stage_map.count(link.second.stage->tensor_->name), 0)
-        << link.second.stage->tensor_->name << " not found in stage_map";
+    PADDLE_ENFORCE_NE(stage_map.count(link.first),
+                      0,
+                      phi::errors::InvalidArgument(
+                          "The stage  is not found in the stage map"));
+    PADDLE_ENFORCE_NE(stage_map.count(link.second.stage->tensor_->name),
+                      0,
+                      phi::errors::InvalidArgument(
+                          "The stage  is not found in the stage map"));
     auto* a = stage_map.at(link.first);
     auto* b = stage_map.at(link.second.stage->tensor_->name);
     After(*a, *b, link.second.level);
@@ -480,7 +501,8 @@ std::vector<Shared<ScheduleGraphNode>> PolyGroupScheduler::Build() {
 
 PolyGroupScheduler::PolyGroupScheduler(const std::vector<Stage*>& stages)
     : stages_(stages) {
-  CHECK_GT(stages.size(), 0) << "No stage is provided";
+  PADDLE_ENFORCE_GT(
+      stages.size(), 0, phi::errors::InvalidArgument("No stage is provided"));
   for (auto* stage : stages) {
     AddStage(*stage);
   }
diff --git a/paddle/cinn/poly/schedule.cc b/paddle/cinn/poly/schedule.cc
@@ -25,12 +25,17 @@
 #include "paddle/cinn/poly/naive_scheduler.h"
 #include "paddle/cinn/poly/poly_scheduler.h"
 #include "paddle/cinn/utils/string.h"
-
+#include "paddle/common/enforce.h"
 namespace cinn {
 namespace poly {
 
 std::string TimeSchedule::__str__() const {
-  CHECK_LE(time_dims_.size(), kMaxDims);
+  PADDLE_ENFORCE_LE(time_dims_.size(),
+                    kMaxDims,
+                    phi::errors::InvalidArgument(
+                        "time_dims_.size() should be less than %d, but got %d",
+                        kMaxDims,
+                        time_dims_.size()));
 
   // generate range: [dup, t0, t1...]
   std::vector<std::string> range_dims, cond_dims;
@@ -75,7 +80,12 @@ std::vector<std::string> TimeSchedule::final_axis_names() const {
 
 TimeSchedule::TimeSchedule(const std::string &id,
                            const std::vector<std::string> &dims) {
-  CHECK_LE(dims.size(), kMaxDims);
+  PADDLE_ENFORCE_LE(dims.size(),
+                    kMaxDims,
+                    phi::errors::InvalidArgument(
+                        "dims.size() should be less than %d, but got %d",
+                        kMaxDims,
+                        dims.size()));
   id_ = id;
   domain_dims = dims;
   for (auto &dim : domain_dims) {
@@ -85,9 +95,20 @@ TimeSchedule::TimeSchedule(const std::string &id,
 }
 
 void TimeSchedule::OrderAfter(const TimeSchedule &other, int level) {
-  CHECK_EQ(space_size(), other.space_size()) << "space not match";
-  CHECK_LT(level, other.space_size()) << other.__str__();
-  CHECK_GE(level, 0);
+  PADDLE_ENFORCE_EQ(space_size(),
+                    other.space_size(),
+                    phi::errors::InvalidArgument("space not match"));
+  PADDLE_ENFORCE_LT(
+      level,
+      other.space_size(),
+      phi::errors::InvalidArgument("level should be less than %d, but got %d",
+                                   other.space_size(),
+                                   level));
+  PADDLE_ENFORCE_GE(
+      level,
+      0,
+      phi::errors::InvalidArgument(
+          "level should be greater than or equal to 0, but got %d", level));
   CHECK(!time_dims_.empty());
 
   root_time_ = std::max(root_time_, other.root_time_);
@@ -116,7 +137,11 @@ const std::string &TimeSchedule::id() const {
 }
 
 void TimeSchedule::ResizeTimeSpace(int size) {
-  CHECK_LE(size, kMaxDims);
+  PADDLE_ENFORCE_LE(
+      size,
+      kMaxDims,
+      phi::errors::InvalidArgument(
+          "size should be less than %d, but got %d", kMaxDims, size));
   for (int i = time_dims_.size(); i < size; i++) {
     time_dims_.emplace_back("0", 0);
   }
@@ -194,7 +219,9 @@ void SchedulerBase::AddStage(const Stage &x) {
   if (!ctx_.get()) {
     ctx_ = x.domain().ctx();
   } else {
-    CHECK_EQ(ctx_.get(), x.domain().ctx().get()) << "isl ctx not match";
+    PADDLE_ENFORCE_EQ(ctx_.get(),
+                      x.domain().ctx().get(),
+                      phi::errors::InvalidArgument("ctx not match"));
   }
 }
 
@@ -243,7 +270,11 @@ std::vector<std::string> SchedulerBase::WrapIteratorNames(
 }
 
 SchedulerBase &SchedulerBase::After(const Stage &a, const Stage &b, int level) {
-  CHECK_LT(level, space_size_);
+  PADDLE_ENFORCE_LT(
+      level,
+      space_size_,
+      phi::errors::InvalidArgument(
+          "level should be less than %d, but got %d", space_size_, level));
   auto *a_node =
       schedule_graph_.RetrieveNode(a.id())->safe_as<ScheduleGraphNode>();
   auto *b_node =
diff --git a/paddle/cinn/poly/stage.cc b/paddle/cinn/poly/stage.cc
