// Copyright (c) 2021 CINN 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 <functional>
#include <string>

#include "paddle/cinn/backends/codegen_cuda_dev.h"
#include "paddle/cinn/backends/llvm/execution_engine.h"
#include "paddle/cinn/cinn.h"
#include "paddle/cinn/common/test_helper.h"
#include "paddle/cinn/hlir/framework/graph_compiler.h"
#include "paddle/cinn/hlir/framework/node.h"
#include "paddle/cinn/hlir/framework/op.h"
#include "paddle/cinn/hlir/framework/op_strategy.h"
#include "paddle/cinn/hlir/op/use_ops.h"
#include "paddle/cinn/hlir/pe/broadcast.h"
#include "paddle/cinn/runtime/flags.h"

DECLARE_bool(cinn_ir_schedule);

namespace cinn {
namespace hlir {
namespace framework {

using CCompute = std::function<std::shared_ptr<ir::Tensor>(const std::vector<ir::Tensor>)>;

TEST(Operator, Operator_ElementWise_Add_Test0) {
  auto add      = Operator::Get("elementwise_add");
  Operator temp = *add;
  auto strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");

  Expr M(100), N(32);
  Placeholder<float> A("A", {M, N});
  Placeholder<float> B("B", {M, N});

  NodeAttr attrs;
  std::vector<ir::Tensor> inputs{A.tensor(), B.tensor()};
  std::vector<Type> type{Float(32)};
  common::Target target = common::DefaultHostTarget();
  auto impl = OpStrategy::SelectImpl(strategy[add](attrs, inputs, type, {{M.as_int32(), N.as_int32()}}, target));
  ASSERT_EQ(impl->name, "strategy.elementwise_add.x86");
  ASSERT_EQ(add->description, "elementwise_add function");

  std::string func_name = "add1";
  Module::Builder builder("module0", target);

  if (FLAGS_cinn_ir_schedule) {
    std::string out_name = "C";
    common::CINNValuePack cinn_input =
        common::CINNValuePack{{common::CINNValue(A), common::CINNValue(B), common::CINNValue(out_name)}};
    std::vector<std::string> input_output_names{"A", "B", out_name};

    auto funcs = framework::GetFuncFromImpl(impl, cinn_input, inputs, input_output_names, func_name, target);

    for (auto func : funcs) {
      LOG(INFO) << "Test Operator_ElementWise_Add_Test0's Strategy, func is :\n" << func;
      builder.AddFunction(func);
    }

  } else {
    common::CINNValuePack cinn_input = common::CINNValuePack{{common::CINNValue(A), common::CINNValue(B)}};
    common::CINNValuePack rets       = impl->fcompute(cinn_input);
    ASSERT_EQ(rets.size(), 2UL);
    rets = impl->fschedule(rets);
    ASSERT_EQ(rets.size(), 2UL);
    // the last element is a StageMap
    for (int i = 0; i < rets->size() - 1; i++) {
      Expr temp = rets[i];
      inputs.push_back(temp.as_tensor_ref());
    }
    auto func = Lower("fn_" + func_name, rets.back(), inputs);
    LOG(INFO) << "Test Strategy Codegen:\n" << func;
    builder.AddFunction(func);
  }

  auto jit    = backends::ExecutionEngine::Create({});
  auto module = builder.Build();
  jit->Link(module);
  auto fn = jit->Lookup("fn_" + func_name);
  CHECK(fn);
  auto fn_ = reinterpret_cast<void (*)(void *, int32_t)>(fn);
  cinn_buffer_t *A_buf;
  cinn_buffer_t *B_buf;
  int set_value = 0;
  if (set_value != 0) {
    A_buf = common::BufferBuilder(Float(32), {M.as_int32(), N.as_int32()}).set_align(512).set_val(set_value).Build();
    B_buf = common::BufferBuilder(Float(32), {M.as_int32(), N.as_int32()}).set_align(512).set_val(set_value).Build();
  } else {
    A_buf = common::BufferBuilder(Float(32), {M.as_int32(), N.as_int32()}).set_align(512).set_random().Build();
    B_buf = common::BufferBuilder(Float(32), {M.as_int32(), N.as_int32()}).set_align(512).set_random().Build();
  }
  auto *C_buf = common::BufferBuilder(Float(32), {M.as_int32(), N.as_int32()}).set_align(512).set_zero().Build();

  cinn_pod_value_t a_arg(A_buf), b_arg(B_buf), c_arg(C_buf);
  cinn_pod_value_t args[] = {a_arg, b_arg, c_arg};
  fn_(args, 3);

  auto *ad = reinterpret_cast<float *>(A_buf->memory);
  auto *bd = reinterpret_cast<float *>(B_buf->memory);
  auto *cd = reinterpret_cast<float *>(C_buf->memory);
  for (int i = 0; i < A_buf->num_elements(); i++) {
    ASSERT_NEAR(cd[i], ad[i] + bd[i], 1e-5);
  }
}
#ifdef CINN_WITH_CUDA
TEST(Operator, Operator_ElementWise_Add_Test1) {
  auto add      = Operator::Get("elementwise_add");
  Operator temp = *add;
  auto strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");

  Expr M(100), N(32);
  Placeholder<float> A("A", {M, N});
  Placeholder<float> B("B", {N});

  NodeAttr attrs;
  attrs.attr_store["axis"] = 1;
  std::vector<ir::Tensor> inputs{A.tensor(), B.tensor()};
  std::vector<Type> type{Float(32)};
  common::Target target = common::DefaultNVGPUTarget();
  auto impl             = OpStrategy::SelectImpl(strategy[add](attrs, inputs, type, {{100, 32}}, target));
  ASSERT_EQ(impl->name, "strategy.elementwise_add.x86");
  ASSERT_EQ(add->description, "elementwise_add function");

  std::string func_name = "add2";
  Module::Builder builder("module", target);

  if (FLAGS_cinn_ir_schedule) {
    std::string out_name = "C";
    common::CINNValuePack cinn_input =
        common::CINNValuePack{{common::CINNValue(A), common::CINNValue(B), common::CINNValue(out_name)}};
    std::vector<std::string> input_output_names{"A", "B", out_name};

    auto funcs = framework::GetFuncFromImpl(impl, cinn_input, inputs, input_output_names, func_name, target);

    for (auto func : funcs) {
      builder.AddFunction(func);
      LOG(INFO) << "Test Operator_ElementWise_Add_Test1's Strategy, func is :\n" << func;
    }
  } else {
    common::CINNValuePack cinn_input = common::CINNValuePack{{common::CINNValue(A), common::CINNValue(B)}};
    common::CINNValuePack rets       = impl->fcompute(cinn_input);
    ASSERT_EQ(rets.size(), 2UL);
    rets = impl->fschedule(rets);
    ASSERT_EQ(rets.size(), 2UL);
    // the last element is a StageMap
    for (int i = 0; i < rets->size() - 1; i++) {
      Expr temp = rets[i];
      inputs.push_back(temp.as_tensor_ref());
    }
    auto func = Lower("fn_" + func_name, rets.back(), inputs);
    LOG(INFO) << "Test Strategy Codegen:\n" << func;
    builder.AddFunction(func);
  }

  backends::CodeGenCUDA_Dev codegen(target);

  auto module      = builder.Build();
  auto source_code = codegen.Compile(module);
  LOG(INFO) << "Operator_ElementWise_Add_Test1 source code:\n" << source_code;
}
#endif

TEST(Operator, Operator_BroadcastTo) {
  auto broadcast_to = Operator::Get("broadcast_to");
  Operator temp     = *broadcast_to;
  auto strategy     = Operator::GetAttrs<StrategyFunction>("CINNStrategy");

  Expr N(1);
  Placeholder<float> B("B", {N});

  NodeAttr attrs;
  std::vector<int> out_shape    = {16};
  attrs.attr_store["out_shape"] = out_shape;

  std::vector<int> broadcast_axes    = {0};
  attrs.attr_store["broadcast_axes"] = broadcast_axes;

  std::vector<ir::Tensor> inputs{B.tensor()};
  std::vector<Type> type{Float(32)};
  common::Target target = common::DefaultHostTarget();

  auto impl = OpStrategy::SelectImpl(strategy[broadcast_to](attrs, inputs, type, {out_shape}, target));

  std::string func_name = "broadcast_to";

  if (FLAGS_cinn_ir_schedule) {
    std::string out_name             = "C";
    common::CINNValuePack cinn_input = common::CINNValuePack{{common::CINNValue(B), common::CINNValue(out_name)}};
    std::vector<std::string> input_output_names{"B", out_name};

    auto funcs = framework::GetFuncFromImpl(impl, cinn_input, inputs, input_output_names, func_name, target);

    for (auto func : funcs) {
      LOG(INFO) << "Test Operator_BroadcastTo's Strategy, func is :\n" << func;
    }
  } else {
    common::CINNValuePack cinn_input = common::CINNValuePack{{common::CINNValue(B)}};
    common::CINNValuePack rets       = impl->fcompute(cinn_input);

    ASSERT_EQ(rets.size(), 2UL);
    rets = impl->fschedule(rets);
    ASSERT_EQ(rets.size(), 2UL);
    // the last element is a StageMap
    for (int i = 0; i < rets->size() - 1; i++) {
      Expr temp = rets[i];
      inputs.push_back(temp.as_tensor_ref());
    }

    auto func = Lower("func" + func_name, rets.back(), inputs);
    LOG(INFO) << "Test Operator_BroadcastTo's Strategy, func is :\n" << func;
  }
}

common::CINNValuePack GetComputeResult(const std::shared_ptr<OpImpl> &impl,
                                       std::vector<common::CINNValue> &cinn_inputs,
                                       const std::string &output_name = "") {
  if (FLAGS_cinn_ir_schedule) {
    cinn_inputs.emplace_back(output_name);
  }
  return impl->fcompute(common::CINNValuePack{cinn_inputs});
}

TEST(Operator, Operator_BroadcastTo_0) {
  auto const_scalar    = Operator::Get("const_scalar");
  auto broadcast_to    = Operator::Get("broadcast_to");
  auto reduce_sum      = Operator::Get("reduce_sum");
  auto elementwise_add = Operator::Get("elementwise_mul");

  auto strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");

  Expr N(16);
  Placeholder<float> A("A", {N, N, N, N});

  NodeAttr attrs;
  attrs.attr_store["value"] = 0.5f;

  std::vector<int> out_shape    = {16};
  attrs.attr_store["out_shape"] = out_shape;

  std::vector<int> broadcast_axes    = {0};
  attrs.attr_store["broadcast_axes"] = broadcast_axes;

  std::vector<int> dim    = {0, 2, 3};
  attrs.attr_store["dim"] = dim;

  std::vector<Type> type{Float(32)};
  common::Target target = common::DefaultHostTarget();

  auto impl_0 =
      OpStrategy::SelectImpl(strategy[const_scalar](attrs, std::vector<ir::Tensor>{}, type, {out_shape}, target));
  std::vector<common::CINNValue> cinn_inputs;
  common::CINNValuePack rets_0 = GetComputeResult(impl_0, cinn_inputs, "out_0");
  ir::Expr out_0               = rets_0[0];
  auto tensor_0                = out_0.as_tensor_ref();
  poly::StageMap stages_0      = rets_0.back();

  auto impl_1 = OpStrategy::SelectImpl(strategy[broadcast_to](attrs, {tensor_0}, type, {out_shape}, target));
  std::vector<common::CINNValue> cinn_inputs_1 = {{common::CINNValue(tensor_0)}};
  common::CINNValuePack rets_1                 = GetComputeResult(impl_1, cinn_inputs_1, "out_1");

  ir::Expr out_1          = rets_1[0];
  auto tensor_1           = out_1.as_tensor_ref();
  poly::StageMap stages_1 = rets_1.back();

  auto impl_2 = OpStrategy::SelectImpl(strategy[reduce_sum](attrs, {A.tensor()}, type, {out_shape}, target));
  std::vector<common::CINNValue> cinn_inputs_2 = {{common::CINNValue(A.tensor())}};
  common::CINNValuePack rets_2                 = GetComputeResult(impl_2, cinn_inputs_2, "out_2");

  ir::Expr out_2          = rets_2[0];
  auto tensor_2           = out_2.as_tensor_ref();
  poly::StageMap stages_2 = rets_2.back();

  std::vector<common::CINNValue> cinn_inputs_4 = {{common::CINNValue(A.tensor())}};
  common::CINNValuePack rets_4                 = GetComputeResult(impl_2, cinn_inputs_4, "out_4");
  ir::Expr out_4                               = rets_4[0];
  auto tensor_4                                = out_4.as_tensor_ref();
  poly::StageMap stages_4                      = rets_4.back();

  auto impl_3 =
      OpStrategy::SelectImpl(strategy[elementwise_add](attrs, {tensor_1, tensor_2}, type, {out_shape}, target));
  std::vector<common::CINNValue> cinn_inputs_3 = {{common::CINNValue(tensor_1), common::CINNValue(tensor_2)}};
  common::CINNValuePack rets_3                 = GetComputeResult(impl_3, cinn_inputs_3, "out_3");

  ir::Expr out_3          = rets_3[0];
  auto tensor_3           = out_3.as_tensor_ref();
  poly::StageMap stages_3 = rets_3.back();

  stages_3->InsertLazily(tensor_0, stages_0[tensor_0]);
  stages_3->InsertLazily(tensor_1, stages_1[tensor_1]);
  stages_3->InsertLazily(tensor_2, stages_2[tensor_2]);
  stages_3->InsertLazily(tensor_4, stages_4[tensor_4]);
  stages_3[tensor_0]->ComputeInline();
  stages_3[tensor_1]->ComputeInline();
  stages_3[tensor_2]->SetBuffer("local");
  stages_3[tensor_4]->SimpleComputeAt(stages_3[tensor_2], 3);
  stages_3[tensor_2]->SimpleComputeAt(stages_3[tensor_3], 0);

  std::vector<ir::Tensor> inputs = {A.tensor(), tensor_3, tensor_4};
  auto func                      = Lower("broadcast_to", stages_3, inputs);
  LOG(INFO) << "Test Strategy Codegen:\n" << func;
}

}  // namespace framework
}  // namespace hlir
}  // namespace cinn