PaddlePaddle
diff --git a/‎paddle/cinn/hlir/op/broadcast.cc
Copy file name to clipboard
+176-90Lines changed: 176 additions & 90 deletions b/‎paddle/cinn/hlir/op/broadcast.cc
Copy file name to clipboard
+176-90Lines changed: 176 additions & 90 deletions
@@ -25,6 +25,8 @@
 #include "paddle/cinn/hlir/pe/schedule.h"
 #include "paddle/cinn/ir/layout.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
+#include "paddle/common/enforce.h"
+#include "paddle/common/errors.h"
 
 namespace cinn {
 namespace hlir {
@@ -67,33 +69,56 @@ std::shared_ptr<OpStrategy> StrategyForBroadcast(
                           const ir::Tensor &B,
                           const std::string &output_name,
                           const Expr &axis)) {
-  framework::CINNCompute binary_compute(
-      [=](lang::Args args, lang::RetValue *ret) {
-        CHECK(!args.empty()) << "The input argument of " << op_name
-                             << " compute is empty! Please check.";
-        CINNValuePack pack_args = args[0];
-        CHECK_GE(pack_args.size(), 2U)
-            << "at least 2 input tensors for " << op_name << " compute";
-        CHECK_GE(pack_args.size(), 3U) << op_name << " 's input is not enough!";
-        CHECK(pack_args[2].is_string());
-        std::string tensor_name = pack_args[2].operator std::string();
-        Expr A_expr = pack_args[0];
-        Expr B_expr = pack_args[1];
-        CHECK(A_expr.as_tensor());
-        CHECK(B_expr.as_tensor());
-        ir::Tensor A = A_expr.as_tensor_ref();
-        ir::Tensor B = B_expr.as_tensor_ref();
-        Expr axis;
-        bool trans_a;
-        for (auto &iter : attrs.attr_store) {
-          if (iter.first == "axis") {
-            axis = Expr(absl::get<int>(iter.second));
-            break;
-          }
-        }
-        auto out = pe_func(A, B, tensor_name, axis);
-        *ret = CINNValuePack{{CINNValue(Expr(out.get()))}};
-      });
+  framework::CINNCompute binary_compute([=](lang::Args args,
+                                            lang::RetValue *ret) {
+    PADDLE_ENFORCE_NE(
+        args.empty(),
+        true,
+        phi::errors::InvalidArgument(
+            "The input argument of %s compute is empty! Please check.",
+            op_name));
+    CINNValuePack pack_args = args[0];
+    PADDLE_ENFORCE_GE(pack_args.size(),
+                      2U,
+                      phi::errors::InvalidArgument(
+                          "At least 2 input tensors for %s compute, but got %d",
+                          op_name,
+                          pack_args.size()));
+    PADDLE_ENFORCE_GE(pack_args.size(),
+                      3U,
+                      phi::errors::InvalidArgument(
+                          "At least 3 input tensors for %s compute, but got %d",
+                          op_name,
+                          pack_args.size()));
+    PADDLE_ENFORCE_EQ(
+        pack_args[2].is_string(),
+        true,
+        phi::errors::InvalidArgument(
+            "Required pack_args[2] must be a string. Please check."));
+    std::string tensor_name = pack_args[2].operator std::string();
+    Expr A_expr = pack_args[0];
+    Expr B_expr = pack_args[1];
+    PADDLE_ENFORCE_NOT_NULL(
+        A_expr.as_tensor(),
+        phi::errors::InvalidArgument(
+            "Required Input must be a tensor. Please check."));
+    PADDLE_ENFORCE_NOT_NULL(
+        B_expr.as_tensor(),
+        phi::errors::InvalidArgument(
+            "Required Input must be a tensor. Please check."));
+    ir::Tensor A = A_expr.as_tensor_ref();
+    ir::Tensor B = B_expr.as_tensor_ref();
+    Expr axis;
+    bool trans_a;
+    for (auto &iter : attrs.attr_store) {
+      if (iter.first == "axis") {
+        axis = Expr(absl::get<int>(iter.second));
+        break;
+      }
+    }
+    auto out = pe_func(A, B, tensor_name, axis);
+    *ret = CINNValuePack{{CINNValue(Expr(out.get()))}};
+  });
 
   auto strategy = std::make_shared<framework::OpStrategy>();
   strategy->AddImpl(binary_compute,
@@ -113,33 +138,56 @@ std::shared_ptr<OpStrategy> StrategyForBroadcastSymbolic(
                           const ir::Tensor &B,
                           const std::string &output_name,
                           const Expr &axis)) {
-  framework::CINNCompute binary_compute(
-      [=](lang::Args args, lang::RetValue *ret) {
-        CHECK(!args.empty()) << "The input argument of " << op_name
-                             << " compute is empty! Please check.";
-        CINNValuePack pack_args = args[0];
-        CHECK_GE(pack_args.size(), 2U)
-            << "at least 2 input tensors for " << op_name << " compute";
-        CHECK_GE(pack_args.size(), 3U) << op_name << " 's input is not enough!";
-        CHECK(pack_args[2].is_string());
-        std::string tensor_name = pack_args[2].operator std::string();
-        Expr A_expr = pack_args[0];
-        Expr B_expr = pack_args[1];
-        CHECK(A_expr.as_tensor());
-        CHECK(B_expr.as_tensor());
-        ir::Tensor A = A_expr.as_tensor_ref();
-        ir::Tensor B = B_expr.as_tensor_ref();
-        Expr axis;
-        bool trans_a;
-        for (auto &iter : attrs.attr_store) {
-          if (iter.first == "axis") {
-            axis = Expr(absl::get<int>(iter.second));
-            break;
-          }
-        }
-        auto out = pe_func(A, B, tensor_name, axis);
-        *ret = CINNValuePack{{CINNValue(Expr(out.get()))}};
-      });
+  framework::CINNCompute binary_compute([=](lang::Args args,
+                                            lang::RetValue *ret) {
+    PADDLE_ENFORCE_NE(
+        args.empty(),
+        true,
+        phi::errors::InvalidArgument(
+            "The input argument of %s compute is empty! Please check.",
+            op_name));
+    CINNValuePack pack_args = args[0];
+    PADDLE_ENFORCE_GE(pack_args.size(),
+                      2U,
+                      phi::errors::InvalidArgument(
+                          "At least 2 input tensors for %s compute, but got %d",
+                          op_name,
+                          pack_args.size()));
+    PADDLE_ENFORCE_GE(pack_args.size(),
+                      3U,
+                      phi::errors::InvalidArgument(
+                          "At least 3 input tensors for %s compute, but got %d",
+                          op_name,
+                          pack_args.size()));
+    PADDLE_ENFORCE_EQ(
+        pack_args[2].is_string(),
+        true,
+        phi::errors::InvalidArgument(
+            "Required pack_args[2] must be a string. Please check."));
+    std::string tensor_name = pack_args[2].operator std::string();
+    Expr A_expr = pack_args[0];
+    Expr B_expr = pack_args[1];
+    PADDLE_ENFORCE_NOT_NULL(
+        A_expr.as_tensor(),
+        phi::errors::InvalidArgument(
+            "Required Input must be a tensor. Please check."));
+    PADDLE_ENFORCE_NOT_NULL(
+        B_expr.as_tensor(),
+        phi::errors::InvalidArgument(
+            "Required Input must be a tensor. Please check."));
+    ir::Tensor A = A_expr.as_tensor_ref();
+    ir::Tensor B = B_expr.as_tensor_ref();
+    Expr axis;
+    bool trans_a;
+    for (auto &iter : attrs.attr_store) {
+      if (iter.first == "axis") {
+        axis = Expr(absl::get<int>(iter.second));
+        break;
+      }
+    }
+    auto out = pe_func(A, B, tensor_name, axis);
+    *ret = CINNValuePack{{CINNValue(Expr(out.get()))}};
+  });
 
   auto strategy = std::make_shared<framework::OpStrategy>();
   strategy->AddImpl(
@@ -155,27 +203,52 @@ std::shared_ptr<OpStrategy> StrategyForBroadcastTo(
     const Target &target) {
   std::vector<int> out_shape;
   std::vector<int> broadcast_axes;
-  CHECK(attrs.attr_store.count("out_shape"));
+  PADDLE_ENFORCE_GE(
+      attrs.attr_store.count("out_shape"),
+      1,
+      phi::errors::InvalidArgument(
+          "The attrs.attr_store doesn't have the attribute of 'out_shape'."));
   out_shape = absl::get<std::vector<int>>(attrs.attr_store.at("out_shape"));
-  CHECK(attrs.attr_store.count("broadcast_axes"));
+  PADDLE_ENFORCE_GE(
+      attrs.attr_store.count("broadcast_axes"),
+      1,
+      phi::errors::InvalidArgument("The attrs.attr_store doesn't have the "
+                                   "attribute of 'broadcast_axes'."));
   broadcast_axes =
       absl::get<std::vector<int>>(attrs.attr_store.at("broadcast_axes"));
   VLOG(3) << "broadcast out shape: " << utils::Join(out_shape, ", ");
   VLOG(3) << "broadcast_axes shape: " << utils::Join(broadcast_axes, ", ");
 
   framework::CINNCompute broadcast_to_compute([=](lang::Args args,
                                                   lang::RetValue *ret) {
-    CHECK(!args.empty())
-        << "The input argument of broadcast_to compute is empty! Please check.";
+    PADDLE_ENFORCE_NE(
+        args.empty(),
+        true,
+        phi::errors::InvalidArgument("The input argument of broadcast_to "
+                                     "compute is empty! Please check."));
     CINNValuePack pack_args = args[0];
-    CHECK(!pack_args.empty())
-        << "The input tensors of broadcast_to compute is empty! Please check.";
-    CHECK_GE(pack_args.size(), 2U);
-    CHECK(pack_args[1].is_string());
+    PADDLE_ENFORCE_NE(
+        pack_args.empty(),
+        true,
+        phi::errors::InvalidArgument("The input tensors of broadcast_to "
+                                     "compute is empty! Please check."));
+    PADDLE_ENFORCE_GE(
+        pack_args.size(),
+        2U,
+        phi::errors::InvalidArgument(
+            "Required at least 2 input tensors, but got %d", pack_args.size()));
+    PADDLE_ENFORCE_EQ(
+        pack_args[1].is_string(),
+        true,
+        phi::errors::InvalidArgument(
+            "Required pack_args[1] must be a string. Please check."));
     std::string tensor_name = pack_args[1].operator std::string();
 
     Expr A_expr = pack_args[0];
-    CHECK(A_expr.as_tensor());
+    PADDLE_ENFORCE_NOT_NULL(
+        A_expr.as_tensor(),
+        phi::errors::InvalidArgument(
+            "Required Input must be a tensor. Please check."));
     ir::Tensor A = A_expr.as_tensor_ref();
     auto out = pe::BroadcastTo(A, out_shape, broadcast_axes, tensor_name);
     *ret = CINNValuePack{{CINNValue(out)}};
@@ -196,41 +269,54 @@ std::shared_ptr<OpStrategy> StrategyForBroadcastToSymbolic(
     const std::vector<Type> &out_type,
     const std::vector<std::vector<ir::Dim>> &output_shapes,
     const Target &target) {
-  CHECK_EQ(output_shapes.size(), 1);
+  PADDLE_ENFORCE_EQ(output_shapes.size(),
+                    1,
+                    phi::errors::InvalidArgument(
+                        "The size of output_shapes must be 1, but got %d.",
+                        output_shapes.size()));
   std::vector<ir::Expr> out_shape(output_shapes[0].size());
   std::transform(output_shapes[0].begin(),
                  output_shapes[0].end(),
                  out_shape.begin(),
                  [](const ir::Dim &dim) { return dim->dim_expr; });
   VLOG(3) << "broadcast out shape: " << utils::Join(out_shape, ", ");
 
-  framework::CINNCompute broadcast_to_compute([=](lang::Args args,
-                                                  lang::RetValue *ret) {
-    CHECK(!args.empty())
-        << "The input argument of broadcast_to compute is empty! Please check.";
-    CINNValuePack pack_args = args[0];
-    CHECK(!pack_args.empty())
-        << "The input tensors of broadcast_to compute is empty! Please check.";
-    std::string tensor_name = [&] {
-      if (pack_args.size() == 2) {
-        return pack_args[1].operator std::string();
-      } else {
-        PADDLE_ENFORCE_EQ(pack_args.size(),
-                          3,
-                          ::common::errors::InvalidArgument(
-                              "The number of input tensors is wrong. "
-                              "The expected inputs is 3, but now is %d.",
-                              pack_args.size()));
-        return pack_args[2].operator std::string();
-      }
-    }();
+  framework::CINNCompute broadcast_to_compute(
+      [=](lang::Args args, lang::RetValue *ret) {
+        PADDLE_ENFORCE_NE(
+            args.empty(),
+            true,
+            phi::errors::InvalidArgument("The input argument of broadcast_to "
+                                         "compute is empty! Please check."));
+        CINNValuePack pack_args = args[0];
+        PADDLE_ENFORCE_NE(
+            pack_args.empty(),
+            true,
+            phi::errors::InvalidArgument("The input tensors of broadcast_to "
+                                         "compute is empty! Please check."));
+        std::string tensor_name = [&] {
+          if (pack_args.size() == 2) {
+            return pack_args[1].operator std::string();
+          } else {
+            PADDLE_ENFORCE_EQ(pack_args.size(),
+                              3,
+                              ::common::errors::InvalidArgument(
+                                  "The number of input tensors is wrong. "
+                                  "The expected inputs is 3, but now is %d.",
+                                  pack_args.size()));
+            return pack_args[2].operator std::string();
+          }
+        }();
 
-    Expr A_expr = pack_args[0];
-    CHECK(A_expr.as_tensor());
-    ir::Tensor A = A_expr.as_tensor_ref();
-    auto out = pe::BroadcastTo(A, out_shape, tensor_name);
-    *ret = CINNValuePack{{CINNValue(out)}};
-  });
+        Expr A_expr = pack_args[0];
+        PADDLE_ENFORCE_NOT_NULL(
+            A_expr.as_tensor(),
+            phi::errors::InvalidArgument(
+                "Required Input must be a tensor. Please check."));
+        ir::Tensor A = A_expr.as_tensor_ref();
+        auto out = pe::BroadcastTo(A, out_shape, tensor_name);
+        *ret = CINNValuePack{{CINNValue(out)}};
+      });
 
   auto strategy = std::make_shared<framework::OpStrategy>();
   strategy->AddImpl(