// 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 "paddle/cinn/hlir/pe/transform.h"

#include <algorithm>

#include "paddle/cinn/common/cas.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/op_util.h"
#include "paddle/cinn/hlir/pe/elementwise.h"
#include "paddle/cinn/hlir/pe/ir_schedule_pe.h"
#include "paddle/cinn/hlir/pe/nn.h"
#include "paddle/cinn/hlir/pe/schedule.h"
#include "paddle/cinn/ir/ir_printer.h"
#include "paddle/cinn/utils/string.h"

DECLARE_bool(cinn_ir_schedule);

namespace cinn {
namespace hlir {
namespace op {
using common::_CINNValuePack_;
using common::CINNValue;
using common::CINNValuePack;
using framework::OpStrategy;
using framework::shape_t;
using framework::StrategyFunction;

std::shared_ptr<OpStrategy> StrategyForMatMul(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  const auto &attr_store = attrs.attr_store;
  bool trans_a = SafeGetAttr(attr_store, "trans_a", false);
  bool trans_b = SafeGetAttr(attr_store, "trans_b", false);
  float alpha = SafeGetAttr(attr_store, "alpha", 1.0f);

  const auto &shape_A = ToPodVector<int>(inputs[0]->shape);
  const auto &shape_B = ToPodVector<int>(inputs[1]->shape);

  const auto &new_shape =
      pe::utils::GetMatmulNewShapes({shape_A, shape_B}, trans_a, trans_b);

  const auto &new_shape_A = new_shape[0];
  const auto &new_shape_B = new_shape[1];
  const auto &output_shape = new_shape[2];

  framework::CINNCompute matmul_compute([=](lang::Args args,
                                            lang::RetValue *ret) {
    CHECK(!args.empty())
        << "The input arguments of Matmul compute is empty! Please check.\n";
    CINNValuePack pack_args = args[0];
    CHECK_GE(pack_args.size(), 2U)
        << "at least 2 input tensors for Matmul compute\n";
    Expr A = pack_args[0];
    Expr B = pack_args[1];
    CHECK(A.as_tensor());
    CHECK(B.as_tensor());

    std::string tensor_name = UniqName("MatMul");
    if (FLAGS_cinn_ir_schedule) {
      CHECK_GE(pack_args.size(), 3);
      CHECK(pack_args[2].is_string());
      tensor_name = pack_args[2].operator std::string();
    }

    auto tensor_A = A.as_tensor_ref();
    auto tensor_B = B.as_tensor_ref();
    auto stages = CreateStages({tensor_A, tensor_B});

    auto new_shape_A_e = ToCinnExprs(new_shape_A);
    auto new_shape_B_e = ToCinnExprs(new_shape_B);

    auto new_A = tensor_A->Reshape(new_shape_A_e, stages);
    auto new_B = tensor_B->Reshape(new_shape_B_e, stages);

    std::vector<ir::Tensor> out;
    if (target.arch == Target::Arch::X86) {
#ifdef CINN_WITH_MKL_CBLAS
      out = pe::MatmulMKL(new_A,
                          new_B,
                          trans_a,
                          trans_b,
                          alpha,
                          UniqName("MatmulMKL_output"),
                          target);
#else
      out = pe::MatmulV2(new_A,
                         new_B,
                         trans_a,
                         trans_b,
                         alpha,
                         UniqName("MatmulV2_output"),
                         target);
#endif
    } else {
      out = pe::Matmul(new_A, new_B, trans_a, trans_b, alpha, tensor_name);
    }

    std::vector<CINNValue> res;
    for (auto &t : out) {
      stages->InsertLazily(t);
    }

    for (auto &t : out) {
      res.push_back(CINNValue(t));
    }
    res.push_back(CINNValue(stages));
    *ret = CINNValuePack{res};
  });

  framework::CINNSchedule matmul_schedule([=](lang::Args args,
                                              lang::RetValue *ret) {
    CHECK(!args.empty())
        << "The input argument of matmul schedule is empty! Please check.\n";
    CINNValuePack arg_pack = args[0];
    if (FLAGS_cinn_ir_schedule) {
      std::vector<CINNValue> results =
          pe::IRCudaScheduleMatMul(arg_pack, output_shape, target);
      *ret = CINNValuePack({results});
    } else {
      CHECK(arg_pack.size() == 2UL || arg_pack.size() == 3UL);
      poly::StageMap stages = arg_pack.back();
      if (target.arch == Target::Arch::NVGPU) {
        Expr out = arg_pack[0];
        CHECK(out.as_tensor());
        pe::MatmulScheduleCUDA(stages, out.as_tensor_ref(), target);
      } else if (target.arch == Target::Arch::X86) {
#ifdef CINN_WITH_MKL_CBLAS
        CHECK_EQ(arg_pack.size(), 3UL);
#else
        CHECK_EQ(arg_pack.size(), 3UL);
        Expr out = arg_pack[0];
        Expr packedB = arg_pack[1];
        CHECK(packedB.as_tensor());
        CHECK(out.as_tensor());
        pe::MatmulScheduleCPU(
            stages, out.as_tensor_ref(), packedB.as_tensor_ref(), target);
#endif
      }
      *ret = arg_pack;
    }
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(matmul_compute, matmul_schedule, "strategy.matmul.x86", 1);

  return strategy;
}

std::vector<std::vector<int>> InferShapeForMatMul(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK_EQ(inputs_shape.size(), 2UL)
      << "The input's shape size should be 2! Please check again.";
  bool trans_a = SafeGetAttr(attrs, "trans_a", false);
  bool trans_b = SafeGetAttr(attrs, "trans_b", false);

  VLOG(4) << "During the matmul shape inference, origin shape_A: "
          << utils::Join(inputs_shape[0], ", ");
  VLOG(4) << "During the matmul shape inference, origin shape_B: "
          << utils::Join(inputs_shape[1], ", ");

  const auto &new_shape =
      pe::utils::GetMatmulNewShapes(inputs_shape, trans_a, trans_b);

  const auto &new_shape_A = new_shape[0];
  const auto &new_shape_B = new_shape[1];
  const auto &output_shape = new_shape[2];

  VLOG(4) << "During the matmul shape inference, new_shape_A: "
          << utils::Join(new_shape_A, ", ");
  VLOG(4) << "During the matmul shape inference, new_shape_B: "
          << utils::Join(new_shape_B, ", ");
  VLOG(4) << "During the matmul shape inference, output_shape: "
          << utils::Join(output_shape, ", ");

  std::vector<std::vector<int>> res{output_shape};
  return res;
}

std::vector<Type> InferDtypeForMatMul(const std::vector<Type> &inputs_type,
                                      const framework::AttrMapType &attrs) {
  CHECK_EQ(inputs_type.size(), 2UL)
      << "The input's type size should be 2! Please check again.";
  CHECK_EQ(inputs_type[0], inputs_type[1])
      << "The input's types should be equal! Please check again.";

  std::vector<Type> res{inputs_type[0]};
  return res;
}

std::vector<std::vector<std::string>> InferLayoutForMatMul(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_EQ(input_layouts.size(), 2U)
      << "The input's layouts size is not 2! Please check again.";
  CHECK_EQ(input_shapes.size(), 2U) << "mul should have 2 input shapes";
  std::vector<std::string> new_input_layouts = input_layouts;
  for (int i = 0; i < input_shapes.size(); i++) {
    if (input_shapes[i].size() > 4) {
      // alter input layout back
      new_input_layouts[i] = "NCHW";
    }
  }

  return {{"", ""}, new_input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForSplit(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  // get attribute
  std::vector<int> sections;
  int axis = 0;
  if (attrs.attr_store.find("num_or_sections") != attrs.attr_store.end()) {
    sections =
        absl::get<std::vector<int>>(attrs.attr_store.at("num_or_sections"));
  }
  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
    axis = absl::get<int>(attrs.attr_store.at("axis"));
  }
  if (axis < 0) axis += static_cast<int>(output_shapes[0].size());

  CHECK(!output_shapes.empty())
      << "The Spilt Op's output shape list should not empty.";
  CHECK_LT(axis, static_cast<int>(output_shapes[0].size()));
  CHECK(!sections.empty())
      << "The Split op doesn't find [num_or_sections] attrbute! It it a "
         "mandatory attribute ! Please check.";

  framework::CINNCompute split_compute(
      [=](lang::Args args, lang::RetValue *ret) {
        CHECK(!args.empty())
            << "The input arguments of split compute is empty! Please check.";
        CINNValuePack pack_args = args[0];
        CHECK(!pack_args.empty())
            << "The input tensors of split compute is empty! Please check.";
        Expr A_expr = pack_args[0];
        CHECK(A_expr.as_tensor());
        ir::Tensor A = A_expr.as_tensor_ref();

        std::vector<std::string> tensor_names;
        if (FLAGS_cinn_ir_schedule) {
          CHECK_EQ(pack_args.size(), output_shapes.size() + 1);
          for (int idx = 1; idx < pack_args.size(); ++idx) {
            CHECK(pack_args[idx].is_string());
            tensor_names.push_back(pack_args[idx].operator std::string());
          }
        } else {
          for (int idx = 0; idx < output_shapes.size(); ++idx) {
            tensor_names.push_back(UniqName("T_Split_Out"));
          }
        }

        auto out = pe::Split(A, axis, output_shapes, tensor_names);
        auto stages = CreateStages(out);

        std::vector<CINNValue> res;
        for (int i = 0; i < out.size(); ++i) {
          res.emplace_back(out[i]);
        }
        res.emplace_back(stages);
        *ret = CINNValuePack{res};
      });

  framework::CINNSchedule split_schedule([=](lang::Args args,
                                             lang::RetValue *ret) {
    if (FLAGS_cinn_ir_schedule) {
      CHECK(!args.empty())
          << "The input argument of split schedule is empty! Please check.";
      CINNValuePack arg_pack = args[0];
      std::vector<Expr> vec_ast;
      for (int i = 0; i < arg_pack.size(); i++) {
        if (arg_pack[i].is_expr()) {
          Expr temp = arg_pack[i];
          vec_ast.emplace_back(temp);
        }
      }
      CHECK(!vec_ast.empty());
      ir::ModuleExpr mod_expr(vec_ast);
      ir::IRSchedule ir_sch(mod_expr);
      ir_sch.MergeExprs();
      pe::IRCudaSplitSchedule(ir_sch, output_shapes, axis, target);
      std::vector<CINNValue> res{
          CINNValue(ir_sch.GetModule().GetExprs().at(0))};
      *ret = CINNValuePack{res};
    } else {
      CHECK(!args.empty())
          << "The input arguments of split schedule is empty! Please check.";
      CINNValuePack arg_pack = args[0];
      CHECK_GE(arg_pack.size(), 2UL)
          << "The input tensor's size of split schedule is " << arg_pack.size()
          << "and it should be greater equal to 2! Please check.";
      pe::CudaSplitSchedule(&arg_pack, output_shapes, axis, target);
      *ret = arg_pack;
    }
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(split_compute, split_schedule, "strategy.split.x86", 1);

  return strategy;
}

std::vector<std::vector<int>> InferShapeForSplit(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  std::vector<int> sections;
  if (attrs.find("num_or_sections") != attrs.end()) {
    sections = absl::get<std::vector<int>>(attrs.at("num_or_sections"));
  } else {
    LOG(FATAL) << "The Split op doesn't find [num_or_sections] attrbute! It it "
                  "a mandatory attribute ! Please check.";
  }

  if (inputs_shape.empty()) {
    std::vector<std::vector<int>> ret;
    if (sections.size() == 1) {
      ret.resize(sections[0]);
    } else {
      ret.resize(sections.size());
    }
    return ret;
  }
  CHECK_GE(inputs_shape.size(), 1U)
      << "The input's shape size should be no less than 1! Please check again.";

  int axis = 0;
  if (attrs.find("axis") != attrs.end()) {
    axis = absl::get<int>(attrs.at("axis"));
    if (axis < 0) {
      axis += inputs_shape[0].size();
    }
  }

  // check sections valid
  int output_size = sections.size();
  int pivot = inputs_shape[0][axis];

  auto real_sections = sections;
  if (output_size == 1) {
    // if the 'sections' is a number, the tensor will split to 'sections'
    // sub-tensor, each sub-tensor length A[axis] / 'sections'
    output_size = sections[0];
    CHECK_EQ(pivot % output_size, 0)
        << "If the attribute 'num_or_sections' is a number, it should be "
           "divisible by the "
           "axis's dimension of inputs A ! Please check.";
    real_sections.assign(output_size, pivot / output_size);
  } else {
    // else the tensor will split to sections.size sub-tensor, each sub-tensor
    // length sections[i] The sections may have at most one '-1' in sections,
    // that means its value should be inferred by others.
    int section_sum = 0, neg_index = -1;
    for (int i = 0; i < output_size; ++i) {
      if (sections[i] > 0) {
        section_sum += sections[i];
      } else if (sections[i] == -1 && neg_index < 0) {
        neg_index = i;
      } else {
        if (sections[i] == 0) {
          LOG(FATAL) << "The attribute 'num_or_sections' should not has 0 ! "
                        "Please check.";
        } else {
          LOG(FATAL) << "The attribute 'num_or_sections' can only have at most "
                        "one '-1' ! Please check.";
        }
      }
    }

    if (neg_index >= 0) {
      // has '-1' in sections
      real_sections[neg_index] = pivot - section_sum;
    } else {
      CHECK_EQ(pivot, section_sum)
          << "The sum of attr sections should be equal with the axis's "
             "dimension value of "
             "inputs A in Split ! Please check.";
    }
  }

  std::vector<std::vector<int>> outputs_shape(output_size, inputs_shape[0]);
  for (int i = 0; i < output_size; ++i) {
    outputs_shape[i][axis] = real_sections[i];
  }
  return outputs_shape;
}

std::vector<Type> InferDtypeForSplit(const std::vector<Type> &inputs_type,
                                     const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";

  std::vector<int> sections;
  if (attrs.find("num_or_sections") != attrs.end()) {
    sections = absl::get<std::vector<int>>(attrs.at("num_or_sections"));
  } else {
    LOG(FATAL) << "The Split op doesn't find [num_or_sections] attrbute! It it "
                  "a mandatory attribute ! Please check.";
  }

  int output_size = sections.size();
  if (output_size == 1) {
    output_size = sections[0];
  }

  std::vector<Type> res(output_size, inputs_type[0]);
  return res;
}

std::vector<std::vector<std::string>> InferLayoutForSplit(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK(!input_layouts.empty())
      << "The input's layout size is 0! Please check again.";
  std::vector<int> sections;
  if (attrs.attr_store.find("num_or_sections") != attrs.attr_store.end()) {
    sections =
        absl::get<std::vector<int>>(attrs.attr_store.at("num_or_sections"));
  } else {
    LOG(FATAL) << "The Split op doesn't find [num_or_sections] attrbute! It it "
                  "a mandatory attribute ! Please check.";
  }

  int output_size = sections.size();
  if (output_size == 1) {
    output_size = sections[0];
  }

  std::vector<std::string> output_layout(output_size, input_layouts[0]);
  return {output_layout, input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForConcat(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  framework::CINNCompute concat_compute([=](lang::Args args,
                                            lang::RetValue *ret) {
    CHECK(!args.empty())
        << "The input arguments of Concat compute is empty! Please check.\n";
    CHECK(!out_type.empty())
        << "Output type of Concat is empty! Please check.\n";
    CINNValuePack pack_args = args[0];
    int input_size =
        FLAGS_cinn_ir_schedule ? pack_args.size() - 1 : pack_args.size();
    CHECK_GE(input_size, 1UL)
        << "at least 2 input tensors for Concat compute\n";
    CHECK(!output_shapes.empty());
    int axis = 0;
    if (attrs.attr_store.count("axis")) {
      axis = absl::get<int>(attrs.attr_store.at("axis"));
    }

    std::vector<ir::Tensor> input_tensors;
    for (int i = 0; i < input_size; i++) {
      Expr tensor = pack_args[i];
      CHECK(tensor.as_tensor());
      input_tensors.push_back(tensor.as_tensor_ref());
    }

    std::string tensor_name = UniqName("Concat_output");
    if (FLAGS_cinn_ir_schedule) {
      CHECK(pack_args[input_size].is_string());
      tensor_name = pack_args[input_size].operator std::string();
    }

    auto stages = CreateStages(input_tensors);
    auto out = pe::Concat(input_tensors, axis, tensor_name);
    stages->InsertLazily(out);

    *ret = CINNValuePack({CINNValue(out), CINNValue(stages)});
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(concat_compute,
                    GetInjectiveScheduleFunc(output_shapes, target, false),
                    "strategy.concat.x86",
                    1);
  return strategy;
}

std::vector<std::vector<int>> InferShapeForConcat(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK_GE(inputs_shape.size(), 1UL)
      << "The input's shape size should be no less than 2! Please check again.";
  int axis = 0;
  for (auto &iter : attrs) {
    if (iter.first == "axis") {
      axis = absl::get<int>(iter.second);
      break;
    }
  }

  if (axis < 0) axis += inputs_shape[0].size();
  std::vector<int> output_shape = inputs_shape[0];
  CHECK(axis >= 0 && axis < inputs_shape[0].size())
      << "In Concat op, the attribute `axis` should be >= 0 and < input "
         "shape's size, please check!";

  int input_dim = inputs_shape[0].size();
  for (int i = 1; i < inputs_shape.size(); i++) {
    CHECK_EQ(inputs_shape[i].size(), input_dim)
        << "Dimensions of inputs tensors in Concat should be equal! Please "
           "check.";

    for (int j = 0; j < input_dim; j++) {
      if (j != axis) {
        CHECK_EQ(inputs_shape[0][j], inputs_shape[i][j])
            << "The " << j << "-th dimension of input[0] and input[" << i
            << "] should be the same, but here input[0].shape=["
            << cinn::utils::Join(inputs_shape[0], ", ") << "], input[" << i
            << "].shape=[" << cinn::utils::Join(inputs_shape[i], ", ")
            << "]! Please check.";
      }
    }

    output_shape[axis] += inputs_shape[i][axis];
  }

  std::vector<std::vector<int>> res{output_shape};
  return res;
}

std::vector<Type> InferDtypeForConcat(const std::vector<Type> &inputs_type,
                                      const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";
  std::vector<Type> res{inputs_type[0]};
  return res;
}

std::vector<std::vector<std::string>> InferLayoutForConcat(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_GE(input_layouts.size(), 1UL)
      << "The input's layout size is less than 2! Please check again.";
  return {{input_layouts[0]}, input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForMul(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  CHECK_EQ(inputs.size(), 2UL) << "mul should have 2 input";
  const auto &attr_store = attrs.attr_store;
  int x_num_col_dims = SafeGetAttr(attr_store, "x_num_col_dims", 1);
  int y_num_col_dims = SafeGetAttr(attr_store, "y_num_col_dims", 1);
  bool is_infer = SafeGetAttr(attr_store, "is_infer", false);

  const auto &shape_A = ToPodVector<int>(inputs[0]->shape);
  const auto &shape_B = ToPodVector<int>(inputs[1]->shape);

  const auto &new_shape = pe::utils::GetMulNewShapes(
      {shape_A, shape_B}, x_num_col_dims, y_num_col_dims, is_infer);

  const auto &new_shape_A = new_shape[0];
  const auto &new_shape_B = new_shape[1];
  const auto &output_shape = new_shape[2];

  framework::CINNCompute mul_compute(
      [=](lang::Args args, lang::RetValue *ret) {
        CHECK(!args.empty())
            << "The input arguments of Mul compute is empty! Please check.\n";
        CINNValuePack pack_args = args[0];
        CHECK_GE(pack_args.size(), 2U)
            << "at least 2 input tensors for Mul compute\n";
        Expr A = pack_args[0];
        Expr B = pack_args[1];
        CHECK(A.as_tensor());
        CHECK(B.as_tensor());

        auto A_tensor = A.as_tensor_ref();
        auto B_tensor = B.as_tensor_ref();
        auto stages = CreateStages({A_tensor, B_tensor});

        auto new_shape_A_e = ToCinnExprs(new_shape_A);
        auto new_shape_B_e = ToCinnExprs(new_shape_B);

        auto new_A = A_tensor->Reshape(new_shape_A_e, stages);
        auto new_B = B_tensor->Reshape(new_shape_B_e, stages);

        std::vector<ir::Tensor> out;
        std::string tensor_name = UniqName("Mul_output");
        if (FLAGS_cinn_ir_schedule) {
          CHECK(pack_args.back().is_string());
          tensor_name = pack_args.back().operator std::string();
        }

        if (target.arch == Target::Arch::X86) {
#ifdef CINN_WITH_MKL_CBLAS
          out = pe::MatmulMKL(
              new_A, new_B, false, is_infer, 1.0f, tensor_name, target);
#else
          out = pe::MatmulV2(
              new_A, new_B, false, is_infer, 1.0f, tensor_name, target);
#endif
        } else {
          out = pe::Matmul(new_A, new_B, false, is_infer, 1.0f, tensor_name);
        }

        std::vector<CINNValue> res;
        for (auto &t : out) {
          stages->InsertLazily(t);
        }

        for (auto &t : out) {
          res.push_back(CINNValue(t));
        }
        res.push_back(CINNValue(stages));
        *ret = CINNValuePack{res};
      });

  framework::CINNSchedule mul_schedule([=](lang::Args args,
                                           lang::RetValue *ret) {
    CHECK(!args.empty())
        << "The input argument of matmul schedule is empty! Please check.\n";
    CINNValuePack arg_pack = args[0];
    if (FLAGS_cinn_ir_schedule) {
      std::vector<CINNValue> results =
          pe::IRCudaScheduleMatMul(arg_pack, output_shape, target);
      *ret = CINNValuePack({results});
    } else {
      CHECK(arg_pack.size() == 2UL || arg_pack.size() == 3UL);
      poly::StageMap stages = arg_pack.back();
      if (target.arch == Target::Arch::NVGPU) {
        Expr out = arg_pack[0];
        CHECK(out.as_tensor());
        pe::MatmulScheduleCUDA(stages, out.as_tensor_ref(), target);
      } else if (target.arch == Target::Arch::X86) {
#ifdef CINN_WITH_MKL_CBLAS
        CHECK_EQ(arg_pack.size(), 3UL);
#else
        CHECK_EQ(arg_pack.size(), 3UL);
        Expr out = arg_pack[0];
        Expr packedB = arg_pack[1];
        CHECK(packedB.as_tensor());
        CHECK(out.as_tensor());
        pe::MatmulScheduleCPU(
            stages, out.as_tensor_ref(), packedB.as_tensor_ref(), target);
#endif
      }
      *ret = arg_pack;
    }
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(mul_compute, mul_schedule, "strategy.mul.x86", 1);

  return strategy;
}

std::vector<std::vector<int>> InferShapeForMul(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK_EQ(inputs_shape.size(), 2U)
      << "The input's shape size should be 2! Please check again.";
  CHECK_GE(inputs_shape[0].size(), 2U)
      << "Input matrix X's dim should be >= 2! Please check.";
  CHECK_GE(inputs_shape[1].size(), 2U)
      << "Input matrix Y's dim should be >= 2! Please check.";

  VLOG(4) << "During the matmul shape inference, origin shape_A: "
          << utils::Join(inputs_shape[0], ", ");
  VLOG(4) << "During the matmul shape inference, origin shape_B: "
          << utils::Join(inputs_shape[1], ", ");

  int x_num_col_dims = SafeGetAttr(attrs, "x_num_col_dims", 1);
  int y_num_col_dims = SafeGetAttr(attrs, "y_num_col_dims", 1);
  bool is_infer = SafeGetAttr(attrs, "is_infer", false);

  const auto &new_shape = pe::utils::GetMulNewShapes(
      inputs_shape, x_num_col_dims, y_num_col_dims, is_infer);

  const auto &new_shape_A = new_shape[0];
  const auto &new_shape_B = new_shape[1];
  const auto &output_shape = new_shape[2];

  VLOG(4) << "During the mul shape inference, new_shape_A: "
          << utils::Join(new_shape_A, ", ");
  VLOG(4) << "During the mul shape inference, new_shape_B: "
          << utils::Join(new_shape_B, ", ");
  VLOG(4) << "During the mul shape inference, output_shape: "
          << utils::Join(output_shape, ", ");

  int a_K = new_shape_A[1];
  int b_K = is_infer ? new_shape_B[1] : new_shape_B[0];

  CHECK_EQ(a_K, b_K) << "The K dimension of mul should be equal.";

  return {output_shape};
}

std::vector<Type> InferDtypeForMul(const std::vector<Type> &inputs_type,
                                   const framework::AttrMapType &attrs) {
  CHECK_EQ(inputs_type.size(), 2U)
      << "The input's type size should be 2! Please check again.";
  CHECK_EQ(inputs_type[0], inputs_type[1])
      << "The input's types should be equal! Please check again.";

  return {inputs_type[0]};
}

std::vector<std::vector<std::string>> InferLayoutForMul(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_EQ(input_layouts.size(), 2U)
      << "The input's layouts size is not 2! Please check again.";
  CHECK_EQ(input_shapes.size(), 2U) << "mul should have 2 input shapes";
  std::vector<std::string> new_input_layouts = input_layouts;
  for (int i = 0; i < input_shapes.size(); i++) {
    if (input_shapes[i].size() > 4) {
      // alter input layout back
      new_input_layouts[i] = "NCHW";
    }
  }

  return {{"", ""}, new_input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForCublasGemm(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  framework::CINNCompute gemm_compute(
      [attrs](lang::Args args, lang::RetValue *ret) {
        auto &attr_store = attrs.attr_store;
        CHECK(attr_store.contains("trans_a"))
            << "The cublas_gemm should have an attr named `trans_a`.";
        CHECK(attr_store.contains("trans_b"))
            << "The cublas_gemm should have an attr named `trans_b`.";
        CHECK(!args.empty())
            << "The input `args` of cublas_gemm is empty! Please check.";

        CINNValuePack input_args = args[0];
        CHECK_GE(input_args.size(), 3U)
            << "The input number of cublas_gemm should be equal to 3.";
        Expr lhs = input_args[0];
        Expr rhs = input_args[1];
        Expr bias = input_args[2];
        CHECK(lhs.as_tensor());
        CHECK(rhs.as_tensor());
        CHECK(bias.as_tensor());
        auto bias_tensor = bias.as_tensor_ref();
        // dummy gemm computation, which will be replaced by
        // cinn_gpu_cublas_gemm in the GemmRewriter pass.

        std::string tensor_name = UniqName("cublas_gemm_output");
        if (FLAGS_cinn_ir_schedule) {
          CHECK_EQ(input_args.size(), 4);
          CHECK(input_args[3].is_string());
          tensor_name = input_args[3].operator std::string();
        }
        auto out = pe::Identity(bias_tensor, tensor_name).front();
        auto stages = CreateStages(
            {lhs.as_tensor_ref(), rhs.as_tensor_ref(), bias_tensor});
        stages->InsertLazily(out);
        std::vector<CINNValue> res{CINNValue(out), CINNValue(stages)};
        *ret = CINNValuePack{res};
      });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(gemm_compute,
                    GetInjectiveScheduleFunc(output_shapes, target),
                    "strategy.cublas.gemm",
                    1);

  return strategy;
}

std::vector<shape_t> InferShapeForCublasGemm(
    const std::vector<std::vector<int>> &input_shapes,
    const framework::AttrMapType &attrs) {
  CHECK_EQ(input_shapes.size(), 3UL)
      << "cublas_gemm should have 3 input shapes";
  CHECK_EQ(input_shapes[0].size(), input_shapes[1].size());
  CHECK_EQ(input_shapes[0].size(), input_shapes[2].size());
  CHECK((input_shapes[0].size() == 2 || input_shapes[0].size() == 3));
  return {input_shapes[2]};
}

std::vector<Type> InferDtypeForCublasGemm(const std::vector<Type> &inputs_type,
                                          const framework::AttrMapType &attrs) {
  CHECK_EQ(inputs_type.size(), 3UL)
      << "The input's type size is 0! Please check again.";
  CHECK_EQ(inputs_type[0], inputs_type[1])
      << "The input A and B's types should be equal! Please check again.";
  CHECK_EQ(inputs_type[0], inputs_type[2])
      << "The input A and C's types should be equal! Please check again.";
  return {inputs_type[0]};
}

std::shared_ptr<OpStrategy> StrategyForLayoutTransform(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  framework::CINNCompute layout_transform_compute([=](lang::Args args,
                                                      lang::RetValue *ret) {
    std::string src_layout;
    std::string dst_layout;
    if (attrs.attr_store.find("src_layout") != attrs.attr_store.end()) {
      src_layout = absl::get<std::string>(attrs.attr_store.at("src_layout"));
    }
    if (attrs.attr_store.find("dst_layout") != attrs.attr_store.end()) {
      dst_layout = absl::get<std::string>(attrs.attr_store.at("dst_layout"));
    }
    CHECK(!args.empty()) << "The input argument of layout_transform compute is "
                            "empty! Please check.\n";
    CINNValuePack input_args = args[0];
    CHECK(!input_args.empty())
        << "at least one input tensor for layout_transform compute\n";
    Expr A = input_args[0];
    CHECK(A.as_tensor());

    std::string tensor_name = UniqName("layout_transform_output");
    if (FLAGS_cinn_ir_schedule) {
      CHECK_EQ(input_args.size(), 2);
      CHECK(input_args[1].is_string());
      tensor_name = input_args[1].operator std::string();
    }

    auto out = pe::LayoutTransform(
        A.as_tensor_ref(), src_layout, dst_layout, tensor_name);
    auto stages = CreateStages({A.as_tensor_ref()});
    std::vector<CINNValue> res;
    stages->InsertLazily(out);
    res = {CINNValue(out), CINNValue(stages)};
    *ret = CINNValuePack{res};
  });

  framework::CINNSchedule layout_transform_schedule(
      [=](lang::Args args, lang::RetValue *ret) {
        if (FLAGS_cinn_ir_schedule) {
          CHECK(!args.empty()) << "The input argument of CublasGemm schedule "
                                  "is empty! Please check.";
          CINNValuePack arg_pack = args[0];
          std::vector<Expr> vec_ast;
          for (int i = 0; i < arg_pack.size(); i++) {
            if (arg_pack[i].is_expr()) {
              Expr temp = arg_pack[i];
              vec_ast.emplace_back(temp);
            }
          }
          CHECK(!vec_ast.empty());
          ir::ModuleExpr mod_expr(vec_ast);
          ir::IRSchedule ir_sch(mod_expr);
          ir_sch.MergeExprs();

          if (target.arch == Target::Arch::X86) {
            pe::IRScheduleInjectiveCPU(ir_sch, output_shapes.front(), target);
          } else {
            CINN_NOT_IMPLEMENTED
          }
          std::vector<CINNValue> res{
              CINNValue(ir_sch.GetModule().GetExprs().at(0))};
          *ret = CINNValuePack{res};
        } else {
          CHECK(!args.empty()) << "The input argument of layout_transform "
                                  "schedule is empty! Please check.\n";
          CINNValuePack arg_pack = args[0];
          CHECK_EQ(arg_pack.size(), 2UL);
          Expr out = arg_pack[0];
          poly::StageMap stages = arg_pack[1];
          CHECK(out.as_tensor());
          auto tensor_out = out.as_tensor_ref();
          std::vector<int> out_shape;
          for (auto shape : tensor_out->shape) {
            out_shape.push_back(shape.as_int32());
          }
          if (target.arch == Target::Arch::X86) {
            pe::ScheduleInjectiveCPU(stages[tensor_out], out_shape, target);
          } else {
            CINN_NOT_IMPLEMENTED
          }
          *ret = arg_pack;
        }
      });

  auto strategy = std::make_shared<framework::OpStrategy>();
  CHECK(out_type.size())
      << "Out_type of layout_transform op is empty! Please check.";
  strategy->AddImpl(layout_transform_compute,
                    layout_transform_schedule,
                    "strategy.layout_transform.x86",
                    1);
  return strategy;
}

std::vector<shape_t> InferShapeForLayoutTransform(
    const std::vector<shape_t> &inputs_shape,
    const framework::AttrMapType &attrs) {
  std::string src_layout;
  std::string dst_layout;
  if (attrs.find("src_layout") != attrs.end()) {
    src_layout = absl::get<std::string>(attrs.at("src_layout"));
  }
  if (attrs.find("dst_layout") != attrs.end()) {
    dst_layout = absl::get<std::string>(attrs.at("dst_layout"));
  }
  CHECK_EQ(inputs_shape.size(), 1UL);

  std::vector<Expr> input_shapes_expr;
  for (int shape : inputs_shape[0]) {
    input_shapes_expr.push_back(Expr(shape));
  }
  absl::flat_hash_map<int, std::vector<int>> split_index_map;
  std::vector<Expr> out_shapes =
      pe::InferShapeLayoutTransform(input_shapes_expr,
                                    ir::Layout(src_layout),
                                    ir::Layout(dst_layout),
                                    &split_index_map);
  VLOG(4) << "out_shapes: " << out_shapes;
  std::vector<int> output_shapes;
  for (auto &shape : out_shapes) {
    output_shapes.push_back(shape.as_int32());
  }
  return {output_shapes};
}

std::vector<Type> InferDtypeForLayoutTransform(
    const std::vector<Type> &inputs_type, const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";
  std::vector<Type> res{inputs_type[0]};
  return res;
}

std::shared_ptr<OpStrategy> StrategyForReverse(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  // check output shape
  CHECK(!output_shapes.empty() && !output_shapes[0].empty())
      << "Output shape is empty! Please check.\n";
  // get axis[0, n_dim)
  std::vector<int> axis;
  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
    axis = absl::get<std::vector<int>>(attrs.attr_store.at("axis"));
    for (auto &e : axis) {
      if (e >= static_cast<int>(output_shapes[0].size()) ||
          e < -1 * static_cast<int>(output_shapes[0].size())) {
        LOG(FATAL) << "axis is not in [0, n_dim), Please check.";
      }
      if (e < 0) {
        e += output_shapes[0].size();
      }
    }
  }

  framework::CINNCompute reverse_compute([=](lang::Args args,
                                             lang::RetValue *ret) {
    CHECK(!args.empty())
        << "The input argument of reverse compute is empty! Please check.\n";
    CINNValuePack input_args = args[0];
    CHECK(!input_args.empty())
        << "at least one input tensor for reverse compute\n";
    Expr A = input_args[0];
    CHECK(A.as_tensor());

    std::string tensor_name = UniqName("Reverse_output");
    if (FLAGS_cinn_ir_schedule) {
      CHECK_EQ(input_args.size(), 2);
      CHECK(input_args[1].is_string());
      tensor_name = input_args[1].operator std::string();
    }

    auto out = pe::Reverse(A.as_tensor_ref(), axis, tensor_name);
    auto stages = CreateStages({A.as_tensor_ref(), out});
    *ret = CINNValuePack{{CINNValue(out), CINNValue(stages)}};
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  CHECK(out_type.size()) << "Out_type of reverse op is empty! Please check.";
  strategy->AddImpl(reverse_compute,
                    GetInjectiveScheduleFunc(output_shapes, target),
                    "strategy.reverse.x86",
                    1);
  return strategy;
}

std::vector<framework::shape_t> InferShapeForReverse(
    const std::vector<framework::shape_t> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK(!inputs_shape.empty() && !inputs_shape[0].empty())
      << "The input's shape size is 0! Please check again.";
  std::vector<framework::shape_t> res{inputs_shape[0]};
  if (attrs.find("axis") != attrs.end()) {
    auto axis = absl::get<std::vector<int>>(attrs.at("axis"));
    for (auto &e : axis) {
      if (e >= static_cast<int>(inputs_shape[0].size()) ||
          e < -1 * static_cast<int>(inputs_shape[0].size())) {
        LOG(FATAL) << "axis is not in [-n_dim, n_dim), Please check.";
      }
      if (e < 0) {
        e += inputs_shape[0].size();
      }
    }
  }
  return res;
}

std::vector<std::vector<std::string>> InferLayoutForReverse(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
    auto axis = absl::get<std::vector<int>>(attrs.attr_store.at("axis"));
    for (auto &e : axis) {
      if (e >= static_cast<int>(input_shapes[0].size()) ||
          e < -1 * static_cast<int>(input_shapes[0].size())) {
        LOG(FATAL) << "axis is not in [-n_dim, n_dim), Please check.";
      }
    }
  }
  CHECK_EQ(input_layouts.size(), 1U)
      << "The input's layout size is not 1! Please check again.";
  return {input_layouts, input_layouts};
}

std::vector<std::vector<std::string>> InferLayoutForLayoutTransform(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_EQ(input_layouts.size(), 1U)
      << "The input's layouts size is not 1! Please check again.";
  std::string dst_layout;
  std::string src_layout;
  if (attrs.attr_store.find("dst_layout") != attrs.attr_store.end()) {
    dst_layout = absl::get<std::string>(attrs.attr_store.at("dst_layout"));
  }
  if (attrs.attr_store.find("src_layout") != attrs.attr_store.end()) {
    src_layout = absl::get<std::string>(attrs.attr_store.at("src_layout"));
  }
  return {{dst_layout}, {src_layout}};
}

std::shared_ptr<OpStrategy> StrategyForTranspose(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  // check output shape
  CHECK(!output_shapes.empty() && !output_shapes[0].empty())
      << "Output shape is empty! Please check.\n";

  std::vector<int> axis;
  auto input_shape = inputs[0]->shape;
  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
    axis = absl::get<std::vector<int>>(attrs.attr_store.at("axis"));
    CHECK_EQ(axis.size(), output_shapes[0].size())
        << "axis size is not equal output_shapes size! Please check setting.\n";
    // check axis and shape
    for (int idx = 0; idx < axis.size(); ++idx) {
      CHECK(axis[idx] >= 0 && axis[idx] < axis.size());
      for (int idy = idx + 1; idy < axis.size(); ++idy) {
        CHECK_NE(axis[idx], axis[idy]) << "axis can't repeat!";
      }
      CHECK_EQ(output_shapes[0][idx], input_shape[axis[idx]].as_int32())
          << "output shape is not equal! Please check!\n";
    }
  } else {
    LOG(FATAL) << "axis is not be set! Please check.";
  }

  framework::CINNCompute transpose_compute([=](lang::Args args,
                                               lang::RetValue *ret) {
    CHECK(!args.empty())
        << "The input argument of transpose compute is empty! Please check.\n";
    CINNValuePack input_args = args[0];
    CHECK(!input_args.empty())
        << "at least one input tensor for transpose compute\n";
    Expr A = input_args[0];
    CHECK(A.as_tensor());
    std::string tensor_name = UniqName("Transpose_output");
    if (FLAGS_cinn_ir_schedule) {
      CHECK_EQ(input_args.size(), 2);
      CHECK(input_args[1].is_string());
      tensor_name = input_args[1].operator std::string();
    }

    auto out = pe::Transpose(A.as_tensor_ref(), axis, tensor_name);
    auto stages = CreateStages({out});
    *ret = CINNValuePack{{CINNValue(out), CINNValue(stages)}};
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(transpose_compute,
                    GetInjectiveScheduleFunc(output_shapes, target),
                    "strategy.transpose.x86",
                    1);
  return strategy;
}

std::vector<framework::shape_t> InferShapeForTranspose(
    const std::vector<framework::shape_t> &inputs_shape,
    const framework::AttrMapType &attrs) {
  std::vector<framework::shape_t> result;
  CHECK(!inputs_shape.empty() && !inputs_shape[0].empty())
      << "The input's shape size is 0! Please check again.";
  if (attrs.find("axis") != attrs.end()) {
    auto axis = absl::get<std::vector<int>>(attrs.at("axis"));
    CHECK_EQ(axis.size(), inputs_shape[0].size())
        << "input size and axis size is not equal!";
    std::vector<int> output_shape;
    for (int idx = 0; idx < axis.size(); ++idx) {
      CHECK(axis[idx] >= 0 && axis[idx] < axis.size());
      for (int idy = idx + 1; idy < axis.size(); ++idy) {
        CHECK_NE(axis[idx], axis[idy]) << "axis can't repeat!";
      }
      output_shape.push_back(inputs_shape[0][axis[idx]]);
    }
    result.push_back(output_shape);
  } else {
    LOG(FATAL) << "axis is not be set! Please check.";
  }
  return result;
}

std::vector<std::vector<std::string>> InferLayoutForTranspose(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_EQ(input_shapes.size(), 1U)
      << "The input's shape size is not 1! Please check again.";
  CHECK_EQ(input_layouts.size(), 1U)
      << "The input's layout size is not 1! Please check again.";

  std::vector<int> axis;
  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
    axis = absl::get<std::vector<int>>(attrs.attr_store.at("axis"));
    for (int idx = 0; idx < axis.size(); ++idx) {
      CHECK(axis[idx] >= 0 && axis[idx] < axis.size());
      for (int idy = idx + 1; idy < axis.size(); ++idy) {
        CHECK_NE(axis[idx], axis[idy]) << "axis can't repeat!";
      }
    }
  } else {
    LOG(FATAL) << "axis is not be set! Please check.";
  }

  std::vector<std::string> new_input_layouts = input_layouts;
  for (int i = 0; i < input_shapes.size(); i++) {
    if (input_shapes[i].size() > 4) {
      // alter input layout back
      new_input_layouts[i] = input_layouts[0].substr(0, 4);
    }
  }

  std::string output_layout = new_input_layouts[0];
  for (int idx = 0; idx < axis.size(); ++idx) {
    output_layout[idx] = new_input_layouts[0][axis[idx]];
  }

  return {{output_layout}, new_input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForGather(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  CHECK(!output_shapes.empty() && !output_shapes[0].empty())
      << "The shape of output is empty! Please check again.";
  VLOG(4) << "The output passed in StrategyForGather: "
          << utils::Join(output_shapes[0], ", ");
  CHECK(!out_type.empty())
      << "The output type of Gather is empty! Please check again.\n";

  int axis = 0;
  if (attrs.attr_store.contains("axis")) {
    axis = absl::get<int>(attrs.attr_store.at("axis"));
  }
  axis = axis < 0 ? axis + static_cast<int>(inputs[0]->shape.size()) : axis;

  std::vector<Expr> output_shape;
  output_shape.reserve(output_shapes[0].size());
  for (int i : output_shapes[0]) {
    output_shape.emplace_back(i);
  }

  framework::CINNCompute gather_compute{
      [axis, output_shape = std::move(output_shape)](lang::Args args,
                                                     lang::RetValue *ret) {
        VLOG(4) << "The axis value used in gather_compute: " << axis;
        CHECK(!args.empty()) << "The input args are empty! Please check again.";
        CINNValuePack input_args = args[0];
        int input_size = input_args.size();
        CHECK_GE(input_size, 2U)
            << "Require 2 input tensors for Gather compute.";
        Expr x = input_args[0];
        CHECK(x.as_tensor());
        Expr index = input_args[1];
        CHECK(index.as_tensor());

        std::string tensor_name = UniqName("gather_output");
        if (FLAGS_cinn_ir_schedule) {
          CHECK_EQ(input_args.size(), 3U);
          CHECK(input_args[2].is_string());
          tensor_name = input_args[2].operator std::string();
        }

        auto out = pe::Gather(x.as_tensor_ref(),
                              index.as_tensor_ref(),
                              output_shape,
                              axis,
                              tensor_name);
        auto stages = CreateStages({x.as_tensor_ref(), index.as_tensor_ref()});
        stages->InsertLazily(out);
        std::vector<CINNValue> res{CINNValue(out), CINNValue(stages)};
        *ret = CINNValuePack{res};
      }};

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(gather_compute,
                    GetInjectiveScheduleFunc(output_shapes, target),
                    "strategy.gather.x86",
                    1);
  return strategy;
}

std::vector<std::vector<int>> InferShapeForGather(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK_EQ(inputs_shape.size(), 2U)
      << "The inputs' shape size should be equal to 2! Please check again.";
  std::vector<int> x_shape = inputs_shape[0];
  std::vector<int> index_shape = inputs_shape[1];
  int axis = absl::get<int>(attrs.at("axis"));
  VLOG(4) << "The axis value used in Gather: " << axis;

  CHECK(axis >= 0 && axis < static_cast<int>(x_shape.size()))
      << "The attribute `axis` in Gather should be >= 0 and < the size of the "
         "first input shape! Please check again.";

  std::vector<int> output_shape = x_shape;
  output_shape[axis] = index_shape[axis];
  VLOG(4) << "The output shape of gather: " << utils::Join(output_shape, ", ");

  return {std::move(output_shape)};
}

std::vector<Type> InferDtypeForGather(const std::vector<Type> &inputs_type,
                                      const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";
  return {inputs_type[0]};
}

std::vector<std::vector<std::string>> InferLayoutForGather(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_EQ(input_layouts.size(), 2U)
      << "The input's layout size is not equal to 2! Please check again.";
  return {{input_layouts[0]}, input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForScatterAssign(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  int axis = 0;
  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
    axis = absl::get<int>(attrs.attr_store.at("axis"));
  }

  framework::CINNCompute scatter_assign_compute([=](lang::Args args,
                                                    lang::RetValue *ret) {
    CHECK(!args.empty()) << "The input arguments of ScatterAssign compute is "
                            "empty! Please check.\n";
    CINNValuePack arg_pack = args[0];
    int input_size = arg_pack.size();
    CHECK_GE(input_size, 3U)
        << "at least 3 input tensors for ScatterAssign compute\n";
    CHECK(!output_shapes.empty());

    Expr expr_input = arg_pack[0];
    CHECK(expr_input.as_tensor());
    auto tensor_input = expr_input.as_tensor_ref();

    Expr expr_updates = arg_pack[1];
    CHECK(expr_updates.as_tensor());
    auto tensor_updates = expr_updates.as_tensor_ref();

    Expr expr_index = arg_pack[2];
    CHECK(expr_index.as_tensor());
    auto tensor_index = expr_index.as_tensor_ref();

    auto stages = CreateStages({tensor_input, tensor_updates, tensor_index});

    std::string tensor_name = UniqName("scatter_assign_output");
    if (FLAGS_cinn_ir_schedule) {
      CHECK_EQ(arg_pack.size(), 4U);
      CHECK(arg_pack[3].is_string());
      tensor_name = arg_pack[3].operator std::string();
    }

    auto out = pe::ScatterAssign(
        tensor_input, tensor_updates, tensor_index, target, axis, tensor_name);

    std::vector<CINNValue> res;
    stages->InsertLazily(out);
    res.push_back(CINNValue(out));
    CHECK(!out_type.empty())
        << "Output type of ScatterAssign is empty! Please check.\n";
    res.push_back(CINNValue(stages));
    *ret = CINNValuePack{res};
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(scatter_assign_compute,
                    GetInjectiveScheduleFunc(output_shapes, target, false),
                    "strategy.scatter_assign.x86",
                    1);
  return strategy;
}

std::vector<std::vector<int>> InferShapeForScatterAssign(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK_GE(inputs_shape.size(), 3U)
      << "The input's shape size should be no less than 3! Please check again.";

  const auto &input_shape = inputs_shape[0];
  const auto &assign_shape = inputs_shape[1];
  const auto &index_shape = inputs_shape[2];

  int axis = 0;
  if (attrs.find("axis") != attrs.end()) {
    axis = absl::get<int>(attrs.at("axis"));
  }

  if (axis < 0) axis += input_shape.size();

  CHECK(axis >= 0 && axis < input_shape.size())
      << "In ScatterAssign op, the attribute `axis` should be >= 0 and < input "
         "shape's size! Please check.";
  CHECK_EQ(index_shape.size(), 1U)
      << "Dimensions of index tensor in ScatterAssign should be 1! Please "
         "check.";
  CHECK_EQ(input_shape.size(), assign_shape.size())
      << "Dimensions of inputs A and B in ScatterAssign should be equal! "
         "Please check.";
  CHECK_EQ(assign_shape[axis], index_shape[0])
      << "The first dimension of input B and index tensor in ScatterAssign "
         "should be equal! Please check.";
  for (int i = 0; i < input_shape.size(); ++i) {
    if (i != axis) {
      CHECK_EQ(input_shape[i], assign_shape[i])
          << "The " << i
          << "-th dimension of input A and B in ScatterAssign should be equal! "
             "Please check.";
    }
  }

  VLOG(4) << "Each input tensor's shape of ScatterAssign: A("
          << cinn::utils::Join(input_shape, ",") << "), B("
          << cinn::utils::Join(assign_shape, ",") << "), index("
          << cinn::utils::Join(index_shape, ",") << ")"
          << " at axis (" << axis << ")";

  return {input_shape};
}

std::vector<Type> InferDtypeForScatterAssign(
    const std::vector<Type> &inputs_type, const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";
  std::vector<Type> res{inputs_type[0]};
  return res;
}

std::vector<std::vector<std::string>> InferLayoutForScatterAssign(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_GE(input_layouts.size(), 3U)
      << "The input's layout size is less than 3! Please check again.";
  return {{input_layouts[0]}, input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForScatterAdd(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  int axis = 0;
  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
    axis = absl::get<int>(attrs.attr_store.at("axis"));
  }

  framework::CINNCompute scatter_add_compute([=](lang::Args args,
                                                 lang::RetValue *ret) {
    CHECK(!args.empty()) << "The input arguments of ScatterAdd compute is "
                            "empty! Please check.\n";
    CINNValuePack arg_pack = args[0];
    int input_size = arg_pack.size();
    CHECK_GE(input_size, 3U)
        << "at least 3 input tensors for ScatterAdd compute\n";
    CHECK(!output_shapes.empty());

    Expr expr_input = arg_pack[0];
    CHECK(expr_input.as_tensor());
    auto tensor_input = expr_input.as_tensor_ref();

    Expr expr_updates = arg_pack[1];
    CHECK(expr_updates.as_tensor());
    auto tensor_updates = expr_updates.as_tensor_ref();

    Expr expr_index = arg_pack[2];
    CHECK(expr_index.as_tensor());
    auto tensor_index = expr_index.as_tensor_ref();

    auto stages = CreateStages({tensor_input, tensor_updates, tensor_index});

    std::string tensor_name = UniqName("scatter_add_output");
    if (FLAGS_cinn_ir_schedule) {
      CHECK_EQ(arg_pack.size(), 4U);
      CHECK(arg_pack[3].is_string());
      tensor_name = arg_pack[3].operator std::string();
    }

    auto out = pe::ScatterAdd(
        tensor_input, tensor_updates, tensor_index, target, axis, tensor_name);

    std::vector<CINNValue> res;
    stages->InsertLazily(out);
    res.push_back(CINNValue(out));
    CHECK(!out_type.empty())
        << "Output type of ScatterAdd is empty! Please check.\n";
    res.push_back(CINNValue(stages));
    *ret = CINNValuePack{res};
  });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(scatter_add_compute,
                    GetInjectiveScheduleFunc(output_shapes, target, false),
                    "strategy.scatter_add.x86",
                    1);
  return strategy;
}

std::vector<std::vector<int>> InferShapeForScatterAdd(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK_GE(inputs_shape.size(), 3U)
      << "The input's shape size should be no less than 3! Please check again.";

  const auto &input_shape = inputs_shape[0];
  const auto &updates_shape = inputs_shape[1];
  const auto &index_shape = inputs_shape[2];

  int axis = 0;
  if (attrs.find("axis") != attrs.end()) {
    axis = absl::get<int>(attrs.at("axis"));
  }

  if (axis < 0) axis += input_shape.size();

  CHECK(axis >= 0 && axis < input_shape.size())
      << "In ScatterAdd op, the attribute `axis` should be >= 0 and < input "
         "shape's size! Please check.";
  CHECK_EQ(index_shape.size(), 1U)
      << "Dimensions of index tensor in ScatterAdd should be 1! Please check.";
  CHECK_EQ(input_shape.size(), updates_shape.size())
      << "Dimensions of inputs A and B in ScatterAdd should be equal! Please "
         "check.";
  CHECK_EQ(updates_shape[axis], index_shape[0])
      << "The first dimension of input B and index tensor in ScatterAdd should "
         "be equal! Please check.";
  for (int i = 0; i < input_shape.size(); ++i) {
    if (i != axis) {
      CHECK_EQ(input_shape[i], updates_shape[i])
          << "The " << i
          << "-th dimension of input A and B in ScatterAdd should be equal! "
             "Please check.";
    }
  }

  VLOG(4) << "Each input tensor's shape of ScatterAdd: A("
          << cinn::utils::Join(input_shape, ",") << "), B("
          << cinn::utils::Join(updates_shape, ",") << "), index("
          << cinn::utils::Join(index_shape, ",") << ")"
          << " at axis (" << axis << ")";

  return {input_shape};
}

std::vector<Type> InferDtypeForScatterAdd(const std::vector<Type> &inputs_type,
                                          const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";
  std::vector<Type> res{inputs_type[0]};
  return res;
}

std::vector<std::vector<std::string>> InferLayoutForScatterAdd(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_GE(input_layouts.size(), 3U)
      << "The input's layout size is less than 3! Please check again.";
  return {{input_layouts[0]}, input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForSlice(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  std::vector<int> starts, ends, axes, strides, decrease_axis;
  if (attrs.attr_store.find("starts") != attrs.attr_store.end()) {
    starts = absl::get<std::vector<int>>(attrs.attr_store.at("starts"));
  }
  if (attrs.attr_store.find("ends") != attrs.attr_store.end()) {
    ends = absl::get<std::vector<int>>(attrs.attr_store.at("ends"));
  }
  if (attrs.attr_store.find("axes") != attrs.attr_store.end()) {
    axes = absl::get<std::vector<int>>(attrs.attr_store.at("axes"));
  }
  if (attrs.attr_store.find("strides") != attrs.attr_store.end()) {
    strides = absl::get<std::vector<int>>(attrs.attr_store.at("strides"));
  }
  if (attrs.attr_store.find("decrease_axis") != attrs.attr_store.end()) {
    decrease_axis =
        absl::get<std::vector<int>>(attrs.attr_store.at("decrease_axis"));
  }

  CHECK(!starts.empty()) << "The Slice op doesn't find [starts] attrbute! It "
                            "it a mandatory attribute, please check.";
  CHECK(!ends.empty()) << "The Slice op doesn't find [ends] attrbute! It it a "
                          "mandatory attribute, please check.";
  CHECK_EQ(starts.size(), ends.size())
      << "The size of [starts] and [ends] must be identical! Please check.";
  if (!axes.empty()) {
    CHECK_EQ(starts.size(), axes.size())
        << "The size of [starts] and [axes] must be identical! Please check.";
  } else {
    for (int i = 0; i < starts.size(); i++) {
      axes.push_back(i);
    }
  }
  if (!strides.empty()) {
    CHECK_EQ(starts.size(), strides.size())
        << "The size of [starts] and [strides] must be identical! Please "
           "check.";
  } else {
    for (int i = 0; i < starts.size(); i++) {
      strides.push_back(1);
    }
  }

  std::vector<Expr> output_shape;
  for (auto &i : output_shapes[0]) {
    output_shape.push_back(Expr(i));
  }

  framework::CINNCompute slice_compute(
      [=](lang::Args args, lang::RetValue *ret) {
        CHECK(!args.empty())
            << "The input arguments of slice compute is empty! Please check.";
        CINNValuePack arg_pack = args[0];
        CHECK(!arg_pack.empty())
            << "The input tensors of slice compute is empty! Please check.";
        Expr A_expr = arg_pack[0];
        CHECK(A_expr.as_tensor());
        ir::Tensor A = A_expr.as_tensor_ref();

        std::string tensor_name = UniqName("Slice_output");
        if (FLAGS_cinn_ir_schedule) {
          CHECK_EQ(arg_pack.size(), 2U);
          CHECK(arg_pack[1].is_string());
          tensor_name = arg_pack[1].operator std::string();
        }

        auto out = pe::Slice(
            A, starts, axes, strides, decrease_axis, output_shape, tensor_name);
        auto stages = CreateStages({out});
        *ret = CINNValuePack{{CINNValue(out), CINNValue(stages)}};
      });

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(slice_compute,
                    GetInjectiveScheduleFunc(output_shapes, target),
                    "strategy.slice.x86",
                    1);

  return strategy;
}

std::vector<std::vector<int>> InferShapeForSlice(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK(!inputs_shape.empty() && !inputs_shape[0].empty())
      << "The input's shape size is 0! Please check again.";
  std::vector<int> starts, ends, axes, strides, decrease_axis, infer_flags;
  for (auto &iter : attrs) {
    if (iter.first == "starts") {
      starts = absl::get<std::vector<int>>(iter.second);
    } else if (iter.first == "ends") {
      ends = absl::get<std::vector<int>>(iter.second);
    } else if (iter.first == "axes") {
      axes = absl::get<std::vector<int>>(iter.second);
    } else if (iter.first == "strides") {
      strides = absl::get<std::vector<int>>(iter.second);
    } else if (iter.first == "decrease_axis") {
      decrease_axis = absl::get<std::vector<int>>(iter.second);
    } else if (iter.first == "infer_flags") {
      infer_flags = absl::get<std::vector<int>>(iter.second);
    } else {
      LOG(ERROR) << "Unsupported attr: " << iter.first << std::endl;
    }
  }
  CHECK(!starts.empty()) << "The Slice op doesn't find [starts] attrbute! It "
                            "it a mandatory attribute, please check.";
  CHECK(!ends.empty()) << "The Slice op doesn't find [ends] attrbute! It it a "
                          "mandatory attribute, please check.";
  CHECK_EQ(starts.size(), ends.size())
      << "The size of [starts] and [ends] must be identical! Please check.";
  if (!axes.empty()) {
    CHECK_EQ(starts.size(), axes.size())
        << "The size of [starts] and [axes] must be identical! Please check.";
  } else {
    for (int i = 0; i < starts.size(); i++) {
      axes.push_back(i);
    }
  }
  if (!strides.empty()) {
    CHECK_EQ(starts.size(), strides.size())
        << "The size of [starts] and [strides] must be identical! Please "
           "check.";
  } else {
    for (int i = 0; i < starts.size(); i++) {
      strides.push_back(1);
    }
  }

  std::vector<int> output_shape = inputs_shape[0];
  for (int i = 0; i < axes.size(); i++) {
    if (ends[i] < 0) {
      ends[i] = output_shape[axes[i]] + ends[i];
    } else if (ends[i] > output_shape[axes[i]]) {
      ends[i] = output_shape[axes[i]];
    }
    if (starts[i] < -output_shape[axes[i]]) {
      starts[i] = 0;
    } else if (starts[i] < 0) {
      starts[i] = output_shape[axes[i]] + starts[i];
    } else if (starts[i] > output_shape[axes[i]]) {
      starts[i] = output_shape[axes[i]] - 1;
    }

    CHECK_NE(strides[i], 0)
        << "The value of [strides] of slice should not be 0 ! Please Check.";
    if (strides[i] > 0) {
      CHECK(ends[i] > starts[i])
          << "[ends] should greater than [starts] when strides > 0 ! But here "
          << ends[i] << " < " << starts[i] << ", Please Check.";
      output_shape[axes[i]] =
          (ends[i] - starts[i] + strides[i] - 1) / strides[i];
    } else {
      CHECK(ends[i] < starts[i])
          << "[ends] should less than [starts] when strides < 0 ! But here "
          << ends[i] << " > " << starts[i] << ",  Please Check.";
      output_shape[axes[i]] =
          (starts[i] - ends[i] + (-strides[i]) - 1) / (-strides[i]);
    }
  }

  if (decrease_axis.size() > 0) {
    std::vector<int> new_shape;
    for (int i = 0; i < output_shape.size(); ++i) {
      if (std::find(decrease_axis.cbegin(), decrease_axis.cend(), i) !=
          decrease_axis.cend()) {
        CHECK_EQ(output_shape[i], 1)
            << "Decrease dim should be 1, but now received " << output_shape[i];
      } else {
        new_shape.emplace_back(output_shape[i]);
      }
    }
    if (new_shape.empty()) {
      // output shape should not empty
      new_shape = {1};
    }

    output_shape = new_shape;
  }

  VLOG(4) << "Output shape of Slice is: "
          << cinn::utils::Join(output_shape, ",");
  std::vector<std::vector<int>> res{output_shape};
  return res;
}

std::vector<Type> InferDtypeForSlice(const std::vector<Type> &inputs_type,
                                     const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";
  std::vector<Type> res{inputs_type[0]};
  return res;
}

std::vector<std::vector<std::string>> InferLayoutForSlice(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_EQ(input_layouts.size(), 1U)
      << "The input's layout size is not 1! Please check again.";
  CHECK_EQ(input_shapes.size(), 1U)
      << "The input's shape size is not 1! Please check again.";
  std::vector<int> starts;
  std::vector<int> ends;
  std::vector<int> axes;
  for (auto &iter : attrs.attr_store) {
    if (iter.first == "starts") {
      starts = absl::get<std::vector<int>>(iter.second);
    } else if (iter.first == "ends") {
      ends = absl::get<std::vector<int>>(iter.second);
    } else if (iter.first == "axes") {
      axes = absl::get<std::vector<int>>(iter.second);
    }
  }
  std::string new_input_layouts = input_layouts[0];
  bool trans_back = false;
  if (input_shapes[0].size() > 4) {
    for (int i = 0; i < axes.size(); i++) {
      if (axes[i] == 1) {
        trans_back = true;
        break;
      }
    }
  }
  if (trans_back) {
    return {{"NCHW"}, {"NCHW"}};
  }
  return {input_layouts, input_layouts};
}

std::shared_ptr<OpStrategy> StrategyForSliceAssign(
    const framework::NodeAttr &attrs,
    const std::vector<ir::Tensor> &inputs,
    const std::vector<Type> &out_type,
    const std::vector<std::vector<int>> &output_shapes,
    const Target &target) {
  CHECK_EQ(inputs.size(), 2)
      << "the number of input tensors must be equal to 2";
  CHECK(!output_shapes.empty() && !output_shapes[0].empty())
      << "The shape of output is empty! Please check again.";
  VLOG(4) << "The output passed in StrategyForSliceAssign: "
          << utils::Join(output_shapes[0], ", ");
  CHECK(!out_type.empty())
      << "The output type of SliceAssign is empty! Please check again.\n";

  std::vector<int> starts, ends, axes, strides;
  if (attrs.attr_store.find("starts") != attrs.attr_store.end()) {
    starts = absl::get<std::vector<int>>(attrs.attr_store.at("starts"));
  }
  if (attrs.attr_store.find("ends") != attrs.attr_store.end()) {
    ends = absl::get<std::vector<int>>(attrs.attr_store.at("ends"));
  }
  if (attrs.attr_store.find("axes") != attrs.attr_store.end()) {
    axes = absl::get<std::vector<int>>(attrs.attr_store.at("axes"));
  }
  if (attrs.attr_store.find("strides") != attrs.attr_store.end()) {
    strides = absl::get<std::vector<int>>(attrs.attr_store.at("strides"));
  }

  CHECK(!starts.empty())
      << "The SliceAssign op doesn't find [starts] attrbute! It it a mandatory "
         "attribute, please check.";
  CHECK(!ends.empty()) << "The SliceAssign op doesn't find [ends] attrbute! It "
                          "it a mandatory attribute, please check.";
  CHECK_EQ(starts.size(), ends.size())
      << "The size of [starts] and [ends] must be identical! Please check.";
  if (!axes.empty()) {
    CHECK_EQ(starts.size(), axes.size())
        << "The size of [starts] and [axes] must be identical! Please check.";
  } else {
    for (int i = 0; i < starts.size(); i++) {
      axes.push_back(i);
    }
  }
  if (!strides.empty()) {
    CHECK_EQ(starts.size(), strides.size())
        << "The size of [starts] and [strides] must be identical! Please "
           "check.";
  } else {
    for (int i = 0; i < starts.size(); i++) {
      strides.push_back(1);
    }
  }

  framework::CINNCompute slice_assign_compute{
      [=](lang::Args args, lang::RetValue *ret) {
        CHECK(!args.empty()) << "The input args are empty! Please check again.";
        CINNValuePack arg_pack = args[0];
        int input_size = arg_pack.size();
        CHECK_GE(input_size, 2U)
            << "Require 2 input tensors for SliceAssign compute.";
        Expr input = arg_pack[0];
        CHECK(input.as_tensor());
        Expr assign = arg_pack[1];
        CHECK(assign.as_tensor());

        std::string tensor_name = UniqName("slice_assign_output");
        if (FLAGS_cinn_ir_schedule) {
          CHECK_EQ(arg_pack.size(), 3U);
          CHECK(arg_pack[2].is_string());
          tensor_name = arg_pack[2].operator std::string();
        }

        auto out = pe::SliceAssign(input.as_tensor_ref(),
                                   assign.as_tensor_ref(),
                                   axes,
                                   starts,
                                   ends,
                                   strides,
                                   tensor_name);
        auto stages = CreateStages({out});
        std::vector<CINNValue> res{CINNValue(out), CINNValue(stages)};
        *ret = CINNValuePack{res};
      }};

  auto strategy = std::make_shared<framework::OpStrategy>();
  strategy->AddImpl(slice_assign_compute,
                    GetInjectiveScheduleFunc(output_shapes, target),
                    "strategy.slice_assign.x86",
                    1);
  return strategy;
}

std::vector<std::vector<int>> InferShapeForSliceAssign(
    const std::vector<std::vector<int>> &inputs_shape,
    const framework::AttrMapType &attrs) {
  CHECK_EQ(inputs_shape.size(), 2U)
      << "The inputs' shape size should be equal to 2! Please check again.";
  return {inputs_shape[0]};
}

std::vector<Type> InferDtypeForSliceAssign(
    const std::vector<Type> &inputs_type, const framework::AttrMapType &attrs) {
  CHECK(!inputs_type.empty())
      << "The input's type size is 0! Please check again.";
  return {inputs_type[0]};
}

std::vector<std::vector<std::string>> InferLayoutForSliceAssign(
    const std::vector<framework::shape_t> &input_shapes,
    const std::vector<std::string> &input_layouts,
    const framework::NodeAttr &attrs,
    const Target &target) {
  CHECK_EQ(input_layouts.size(), 2U)
      << "The input's layout size is not equal to 2! Please check again.";
  return {{input_layouts[0]}, {""}};
}

}  // namespace op
}  // namespace hlir
}  // namespace cinn

CINN_REGISTER_HELPER(transform_ops) {
  CINN_REGISTER_OP(matmul)
      .describe(
          "This operator is used to perform (batched) matrix multiplication "
          "over the last two dimensions of the input "
          "tensors X and Y.")
      .set_num_inputs(2)
#ifdef CINN_WITH_CUDA
      .set_num_outputs(1)
#else
      .set_num_outputs(2)
#endif
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForMatMul)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForMatMul))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForMatMul))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForMatMul))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
      .set_support_level(4);

  CINN_REGISTER_OP(split)
      .describe(
          "This operator is used to split tensors X to 'sections' sub-tensor "
          "on specified axis.")
      .set_num_inputs(1)
      .set_num_outputs(0)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForSplit)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForSplit))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForSplit))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForSplit))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
      .set_support_level(4);

  CINN_REGISTER_OP(concat)
      .describe(
          "This operator is used to concat two input tensors X and Y on "
          "specified axis.")
      .set_num_inputs(2)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForConcat)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForConcat))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForConcat))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForConcat))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kInjective)
      .set_support_level(4);

