[PIR] Support AssertOp for pir (#60044)

MarioLulab · web-flow · commit 2d38689f1917 · 2023-12-20T16:18:39.000+08:00
* delete assert from ops.yaml

* add pd_op.assert api and pir op

* update executor

* fix compile error

* fix assert bug

* update testcases

* add OpYamlInfoInterface for transcribe
diff --git a/paddle/fluid/framework/new_executor/instruction/CMakeLists.txt b/paddle/fluid/framework/new_executor/instruction/CMakeLists.txt
@@ -4,6 +4,7 @@ cc_library(
        phi_kernel_instruction.cc
        legacy_kernel_instruction.cc
        if_instruction.cc
+       assert_instruction.cc
        while_instruction.cc
        select_input_instruction.cc
        has_elements_instruction.cc
diff --git a/paddle/fluid/framework/new_executor/instruction/assert_instruction.cc b/paddle/fluid/framework/new_executor/instruction/assert_instruction.cc
@@ -0,0 +1,93 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/new_executor/instruction/assert_instruction.h"
+#include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
+#include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
+#include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
+#include "paddle/phi/kernels/funcs/tensor_formatter.h"
+
+namespace paddle {
+namespace framework {
+AssertInstruction::AssertInstruction(size_t id,
+                                     const platform::Place& place,
+                                     ::pir::Operation* op,
+                                     ValueExecutionInfo* value_exe_info)
+    : InstructionBase(id, place), op_(op), value_exe_info_(value_exe_info) {
+  PADDLE_ENFORCE(op->isa<paddle::dialect::AssertOp>(),
+                 phi::errors::PreconditionNotMet(
+                     "Assert instruction only support assert op"));
+
+  auto assert_op = op->dyn_cast<paddle::dialect::AssertOp>();
+  VLOG(6) << "construct assert instruction for: " << assert_op->name();
+
+  std::unordered_map<pir::Value, std::vector<int>> inputs;
+  inputs.emplace(assert_op.cond(),
+                 GetValueIds(assert_op.cond(), *value_exe_info_));
+  inputs.emplace(assert_op.data(),
+                 GetValueIds(assert_op.data(), *value_exe_info_));
+  SetInputs(inputs);
+
+  op_ = op;
+  SetKernelType(AnalyseOpFuncType(op, place));
+  VLOG(6) << "finish process analyse kernel type";
+
+  auto cond_value = assert_op.operand_source(0);
+  cond_var_ = value_exe_info_->GetVarByValue(cond_value);
+  auto data_value = assert_op.operand_source(1);
+  data_var_ = value_exe_info_->GetVarByValue(data_value);
+  VLOG(6) << "finish process cond_var and data_var";
+}
+
+void AssertInstruction::Run() {
+  DeviceContext().Wait();
+  const phi::DenseTensor& cond = cond_var_->Get<phi::DenseTensor>();
+
+  PADDLE_ENFORCE_EQ(
+      cond.numel(),
+      1,
+      platform::errors::InvalidArgument(
+          "The numel of Input(Condition) of AssertOp must be 1. But now "
+          "the Condition's shape is %s.",
+          cond.dims().to_str()));
+
+  bool cond_data = GetCondData(cond);
+  if (cond_data) {
+    return;
+  }
+
+  funcs::TensorFormatter formatter;
+  formatter.SetSummarize(
+      op_->attribute<::pir::Int64Attribute>("summarize").data());
+
+  const std::vector<pir::Value>& inputs_data_val =
+      op_->dyn_cast<paddle::dialect::AssertOp>()
+          .data()
+          .defining_op<::pir::CombineOp>()
+          .inputs();
+  for (pir::Value val : inputs_data_val) {
+    const std::string& name = value_exe_info_->GetVarName(val);
+    const phi::DenseTensor& tensor =
+        value_exe_info_->GetVarByValue(val)->Get<phi::DenseTensor>();
+    formatter.Print(tensor, name);
+  }
+
+  PADDLE_THROW(platform::errors::InvalidArgument(
+      "The condition variable '%s' of AssertOp must be "
+      "true, but received false",
+      value_exe_info_->GetVarName(cond_var_)));
+}
+
+}  // namespace framework
+}  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/instruction/assert_instruction.h b/paddle/fluid/framework/new_executor/instruction/assert_instruction.h
@@ -0,0 +1,52 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <string>
+#include "paddle/fluid/framework/new_executor/instruction/instruction_base.h"
+#include "paddle/fluid/framework/new_executor/new_executor_defs.h"
+
+namespace paddle {
+namespace framework {
+class ValueExecutionInfo;
+
+class AssertInstruction : public InstructionBase {
+ public:
+  AssertInstruction(size_t id,
+                    const platform::Place& place,
+                    ::pir::Operation* op,
+                    ValueExecutionInfo* value_exe_info);
+
+  void Run() override;
+
+  const std::string& Name() const override { return name_; }
+
+  ::pir::Operation* Operation() const override { return op_; }
+
+ private:
+  ::pir::Operation* op_;
+
+  OpFuncType type_;
+
+  std::string name_{"assert_instruction"};
+
+  ValueExecutionInfo* value_exe_info_;  // not owned
+
+  Variable* cond_var_;
+  Variable* data_var_;
+};
+
+}  // namespace framework
+}  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -45,6 +45,7 @@
 #include "paddle/fluid/framework/new_executor/instruction/cinn_jit_instruction.h"
 #endif
 
+#include "paddle/fluid/framework/new_executor/instruction/assert_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/builtin_combine_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/has_elements_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/if_instruction.h"
@@ -703,6 +704,8 @@ void PirInterpreter::BuildInstruction() {
                  ->BodyInterpreter()});
       } else if (op.isa<paddle::dialect::HasElementsOp>()) {
         CREATE_INSTR(HasElementsInstruction);
+      } else if (op.isa<paddle::dialect::AssertOp>()) {
+        CREATE_INSTR(AssertInstruction);
       } else if (op.isa<paddle::dialect::SelectInputOp>()) {
         CREATE_INSTR(SelectInputInstruction);
       } else {
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -13,7 +13,8 @@
 // limitations under the License.
 #ifdef GET_OP_LIST
 #undef GET_OP_LIST
-paddle::dialect::IfOp, paddle::dialect::WhileOp, paddle::dialect::HasElementsOp
+paddle::dialect::IfOp, paddle::dialect::WhileOp, paddle::dialect::HasElementsOp,
+    paddle::dialect::AssertOp
 #else
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 
@@ -465,11 +466,111 @@ void HasElementsOp::VerifySig() {
              "The type of cf.has_elements' output is not correct.");
 }
 
