ir_program_test.cc

// 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
//
//     http://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 <gtest/gtest.h>

#include "paddle/fluid/ir/dialect/pd_dialect.h"
#include "paddle/fluid/ir/dialect/pd_type.h"
#include "paddle/fluid/ir/dialect/utils.h"
#include "paddle/fluid/ir/interface/op_yaml_info.h"
#include "paddle/ir/core/block.h"
#include "paddle/ir/core/builtin_attribute.h"
#include "paddle/ir/core/builtin_dialect.h"
#include "paddle/ir/core/builtin_op.h"
#include "paddle/ir/core/ir_context.h"
#include "paddle/ir/core/program.h"
#include "paddle/ir/core/utils.h"
#include "paddle/phi/core/meta_tensor.h"
#include "paddle/phi/infermeta/binary.h"
#include "paddle/phi/kernels/elementwise_add_kernel.h"
// NOTE(zhangbo9674): File pd_op.h is generated by op_gen.py, see details in
// paddle/fluid/ir/dialect/CMakeLists.txt.
#include "paddle/fluid/ir/dialect/pd_op.h"

class AddOp : public ir::Op<AddOp> {
 public:
  using Op::Op;
  static const char *name() { return "test.add"; }
  static constexpr const char **attributes_name = nullptr;
  static constexpr uint32_t attributes_num = 0;
  void Verify();
  static void Build(ir::Builder &builder,             // NOLINT
                    ir::OperationArgument &argument,  // NOLINT
                    ir::OpResult l_operand,
                    ir::OpResult r_operand,
                    ir::Type sum_type);
};
void AddOp::Verify() {
  if (num_operands() != 2) {
    throw("The size of inputs must be equal to 2.");
  }
  if (num_results() != 1) {
    throw("The size of outputs must be equal to 1.");
  }
}
void AddOp::Build(ir::Builder &,
                  ir::OperationArgument &argument,
                  ir::OpResult l_operand,
                  ir::OpResult r_operand,
                  ir::Type sum_type) {
  argument.AddOperand(l_operand);
  argument.AddOperand(r_operand);
  argument.AddOutput(sum_type);
}
IR_DECLARE_EXPLICIT_TYPE_ID(AddOp)
IR_DEFINE_EXPLICIT_TYPE_ID(AddOp)