  CINN_REGISTER_OP(reverse)
      .describe("This operator implements the meta op reverse.")
      .set_num_inputs(1)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForReverse)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForReverse))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForLayoutTransform))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForReverse))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kInjective)
      .set_support_level(4);

  CINN_REGISTER_OP(transpose)
      .describe("This operator implements the meta op transpose.")
      .set_num_inputs(1)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForTranspose)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForTranspose))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForLayoutTransform))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForTranspose))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kInjective)
      .set_support_level(4);

  CINN_REGISTER_OP(mul)
      .describe(
          "This operator is used to perform matrix multiplication for input X "
          "and Y.")
      .set_num_inputs(2)
      .set_num_outputs(2)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForMul)
      .set_attr("infershape", MakeOpFunction(cinn::hlir::op::InferShapeForMul))
      .set_attr("inferdtype", MakeOpFunction(cinn::hlir::op::InferDtypeForMul))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForMul))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
      .set_support_level(4);

#ifdef CINN_WITH_CUDA
  CINN_REGISTER_OP(cublas_gemm)
      .describe("This operator uses cublas to compute the gemm.")
      .set_num_inputs(3)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForCublasGemm)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForCublasGemm))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForCublasGemm))
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
      .set_support_level(4);

  CINN_REGISTER_OP(cublas_matmul)
      .describe("This operator uses cublas to compute the matmul.")
      .set_num_inputs(2)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForMatMul)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForMatMul))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForMatMul))
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
      .set_support_level(4);