+const char *AssertOp::attributes_name[1] = {"summarize"};
+
+void AssertOp::Build(pir::Builder &builder,             // NOLINT
+                     pir::OperationArgument &argument,  // NOLINT
+                     pir::Value cond_,
+                     pir::Value data_,
+                     int64_t summarize) {
+  VLOG(4) << "Start build AssertOp";
+
+  VLOG(4) << "Builder construction inputs";
+  std::vector<pir::Value> argument_inputs = {cond_, data_};
+  argument.AddInputs(argument_inputs);
+
+  VLOG(4) << "Builder construction attributes";
+  pir::Attribute attr_summarize =
+      pir::Int64Attribute::get(pir::IrContext::Instance(), summarize);
+  argument.AddAttribute("summarize", attr_summarize);
+}
+
+OpInfoTuple AssertOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      paddle::dialect::OpInputInfo("cond",
+                                   "paddle::dialect::DenseTensorType",
+                                   false,
+                                   false,
+                                   false,
+                                   false),
+      paddle::dialect::OpInputInfo(
+          "data",
+          "pir::VectorType<paddle::dialect::DenseTensorType>",
+          false,
+          false,
+          false,
+          false)};
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {
+      paddle::dialect::OpAttributeInfo("summarize", "pir::Int64Attribute", "")};
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {};
+  paddle::dialect::OpRunTimeInfo run_time_info = paddle::dialect::OpRunTimeInfo(
+      "", {""}, "assert", {"cond", "data", "summarize"}, {"cond"}, {}, {}, {});
+  return std::make_tuple(inputs, attributes, outputs, run_time_info, "assert");
+}
+
+void AssertOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: AssertOp.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    IR_ENFORCE(input_size == 2u,
+               "The size %d of inputs must be equal to 2.",
+               input_size);
+
+    if ((*this)->operand_source(0).type().isa<pir::DenseTensorType>()) {
+      IR_ENFORCE((*this)
+                     ->operand_source(0)
+                     .type()
+                     .dyn_cast<pir::DenseTensorType>()
+                     .dtype()
+                     .isa<pir::BoolType>(),
+                 "Type validation failed for the 0th input, it should be a "
+                 "bool DenseTensorType.");
+    }
+
+    if (auto vec_type =
+            (*this)->operand(1).type().dyn_cast<pir::VectorType>()) {
+      for (size_t i = 0; i < vec_type.size(); ++i) {
+        IR_ENFORCE(vec_type[i].isa<paddle::dialect::DenseTensorType>() ||
+                       vec_type[i].isa<paddle::dialect::SelectedRowsType>() ||
+                       vec_type[i].isa<AllocatedDenseTensorType>(),
+                   "Type validation failed for the 1th input.");
+      }
+    } else {
+      IR_ENFORCE(
+          (*this)->operand(1).type().isa<paddle::dialect::DenseTensorType>() ||
+              (*this)
+                  ->operand(1)
+                  .type()
+                  .isa<paddle::dialect::SelectedRowsType>(),
+          (*this)->operand(1).type().isa<AllocatedDenseTensorType>(),
+          "Type validation failed for the 1th input.");
+    }
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto &attributes = this->attributes();
+    IR_ENFORCE(attributes.count("summarize") > 0, "summarize does not exist.");
+    IR_ENFORCE(attributes.at("summarize").isa<pir::Int64Attribute>(),
+               "Type of attribute: summarize is not pir::Int64Attribute.");
+  }
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    IR_ENFORCE(output_size == 0u,
+               "The size %d of outputs must be equal to 0.",
+               output_size);
+    // Outputs num is 0, not need to check outputs type.
+  }
+  VLOG(4) << "End Verifying for: AssertOp.";
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::IfOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::WhileOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::HasElementsOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AssertOp)
 
 #endif
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
@@ -15,6 +15,7 @@
 #pragma once
 #include <vector>
 
