paddlepaddle--paddle/paddle/phi/kernels/impl/einsum_kernel_impl.h

// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once

#include <set>

#include "glog/logging.h"
#include "paddle/common/flags.h"

#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/diagonal_kernel.h"
#include "paddle/phi/kernels/fill_diagonal_tensor_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/matmul_kernel.h"
#include "paddle/phi/kernels/reduce_sum_kernel.h"
#include "paddle/phi/kernels/tile_kernel.h"
#include "paddle/phi/kernels/transpose_kernel.h"
#include "paddle/utils/string/string_helper.h"

PD_DECLARE_bool(einsum_opt);

namespace phi {

// check the validation of the Einsum equation.
// 1. the label must between 'a' - 'z'.
// 2. the dim of the same label must be same.
// 3. the broad cast dims in two operands is broadcastable.
// 4. there must exist '->' and the default output is complete in python.
// may be we can skip validation check in C++ and just put it in python.
inline static void ValidationCheck(const std::string& equation) {
  auto n_part = paddle::string::split_string(equation, "->").size();
  PADDLE_ENFORCE_EQ(n_part,
                    2,
                    common::errors::InvalidArgument(
                        "Required at least one `->` in equation of EinsumOp."));
  size_t pos;
  auto trimmed_equ = equation;
  if ((pos = trimmed_equ.find("->", 0)) != std::string::npos) {
    trimmed_equ.replace(pos, 2, "");
  }
  auto is_valid_char = [](char c) {
    if (c >= 'a' && c <= 'z') return true;
    if (c == ',') return true;
    return false;
  };
  for (auto c : trimmed_equ) {
    if (!is_valid_char(c))
      PADDLE_THROW(common::errors::InvalidArgument(
          "Found invalid char in equation. Einsum only accept `a`-`z` and `...`"
          "but get:`%c`",
          c));
  }
}

enum LabelType {
  ALL_TYPE = 0,
  Batch = 1,    // ABO
  AO,           // AO --  free label
  BO,           // BO --  free label
  Contraction,  // AB
  Reduction,    // A, B
};

// map a label('a' - 'z') -> int64_t, O(1) speed.
class LabelMap {
  constexpr static int N =
      26 + 1;  // 'a' - 'z' + '.', '.' is for broadcast dims
  int64_t default_value;
  int64_t map[N];

 public:
  explicit LabelMap(int64_t default_value = 0) {
    this->default_value = default_value;
    for (size_t i = 0; i < N; ++i) map[i] = default_value;
  }
  int64_t& operator[](int label) {
    int i = label - 'a';
    return map[i];
  }
  int64_t operator[](int label) const {
    int i = label - 'a';
    return map[i];
  }
  bool exist(char label) { return !is_default(label); }