#endif

  CINN_REGISTER_OP(layout_transform)
      .describe("This operator is used to transform op's layouts")
      .set_num_inputs(1)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForLayoutTransform)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForLayoutTransform))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForLayoutTransform))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForLayoutTransform))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kInjective)
      .set_support_level(4);

  CINN_REGISTER_OP(slice)
      .describe("This operator implements the slice layer")
      .set_num_inputs(1)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForSlice)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForSlice))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForSlice))
#ifndef CINN_WITH_CUDA
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForSlice))
#endif
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kInjective)
      .set_support_level(4);

  CINN_REGISTER_OP(slice_assign)
      .describe(
          "This operator is used to perform slice assign for tensor input and "
          "tensor assign.")
      .set_num_inputs(2)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForSliceAssign)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForSliceAssign))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForSliceAssign))
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForSliceAssign))
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kElementWise)
      .set_support_level(4);

  CINN_REGISTER_OP(gather)
      .describe(
          "This operator is used to create a new tensor which indexes the "
          "`input` tensor along dimension `axis` using "
          "the entries in `index`.")
      .set_num_inputs(2)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForGather)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForGather))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForGather))
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForGather))
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kInjective)
      .set_support_level(4);

  CINN_REGISTER_OP(scatter_assign)
      .describe(
          "This operator is used to assign tensor B to tensor A by index.")
      .set_num_inputs(3)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForScatterAssign)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForScatterAssign))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForScatterAssign))
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForScatterAssign))
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kInjective)
      .set_support_level(4);

  CINN_REGISTER_OP(scatter_add)
      .describe(
          "This operator is used to add update tensor B into tensor A by "
          "index.")
      .set_num_inputs(3)
      .set_num_outputs(1)
      .set_attr<cinn::hlir::framework::StrategyFunction>(
          "CINNStrategy", cinn::hlir::op::StrategyForScatterAdd)
      .set_attr("infershape",
                MakeOpFunction(cinn::hlir::op::InferShapeForScatterAdd))
      .set_attr("inferdtype",
                MakeOpFunction(cinn::hlir::op::InferDtypeForScatterAdd))
      .set_attr("inferlayout",
                MakeOpFunction(cinn::hlir::op::InferLayoutForScatterAdd))
      // Because the scatter_add operator calls the external function by passing
      // pointers, the code generated by operator fusion will have out-of-bounds
      // access. It should not fuse with any other injective operators, though
      // scatter_add is injective. turn KNonFusible to kInjective will fail
      // /Paddle/python/paddle/fluid/tests/unittests/test_index_select_op.py
      .set_attr<cinn::hlir::framework::OpPatternKind>(
          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
      .set_support_level(4);

  return true;
}