TEST(program_test, program) {
  // (1) Init environment.
  ir::IrContext *ctx = ir::IrContext::Instance();
  ir::Dialect *builtin_dialect =
      ctx->GetOrRegisterDialect<ir::BuiltinDialect>();
  builtin_dialect->RegisterOp<AddOp>();
  ir::Dialect *paddle_dialect =
      ctx->GetOrRegisterDialect<paddle::dialect::PaddleDialect>();

  // (2) Create an empty program object
  ir::Program program(ctx);

  // (3) Create a float32 DenseTensor Parameter and save into Program
  ir::Type fp32_dtype = ir::Float32Type::get(ctx);
  phi::DDim dims = {2, 2};
  phi::DataLayout data_layout = phi::DataLayout::NCHW;
  phi::LoD lod = {{0, 1, 2}};
  size_t offset = 0;
  ir::Type dense_tensor_dtype = paddle::dialect::DenseTensorType::get(
      ctx, fp32_dtype, dims, data_layout, lod, offset);

  std::vector<float> data_a = {1, 2, 3, 4};
  std::unique_ptr<ir::Parameter> parameter_a =
      std::make_unique<ir::Parameter>(reinterpret_cast<void *>(data_a.data()),
                                      4 * sizeof(float),
                                      dense_tensor_dtype);
  program.SetParameter("a", std::move(parameter_a));
  EXPECT_EQ(program.parameters_num() == 1, true);

  std::vector<float> data_b = {5, 6, 7, 8};
  std::unique_ptr<ir::Parameter> parameter_b =
      std::make_unique<ir::Parameter>(reinterpret_cast<void *>(data_b.data()),
                                      4 * sizeof(float),
                                      dense_tensor_dtype);
  program.SetParameter("b", std::move(parameter_b));
  EXPECT_EQ(program.parameters_num() == 2, true);

  // (4) Def a = GetParameterOp("a"), and create DenseTensor for a.
  ir::Builder builder(ctx, program.block());
  auto op1 = builder.Build<ir::GetParameterOp>("a", dense_tensor_dtype);

  EXPECT_EQ(&program, op1->GetParentProgram());
  EXPECT_EQ(op1->result(0).type().dialect().id(), paddle_dialect->id());
  using Interface = paddle::dialect::ParameterConvertInterface;
  Interface *a_interface =
      op1->result(0).type().dialect().GetRegisteredInterface<Interface>();
  std::shared_ptr<paddle::framework::Variable> a_var =
      a_interface->ParameterToVariable(program.GetParameter("a"));
  const phi::DenseTensor &a_tensor = a_var->Get<phi::DenseTensor>();
  EXPECT_EQ(a_tensor.numel(), 4);
  EXPECT_EQ(a_tensor.dims(), dims);
  EXPECT_EQ(a_tensor.dtype(), paddle::dialect::TransToPhiDataType(fp32_dtype));
  EXPECT_EQ(a_tensor.layout(), data_layout);
  EXPECT_EQ(a_tensor.lod(), lod);
  EXPECT_EQ(a_tensor.offset(), offset);
  for (int64_t i = 0; i < a_tensor.numel(); i++) {
    EXPECT_EQ(*(a_tensor.data<float>() + i), data_a[i]);
  }

  // (5) Def b = GetParameterOp("b"), and create DenseTensor for b.
  auto op2 = builder.Build<ir::GetParameterOp>("b", dense_tensor_dtype);

  EXPECT_EQ(op2->result(0).type().dialect().id(), paddle_dialect->id());
  Interface *b_interface =
      op2->result(0).type().dialect().GetRegisteredInterface<Interface>();
  std::shared_ptr<paddle::framework::Variable> b_var =
      b_interface->ParameterToVariable(program.GetParameter("b"));
  const phi::DenseTensor &b_tensor = b_var->Get<phi::DenseTensor>();
  EXPECT_EQ(b_tensor.numel(), 4);
  EXPECT_EQ(b_tensor.dims(), dims);
  EXPECT_EQ(b_tensor.dtype(), paddle::dialect::TransToPhiDataType(fp32_dtype));
  EXPECT_EQ(b_tensor.layout(), data_layout);
  EXPECT_EQ(b_tensor.lod(), lod);
  EXPECT_EQ(b_tensor.offset(), offset);
  for (int64_t i = 0; i < b_tensor.numel(); i++) {
    EXPECT_EQ(*(b_tensor.data<float>() + i), data_b[i]);
  }

  // (6) Def c = AddOp(a, b), execute this op.
  auto op3 =
      builder.Build<AddOp>(op1->result(0), op2->result(0), dense_tensor_dtype);
  phi::CPUContext *dev_ctx = static_cast<phi::CPUContext *>(
      paddle::platform::DeviceContextPool::Instance().Get(
          paddle::platform::CPUPlace()));
  phi::DenseTensor c_tensor =
      phi::Add<float, phi::CPUContext>(*dev_ctx, a_tensor, b_tensor);
  std::shared_ptr<paddle::framework::Variable> variable_c =
      std::make_shared<paddle::framework::Variable>();
  auto *dst_tensor = variable_c->GetMutable<phi::DenseTensor>();
  *dst_tensor = c_tensor;
  EXPECT_EQ(dst_tensor->numel(), b_tensor.numel());
  EXPECT_EQ(dst_tensor->dims(), b_tensor.dims());
  EXPECT_EQ(dst_tensor->dtype(), b_tensor.dtype());
  EXPECT_EQ(dst_tensor->layout(), b_tensor.layout());
  EXPECT_EQ(dst_tensor->lod(), b_tensor.lod());
  EXPECT_EQ(dst_tensor->offset(), b_tensor.offset());
  for (int64_t i = 0; i < dst_tensor->numel(); i++) {
    EXPECT_EQ(*(dst_tensor->data<float>() + i), data_a[i] + data_b[i]);
  }

  // (7) Def AbsOp(b)
  auto abs_op = builder.Build<paddle::dialect::AbsOp>(op1->result(0));
  paddle::dialect::OpYamlInfoInterface interface =
      abs_op->dyn_cast<paddle::dialect::OpYamlInfoInterface>();
  EXPECT_EQ(std::get<0>(interface.GetOpInfo())[0].name == "x", true);

  // (8) Def SetParameterOp(c, "c")
  auto op4 = builder.Build<ir::SetParameterOp>(op3->result(0), "c");

  EXPECT_EQ(op4->op_operand(0).type().dialect().id(), paddle_dialect->id());
  Interface *c_interface =
      op4->op_operand(0).type().dialect().GetRegisteredInterface<Interface>();
  //   ir::Parameter *parameter_c =
  //       c_interface->VariableToParameter(variable_c.get());
  std::unique_ptr<ir::Parameter> parameter_c =
      c_interface->VariableToParameter(variable_c.get());
  EXPECT_EQ(parameter_c->type(), dense_tensor_dtype);
  for (int64_t i = 0; i < dst_tensor->numel(); i++) {
    EXPECT_EQ(*(dst_tensor->data<float>() + i),
              *(static_cast<float *>(parameter_c->data()) + i));
  }
  program.SetParameter("c", std::move(parameter_c));

  // (8) Traverse Program
  EXPECT_EQ(program.block()->size() == 5, true);
  EXPECT_EQ(program.parameters_num() == 3, true);

  program.Print(std::cout);
}