 private:
  // non-exist is present by is_default
  bool is_default(char label) {
    return (*this)[static_cast<int>(label)] == default_value;
  }
};

inline std::string label_to_string(const std::vector<char>& all_labels,
                                   const LabelMap& label2type) {
  std::string str;
  for (int a : all_labels) {
    std::stringstream ss;
    ss << label2type[a];
    str += ss.str();
  }
  return str;
}

template <typename CharIterable1, typename CharIterable2>
inline std::vector<char> union_labels(const CharIterable1& a,
                                      const CharIterable2& b) {
  LabelMap counter(0);
  std::vector<char> res;
  auto f = [&](char c) {
    if (counter[static_cast<int>(c)] == 0) {
      res.push_back(c);
    }
    counter[static_cast<int>(c)] += 1;
  };
  std::for_each(a.begin(), a.end(), f);
  std::for_each(b.begin(), b.end(), f);
  return res;
}

template <typename CharIterable>
inline std::vector<char> unique_labels(const CharIterable& a) {
  return union_labels(a, CharIterable());
}

// Apply transforms to all_labels and get another all_labels
inline std::vector<char> TransformLabelsOrder(
    const std::vector<char>& all_labels,
    const LabelMap& type,
    std::vector<LabelType> new_order) {
  std::vector<char> ret;
  for (auto cnt_type : new_order) {
    std::vector<char> tmp;
    for (int c : all_labels) {
      if (type[c] == cnt_type) tmp.push_back(c);
    }
    ret.insert(ret.end(), tmp.begin(), tmp.end());
  }
  return ret;
}

inline static void GlobalInfo(const std::vector<std::string>& op_labels,
                              const std::string& right,
                              LabelMap* label2type,
                              std::vector<char>* sorted_labels) {
  std::vector<char> all;
  LabelMap counter(0);
  for (auto& ch : right) {  // char
    int c = ch;
    (*label2type)[c] = LabelType::BO;
  }

  for (auto& op : op_labels) {
    for (auto& ch : unique_labels(op)) {  // char
      int c = ch;
      if (!counter.exist(c)) {
        all.push_back(ch);
      }
      counter[c] += 1;
      if ((*label2type)[c] != LabelType::BO && counter[c] == 2)
        (*label2type)[c] = LabelType::Contraction;
      else if (counter[c] == 2)
        (*label2type)[c] = LabelType::Batch;
    }
  }

  // BO is represent Free, so we need find the AO.
  for (int c : op_labels[0]) {
    if ((*label2type)[c] == LabelType::BO) (*label2type)[c] = LabelType::AO;
  }

  if (sorted_labels->size()) {
    std::set<char> exist(all.begin(), all.end());
    all.clear();
    std::for_each(
        sorted_labels->begin(), sorted_labels->end(), [&exist, &all](char c) {
          if (exist.count(c)) all.push_back(c);
        });
  }

  *sorted_labels = TransformLabelsOrder(all,
                                        *label2type,
                                        {LabelType::Batch,
                                         LabelType::AO,
                                         LabelType::BO,
                                         LabelType::Contraction,
                                         LabelType::Reduction});

  VLOG(5) << "GlobalInfo: sorted_labels after: "
          << paddle::string::join_strings(*sorted_labels, ",");
}

inline static void InferLabelShape(
    const std::vector<std::string>& op_labels,
    const std::vector<DDim>& inputs,
    LabelMap* labelshape,
    std::vector<std::vector<int64_t>>* broadcast_shapes,
    LabelMap* labeltype) {
  LabelMap labelshape_copy = *labelshape;
  VLOG(5) << "Start InferLabelShape";
  for (size_t i = 0; i < op_labels.size(); ++i) {
    auto& op_str = op_labels[i];
    auto& op_dim = inputs[i];
    VLOG(5) << "i = " << i << " op_str " << op_str << " op_dim " << op_dim;
    int dim_ptr = 0;
    for (auto& c : op_str) {
      if (!labelshape->exist(c) || abs((*labelshape)[c]) == 1) {
        VLOG(5)
            << "if (!labelshape->exist(c) || abs((*labelshape)[c]) == 1) c = "
            << c << " (*labelshape)[c] " << (*labelshape)[c]
            << " op_dim[dim_ptr] " << op_dim[dim_ptr];
        (*labelshape)[c] = static_cast<int>(op_dim[dim_ptr]);
      } else if (abs(op_dim[dim_ptr]) != 1) {
        VLOG(5) << "if (abs(op_dim[dim_ptr]) != 1) c = " << c
                << " (*labelshape)[c] " << (*labelshape)[c]
                << " op_dim[dim_ptr] " << op_dim[dim_ptr];
        PADDLE_ENFORCE_EQ(
            (*labelshape)[c],
            op_dim[dim_ptr],
            common::errors::InvalidArgument(
                "Same label have different shapes for label: `%c`", c));
      }
      dim_ptr++;
    }
  }
  for (size_t i = 0; i < op_labels.size(); ++i) {
    for (auto& c : op_labels[i]) {
      // Note: When broadcasting is involved, ensure the gradient is calculated
      // with respect to the broadcasted shape. For example, in
      // einsum("ij,ij->j", x(2,2), y(1,2)), y is broadcast to (2,2). The
      // gradient calculation for x must use this broadcasted shape of y.
      if (labelshape_copy.exist(c) && labelshape_copy[c] > (*labelshape)[c]) {
        // Strict check for the situation.
        PADDLE_ENFORCE_EQ(
            (*labelshape)[c] == 1 && ((*labeltype)[c] == LabelType::AO ||
                                      (*labeltype)[c] == LabelType::BO),
            true,
            common::errors::InvalidArgument(
                "Broadcast dims must be 1 for label: `%c`", c));
        (*labelshape)[c] = labelshape_copy[c];
      }
      (*broadcast_shapes)[i].push_back((*labelshape)[c]);
    }
  }
  for (size_t i = 0; i < op_labels.size(); ++i) {
    VLOG(5) << "InferLabelShape: After broadcast shape is:"
            << paddle::string::join_strings((*broadcast_shapes)[i], ",");
  }
}

template <class CharIterable>
inline static void InferLabelPerm(const CharIterable& op,
                                  LabelMap* label2perm) {
  int cur = 0;
  for (int c : op) {
    if (!label2perm->exist(
            c))  // can appear repeatedly. we just record the first position.
      (*label2perm)[c] = cur;
    cur += 1;
  }
}

inline static void InferOutputDims(const std::string& right,
                                   const LabelMap& labelshape,
                                   std::vector<int64_t>* output_dims) {
  for (int c : right) {
    output_dims->push_back(labelshape[c]);
  }
}
//
inline static void ParseEinsumEquation(
    const std::string& equation,
    const std::vector<DDim>& inputs,
    LabelMap* labelshape,
    LabelMap* labeltype,
    std::vector<char>* all_labels,
    std::vector<LabelMap>* label2perms,
    std::vector<std::vector<int64_t>>* broadcast_shapes,
    std::vector<int64_t>* output_dims,
    std::string* right,
    std::vector<std::string>* input_strs) {
  VLOG(5) << "Start ParseEinsumEquation " << equation;
  auto results = paddle::string::split_string(equation, "->");
  auto left = results[0];
  *right = results[1];
  auto op_labels = paddle::string::split_string(left, ",");
  // split_string("i,") -> ["i", ""], we push back a "".
  // split_string("->") -> [], we push back a "".
  if (op_labels.empty()) op_labels.emplace_back("");
  GlobalInfo(op_labels, *right, labeltype, all_labels);
  InferLabelShape(op_labels, inputs, labelshape, broadcast_shapes, labeltype);
  VLOG(5) << "Einsum Infershape: right:" << *right;
  VLOG(5) << "Einsum Infershape: left :"
          << paddle::string::join_strings(op_labels, '\n');
  InferOutputDims(*right, *labelshape, output_dims);
  for (size_t i = 0; i < inputs.size(); ++i) {
    InferLabelPerm(op_labels[i], &((*label2perms)[i]));
    (*input_strs).push_back(std::move(op_labels[i]));
  }
}

template <typename T>
std::vector<T> GetLabelIndexByType(const std::vector<char>& all_labels,
                                   const LabelMap& type,
                                   const LabelMap& perm,
                                   LabelType filter) {
  std::vector<T> res;
  for (T c : all_labels) {
    if ((filter == LabelType::ALL_TYPE || type[c] == filter) && perm[c] != -1) {
      res.push_back(perm[c]);
    }
  }
  return res;
}

template <typename T>
std::vector<T> GetShapeByType(const std::vector<char>& all_labels,
                              const LabelMap& type,
                              const LabelMap& perm,
                              const LabelMap& label2shape,
                              std::set<LabelType> filter) {
  std::vector<T> res;
  for (T c : all_labels) {
    if ((filter.count(LabelType::ALL_TYPE) ||
         filter.count(LabelType(type[c]))) &&
        perm[c] != -1) {
      res.push_back(label2shape[c]);
    }
  }
  return res;
}

inline static std::vector<int> perm_moveto(int n, int from, int to) {
  // a permutation means moving `from` to `to`.
  /*
  f => t   permutation
  --------------------
           0 1 2 3 4 5
  5 => 2 : 0 2 5 2 3 4
  2 => 5 : 0 1 3 4 5 2
  we can conclude the following rules.
  */
  if (from < 0) from = n + from;
  if (to < 0) to = n + to;
  std::vector<int> res(n);
  for (int i = 0; i < n; ++i) {
    res[i] = i;
  }
  res[to] = from;
  auto offset = from > to ? -1 : 1;
  auto start = from > to ? to + 1 : from;
  auto end = from > to ? from : to - 1;
  for (int i = start; i <= end; ++i) {
    res[i] += offset;
  }
  return res;
}

template <typename T, typename Context>
DenseTensor Undiagonal(const Context& dev_ctx,
                       const DenseTensor& tensor,
                       size_t insert_pos,
                       size_t axis) {
  // tensor with shape (3, 4, 5, 2, 1), insert_pos = 5, axis = 2.
  // output is (3, 4, 5, 2, 1, 5)
  VLOG(5) << "Start undiagonal with args: insert_pos = " << insert_pos
          << ", axis = " << axis;
  std::vector<int64_t> shape(tensor.dims().size() + 1);
  int point = 0;  // point to the tensor.dims()
  for (size_t i = 0; i < shape.size(); ++i) {
    if (i == insert_pos)
      shape[i] = tensor.dims()[axis];
    else
      shape[i] = tensor.dims()[point++];
  }
  auto zeros = Full<T, Context>(dev_ctx, shape, 0);
  auto diags = Transpose<T, Context>(
      dev_ctx, tensor, perm_moveto(tensor.dims().size(), axis, -1));
  return FillDiagonalTensor<T, Context>(
      dev_ctx, zeros, diags, 0, insert_pos, axis + (insert_pos <= axis));
}

template <typename T, typename Context>
DenseTensor PerformUndiagonal(const Context& dev_ctx,
                              const DenseTensor& tensor,
                              const std::string& equ) {
  //  if the equ is 'iijjkij', then the tensor must be 'ijk', so we have enough
  //  information to do un-diagonal with equ.
  auto res = tensor;
  LabelMap label2perm(-1);
  InferLabelPerm(equ, &label2perm);
  // Un-Diagonal
  int tot = equ.size();
  int cur = tot - 1;
  for (auto it = equ.rbegin(); it != equ.rend(); ++it) {
    char c = *it;
    if (cur != label2perm[c]) {
      // do diagonal, followed by movedim().
      auto insert_pos = cur - tot + res.dims().size() + 1;
      res = Undiagonal<T, Context>(dev_ctx, res, insert_pos, label2perm[c]);
    }
    --cur;
  }
  return res;
}

template <typename T, typename Context>
DenseTensor PerformDiagonalAndReduction(
    const Context& dev_ctx,
    const DenseTensor& tensor,
    const std::string& equ,
    const LabelMap& label2perm,
    const std::vector<char>& all_labels,
    const std::vector<int64_t>& broadcast_shape,
    const LabelMap& label2type) {
  auto res = tensor;
  int tot = equ.size();
  // tiling tensor for broadcast
  std::vector<int64_t> repeat_times;
  auto tensor_origin_shape = vectorize(tensor.dims());
  for (size_t i = 0; i < tensor_origin_shape.size(); ++i) {
    VLOG(4) << "broadcast shape is " << broadcast_shape[i]
            << ", tensor shape is " << tensor_origin_shape[i];
    repeat_times.push_back(broadcast_shape[i] / tensor_origin_shape[i]);
  }
  DenseTensor after_tile;
  bool is_all_ones = std::all_of(repeat_times.begin(),
                                 repeat_times.end(),
                                 [](int64_t x) { return x == 1; });
  if (!is_all_ones) {
    TileKernel<T, Context>(dev_ctx, res, repeat_times, &after_tile);
    res = after_tile;
  }
  // Diagonal
  int cur = tot - 1;
  for (auto it = equ.rbegin(); it != equ.rend(); ++it) {
    char c = *it;
    if (cur != label2perm[c]) {
      // do diagonal, followed by movedim().
      VLOG(5) << "Do diagonal with shape="
              << paddle::string::join_strings(vectorize<int64_t>(res.dims()),
                                              ',')
              << ", axis1=" << cur << ", axis2=" << label2perm[c];
      res = Diagonal<T, Context>(dev_ctx, res, 0, cur, label2perm[c]);
      res = Transpose<T, Context>(
          dev_ctx, res, perm_moveto(res.dims().size(), -1, label2perm[c]));
    }
    --cur;
  }
  // reduction
  auto indices = GetLabelIndexByType<int64_t>(
      all_labels, label2type, label2perm, LabelType::Reduction);
  VLOG(5) << "call PerformDiagonalAndReduction: with axis: "
          << paddle::string::join_strings(indices, ",");
  if (indices.empty()) return res;
  return Sum<T, Context>(dev_ctx, res, IntArray(indices), res.dtype(), true);
}

inline bool is_no_need_transpose(const std::vector<int>& axis) {
  for (size_t i = 0; i < axis.size(); ++i) {
    if (i != static_cast<size_t>(axis[i])) return false;
  }
  return true;
}

template <typename T, typename Context>
DenseTensor PerformTranspose(const Context& dev_ctx,
                             const DenseTensor& tensor,
                             const LabelMap& label2perm,
                             const std::vector<char>& all_labels,
                             const LabelMap& label2type) {
  auto axis = GetLabelIndexByType<int>(
      all_labels, label2type, label2perm, LabelType::ALL_TYPE);
  VLOG(5) << "PerformTranspose: " << paddle::string::join_strings(axis, ",");
  if (is_no_need_transpose(axis)) {
    return tensor;
  }
  auto ret = Transpose<T, Context>(dev_ctx, tensor, axis);
  VLOG(5) << "PerformTranspose: do_transpose()";
  return ret;
}

template <typename T, typename Context>
DenseTensor PerformContraction(
    const Context& dev_ctx,
    const std::vector<const DenseTensor*>& operands,
    const std::vector<std::string>& input_strs,
    const std::vector<LabelMap>& label2perm,
    const std::vector<char>& all_labels,
    const LabelMap& label2type,
    const LabelMap& label2shape,
    const std::vector<std::vector<int64_t>>& broadcast_shapes,
    std::vector<DenseTensor*> cache,
    bool use_cache) {
  auto all_valid = LabelMap(1);
  auto recover_dim = GetShapeByType<int64_t>(
      all_labels, label2type, all_valid, label2shape, {LabelType::Batch});
  auto preprocess = [&](const DenseTensor& t,
                        const LabelMap& perm,
                        const std::vector<int64_t>& broadcast,
                        int operand_idx) -> DenseTensor {
    // reshape
    auto frees = GetShapeByType<int64_t>(all_labels,
                                         label2type,
                                         perm,
                                         label2shape,
                                         {LabelType::AO, LabelType::BO});
    auto conts = GetShapeByType<int64_t>(
        all_labels, label2type, perm, label2shape, {LabelType::Contraction});
    std::vector<char> reordered_all_labels = all_labels;
    if (operand_idx == 1) {
      reordered_all_labels = TransformLabelsOrder(all_labels,
                                                  label2type,
                                                  {LabelType::Batch,
                                                   LabelType::Contraction,
                                                   LabelType::AO,
                                                   LabelType::BO,
                                                   LabelType::Reduction});
    }
    // reduction
    DenseTensor trans_t;
    if (use_cache && cache[operand_idx] != nullptr &&
        cache[operand_idx]->IsInitialized()) {
      trans_t.ShareBufferWith(*(cache[operand_idx]));
      VLOG(5) << "Cache Used!";
    } else {
      auto reduct_t =
          PerformDiagonalAndReduction<T, Context>(dev_ctx,
                                                  t,
                                                  input_strs[operand_idx],
                                                  perm,
                                                  all_labels,
                                                  broadcast_shapes[operand_idx],
                                                  label2type);
      trans_t = PerformTranspose<T, Context>(
          dev_ctx, reduct_t, perm, reordered_all_labels, label2type);
      if (cache[operand_idx] != nullptr) {
        std::vector<int64_t> broadcast_shapes_restore(
            broadcast_shapes[operand_idx].size());

        auto contraction_dim1 =
            [&](const std::vector<int64_t>& broadcast_shapes,
                const std::vector<int64_t>& original_shapes) -> bool {
          bool found = false;
          for (size_t i = 0; i < broadcast_shapes.size(); ++i) {
            if (broadcast_shapes[i] != original_shapes[i] &&
                label2type[input_strs[operand_idx][i]] ==
                    LabelType::Contraction) {
              broadcast_shapes_restore[i] = original_shapes[i];
              found = true;
            } else {
              broadcast_shapes_restore[i] = broadcast_shapes[i];
            }
          }
          return found;
        };
        if (!contraction_dim1(broadcast_shapes[operand_idx],
                              vectorize<int64_t>(t.dims()))) {
          cache[operand_idx]->ShareBufferWith(trans_t);
          cache[operand_idx]->Resize(trans_t.dims());
          VLOG(5) << "Set dims of cache[" << operand_idx
                  << "]: " << trans_t.dims();
        } else {
          auto reduct_t_for_cache =
              PerformDiagonalAndReduction<T, Context>(dev_ctx,
                                                      t,
                                                      input_strs[operand_idx],
                                                      perm,
                                                      all_labels,
                                                      broadcast_shapes_restore,
                                                      label2type);
          DenseTensor trans_t_for_cache;
          trans_t_for_cache = PerformTranspose<T, Context>(dev_ctx,
                                                           reduct_t_for_cache,
                                                           perm,
                                                           reordered_all_labels,
                                                           label2type);
          cache[operand_idx]->ShareBufferWith(trans_t_for_cache);
          cache[operand_idx]->Resize(trans_t_for_cache.dims());
          VLOG(5) << "Set dims of cache[" << operand_idx
                  << "]: " << trans_t_for_cache.dims();
        }
      }
    }
    auto mul_dims = GetShapeByType<int64_t>(
        all_labels, label2type, perm, label2shape, {LabelType::Batch});
    recover_dim.insert(recover_dim.end(), frees.begin(), frees.end());
    if (operand_idx == 0) {
      mul_dims.push_back(std::accumulate(
          frees.begin(), frees.end(), 1, std::multiplies<int64_t>()));
      mul_dims.push_back(std::accumulate(
          conts.begin(), conts.end(), 1, std::multiplies<int64_t>()));
    } else {
      mul_dims.push_back(std::accumulate(
          conts.begin(), conts.end(), 1, std::multiplies<int64_t>()));
      mul_dims.push_back(std::accumulate(
          frees.begin(), frees.end(), 1, std::multiplies<int64_t>()));
    }
    VLOG(5) << "PerformContraction: mul_dims: "
            << paddle::string::join_strings(mul_dims, ",");
    trans_t.Resize(mul_dims);
    return trans_t;
  };

  // Reduction, Reshape and Matmul
  DenseTensor after_contraction;
  if (operands.size() == 2) {
    auto trans_a =
        preprocess(*(operands[0]), label2perm[0], broadcast_shapes[0], 0);
    auto trans_b =
        preprocess(*(operands[1]), label2perm[1], broadcast_shapes[1], 1);
    after_contraction =
        Matmul<T, Context>(dev_ctx, trans_a, trans_b, false, false);
  } else if (operands.size() == 1) {
    after_contraction =
        preprocess(*(operands[0]), label2perm[0], broadcast_shapes[0], 0);
  }
  if (recover_dim.empty()) recover_dim.push_back(1);
  VLOG(5) << "PerformContraction: recover_dim: "
          << paddle::string::join_strings(recover_dim, ",");
  after_contraction.Resize(recover_dim);
  return after_contraction;
}

template <typename T, typename Context>
DenseTensor TransposeToOutput(const Context& dev_ctx,
                              const DenseTensor& to_trans,
                              const std::vector<char>& right,
                              const std::vector<char>& all_labels) {
  std::vector<int> axis;
  for (char c : right) {
    auto it = std::find(all_labels.begin(), all_labels.end(), c);
    PADDLE_ENFORCE_NE(it,
                      all_labels.end(),
                      common::errors::InvalidArgument("Must in all_labels."));
    axis.push_back(it - all_labels.begin());
  }
  if (is_no_need_transpose(axis)) {
    return to_trans;
  }
  VLOG(5) << "call TransposeToOutput: with axis: "
          << paddle::string::join_strings(axis, ",")
          << "  to trans dims is: " << to_trans.dims();
  auto output = Transpose<T, Context>(dev_ctx, to_trans, axis);
  VLOG(5) << "After Transpose.";
  return output;
}

template <typename T, typename Context>
void EinsumKernelImpl(const Context& dev_ctx,
                      const std::vector<char>& forward_all_labels,
                      const LabelMap& forward_label_shape,
                      const std::vector<const DenseTensor*>& inputs,
                      const std::string& equation,
                      DenseTensor* out,
                      std::vector<DenseTensor*> cache,
                      bool is_forward = true) {
  VLOG(5) << "Start EinsumKernelImpl with inputs(" << inputs.size() << "): ";
  for (auto& i : inputs) {
    VLOG(5) << "      inputs [ " << i << " ].shape=" << i->dims();
  }
  ValidationCheck(equation);
  // collect the following information to prepare einsum.
  LabelMap labelshape(0);
  LabelMap labeltype(LabelType::Reduction);
  std::vector<LabelMap> label2perms(inputs.size(), LabelMap(-1));
  std::vector<char> all_labels;  // order: ABO, AO, BO, AB, Reduce
  std::vector<std::vector<int64_t>> broadcast_shapes(2);
  std::vector<int64_t> output_dims;

  std::vector<DDim> input_dims;
  for (auto& i : inputs) {
    input_dims.push_back(i->dims());
  }
  std::vector<std::string> input_strs;
  std::string right;
  if (!is_forward) {
    all_labels = forward_all_labels;
    labelshape = forward_label_shape;
  }
  ParseEinsumEquation(equation,
                      input_dims,
                      &labelshape,
                      &labeltype,
                      &all_labels,
                      &label2perms,
                      &broadcast_shapes,
                      &output_dims,
                      &right,
                      &input_strs);
  if (inputs.size() > 2) {
    PADDLE_THROW(common::errors::InvalidArgument(
        "EinsumOp kernel only support len(operands) between (0, 2]. Use "
        "opt_einsum first to convert multi-variable to binary-variable."));
  }
  auto after_contraction = PerformContraction<T, Context>(dev_ctx,
                                                          inputs,
                                                          input_strs,
                                                          label2perms,
                                                          all_labels,
                                                          labeltype,
                                                          labelshape,
                                                          broadcast_shapes,
                                                          cache,
                                                          !is_forward);
  *out = TransposeToOutput<T, Context>(
      dev_ctx, after_contraction, unique_labels(right), all_labels);
  *out = PerformUndiagonal<T, Context>(dev_ctx, *out, right);
  out->Resize(output_dims);
}

template <typename T, typename Context>
void EinsumKernel(const Context& dev_ctx,
                  const std::vector<const DenseTensor*>& inputs,
                  const std::string& equation,
                  DenseTensor* out,
                  std::vector<DenseTensor*> cache,
                  std::vector<DenseTensor*> xshape UNUSED) {
  for (const auto& input : inputs) {
    if (input->numel() == 0) {
      dev_ctx.template Alloc<T>(out);
      if (out->numel() > 0) {
        std::vector<int64_t> vec_dims = vectorize(out->dims());
        Full<T, Context>(dev_ctx, IntArray(vec_dims), static_cast<T>(0), out);
      }
      return;
    }
  }
  std::vector<char> tmp;
  LabelMap labelshape_holder;
  // for the sake of compatibility, we may load and run v2.3 EinsumOp. Output
  // may have nullptr and the cache.size() is not equal to inputs.size(). refer
  // to BuildPhiKernelContext for details.
  int diff = inputs.size() - cache.size();
  for (int i = 0; i < diff; ++i) {
    cache.push_back(nullptr);
  }
  EinsumKernelImpl<T, Context>(dev_ctx,
                               tmp,
                               labelshape_holder,
                               inputs,
                               equation,
                               out,
                               cache,
                               /*forward=*/true);
}

template <typename T, typename Context>
void EinsumInferKernel(const Context& dev_ctx,
                       const std::vector<const DenseTensor*>& inputs,
                       const std::string& equation,
                       DenseTensor* out) {
  std::vector<char> place_holder;
  LabelMap labelshape_holder;
  std::vector<DenseTensor*> cache_tensor(
      inputs.size());  // set empty; TA, TB, TdC
  for (size_t i = 0; i < inputs.size(); ++i) {
    cache_tensor[i] = nullptr;
  }
  EinsumKernelImpl<T, Context>(dev_ctx,
                               place_holder,
                               labelshape_holder,
                               inputs,
                               equation,
                               out,
                               cache_tensor,
                               true);
}

}  // namespace phi