+#include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
 #include "paddle/fluid/pir/dialect/operator/interface/vjp.h"
 #include "paddle/pir/core/block.h"
 #include "paddle/pir/core/op_base.h"
@@ -128,9 +129,41 @@ class HasElementsOp : public pir::Op<HasElementsOp> {
   pir::Value out() { return result(0); }
 };
 
+///
+/// \brief The AssertOp is an operation that asserts the given condition is
+/// true. If the condition is false, prints the tensors in data. ``summarize``
+/// specifies the number of the elements in the tensors to print.
+
+/// It takes two inputs: cond and data, and takes one attribute: summarize. The
+/// semantics of AssertOp[assert_op(cond, data, summarize)] are as below:
+///   if(!cond){
+///      print(summarize number of elements in data)
+///   }
+///
+class AssertOp : public pir::Op<AssertOp, OpYamlInfoInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd_op.assert"; }
+  static constexpr uint32_t attributes_num = 1;
+  static const char *attributes_name[1];
+
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value cond_,
+                    pir::Value data_,
+                    int64_t summarize);
+
+  static OpInfoTuple GetOpInfo();
+  void VerifySig();
+
+  pir::Value cond() { return operand_source(0); }
+  pir::Value data() { return operand_source(1); }
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::IfOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::WhileOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::HasElementsOp);
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AssertOp);
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -85,14 +85,6 @@
   kernel :
     func : any
 
-- op : assert
-  args : (Tensor cond, Tensor[] data, int64_t summarize = -1)
-  output :
-  kernel :
-    func : assert
-    param : [cond, data, summarize]
-    data_type : cond
-
 - op : assign
   args : (Tensor x)
   output : Tensor
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -105,6 +105,7 @@ const std::unordered_set<std::string> SpecialLowerOps = {
     pir::TuplePushOp::name(),
     pir::TuplePopOp::name(),
     HasElementsOp::name(),
+    AssertOp::name(),
     SelectInputOp::name(),
     "cinn_runtime.jit_kernel"};
 
@@ -1214,6 +1215,15 @@ void HandleForSpecialOp(
     }
   }
 
+  if (op_item->isa<AssertOp>()) {
+    for (size_t i = 0; i < op_item->num_operands(); ++i) {
+      auto cur_in = op_item->operand_source(i);
+      auto new_in = GetNewInput(
+          cur_in, *map_value_pair, static_cast<int>(i), op_item->name());
+      vec_inputs.push_back(new_in);
+    }
+  }
+
   if (op_item->isa<::pir::TuplePopOp>()) {
     for (size_t i = 0; i < op_item->num_operands(); ++i) {
       auto cur_in = op_item->operand_source(i);
diff --git a/paddle/fluid/pybind/control_flow_api.cc b/paddle/fluid/pybind/control_flow_api.cc
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
diff --git a/test/ir/pir/CMakeLists.txt b/test/ir/pir/CMakeLists.txt
diff --git a/test/ir/pir/test_assert_api.py b/test/ir/pir/test_assert_api.py