TEST(program_test, slice_combine_test) {
  // (1) Init environment.
  ir::IrContext *ctx = ir::IrContext::Instance();
  ctx->GetOrRegisterDialect<ir::BuiltinDialect>();

  // (2) Create an empty program object
  ir::Program program(ctx);
  //   ir::Program *program = new ir::Program();
  EXPECT_EQ(program.block()->empty(), true);

  // (3) Create a float32 DenseTensor Parameter and save into Program
  ir::Type fp32_dtype = ir::Float32Type::get(ctx);

  // (4) Def a = GetParameterOp("a")
  std::string op1_name = ir::GetParameterOp::name();
  ir::OpInfo op1_info = ctx->GetRegisteredOpInfo(op1_name);
  std::unordered_map<std::string, ir::Attribute> op1_attribute{
      {"parameter_name", ir::StrAttribute::get(ctx, "a")}};
  ir::Operation *op1 =
      ir::Operation::Create({}, op1_attribute, {fp32_dtype}, op1_info);
  program.block()->push_back(op1);

  // (5) Def b = Constant("b")
  std::string op2_name = std::string(ir::ConstantOp::name());
  ir::OpInfo op2_info = ctx->GetRegisteredOpInfo(op2_name);
  ir::AttributeMap attr_map;
  attr_map.insert(std::pair<std::string, ir::Attribute>(
      "value", ir::FloatAttribute::get(ctx, 2.0)));
  ir::Operation *op2 =
      ir::Operation::Create({}, attr_map, {fp32_dtype}, op2_info);
  program.block()->push_back(op2);

  // (6) Def combine_op = CombineOp("a", "b")
  std::string combine_op_name = std::string(ir::CombineOp::name());
  ir::OpInfo combine_op_info = ctx->GetRegisteredOpInfo(combine_op_name);
  ir::Type output_type =
      ir::VectorType::get(ctx, std::vector<ir::Type>({fp32_dtype, fp32_dtype}));
  ir::Operation *combine_op = ir::Operation::Create(
      {op1->result(0), op2->result(0)}, {}, {output_type}, combine_op_info);
  ir::CombineOp combine_op_type = combine_op->dyn_cast<ir::CombineOp>();
  EXPECT_TRUE(combine_op_type.out());
  program.block()->push_back(combine_op);

  // (7) Def slice_op = SliceOp(combine_op, 0)
  std::string slice_op_name = std::string(ir::SliceOp::name());
  ir::OpInfo slice_op_info = ctx->GetRegisteredOpInfo(slice_op_name);
  ir::Attribute index_attr = ir::Int32Attribute::get(ctx, 0);
  ir::Operation *slice_op = ir::Operation::Create({combine_op->result(0)},
                                                  {{"index", index_attr}},
                                                  {fp32_dtype},
                                                  slice_op_info);
  program.block()->push_back(slice_op);

  // (8) Traverse Program
  EXPECT_EQ(program.block()->size() == 4, true);
}

TEST(program_test, builder) {
  ir::IrContext *ctx = ir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::PaddleDialect>();
  ir::Program program(ctx);
  ir::Builder builder = ir::Builder(ctx, program.block());

  paddle::dialect::FullOp full_op = builder.Build<paddle::dialect::FullOp>(
      std::vector<int64_t>{2, 2}, 1.5, phi::DataType::FLOAT32, phi::CPUPlace());
  ir::Type full_op_output = full_op->result(0).type();
  EXPECT_EQ(program.block()->size(), 1u);
  EXPECT_EQ(program.block()->back(), full_op.operation());
  EXPECT_EQ(full_op.num_operands(), 0u);
  EXPECT_EQ(full_op.num_results(), 1u);
  EXPECT_EQ(full_op.attributes().size(), 4u);
  EXPECT_EQ(
      full_op_output.dyn_cast<paddle::dialect::DenseTensorType>().offset() == 0,
      true);
  for (auto dim : phi::vectorize(
           full_op_output.dyn_cast<paddle::dialect::DenseTensorType>()
               .dims())) {
    EXPECT_EQ(dim == 2, true);
  }

  ir::ConstantOp constant = builder.Build<ir::ConstantOp>(
      ir::Int32Attribute::get(ctx, 2), ir::Int32Type::get(ctx));
  EXPECT_EQ(program.block()->size() == 2, true);
  EXPECT_EQ(constant.value().dyn_cast<ir::Int32Attribute>().data() == 2, true);
}