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

#include "paddle/phi/backends/gpu/gpu_device_function.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/fused_gate_attention.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/fusion/gpu/attn_gemm.h"
#include "paddle/utils/optional.h"

namespace phi {
namespace fusion {

template <typename T>
struct SigmoidMultiplyFunctor {
  using MT = typename MPTypeTrait<T>::Type;
  MT one = static_cast<MT>(1.0f);

  // sigmoid(x) = 1 / (1 + exp(-x))
  // out = sigmoid(x) * y
  inline HOSTDEVICE T operator()(T x, T y) const {
    MT x_mp = static_cast<MT>(x);
    T sigmoid_out = static_cast<T>(one / (one + exp(-x_mp)));
    return sigmoid_out * y;
  }
};

template <typename T>
void ComputeMergedQKVMatmulForward(const GPUContext &dev_ctx,
                                   const funcs::GateAttentionConfig<T> &config,
                                   const DenseTensor *query,
                                   DenseTensor *qkv_out,
                                   const DenseTensor &qkv_weight_in) {
  // query: shape=[batch_size, seq_len_m, seq_len_r, qkv_dim]
  // qkv_weight: shape=[3, num_heads, head_dim, qkv_dim]
  // qkv_out: shape=[batch_size, seq_len_m, seq_len_r, 3, num_heads, head_dim]
  auto *qkv_weight = &qkv_weight_in;

  // qkv_out = GEMM(query, qkv_weight^T)
  int m = config.batch_size * config.seq_len_m * config.seq_len_r;
  int n = 3 * config.num_heads * config.head_dim;
  int k = config.q_dim;
  auto qkv_compute =
      fusion::AttnMatMul<T>(dev_ctx, false, true, m, n, k, false);
  qkv_compute.ComputeForward(qkv_weight, query, nullptr, qkv_out, nullptr);
}

template <typename T>
void ComputeSeparatedQKVMatmulForward(
    const GPUContext &dev_ctx,
    const funcs::GateAttentionConfig<T> &config,
    const DenseTensor *query,
    const DenseTensor *key,
    DenseTensor *query_out,
    DenseTensor *key_out,
    DenseTensor *value_out,
    const DenseTensor &query_weight_in,
    const DenseTensor &key_weight_in,
    const DenseTensor &value_weight_in) {
  auto *query_weight = &query_weight_in;
  auto *key_weight = &key_weight_in;
  auto *value_weight = &value_weight_in;

  // query_out = GEMM(query, query_weight)
  // query: shape=[batch_size, seq_len_m, seq_len_r, q_dim]
  // query_weight: shape=[q_dim, num_heads, head_dim]
  // query_out: shape=[batch_size, seq_len_m, seq_len_r, num_heads, head_dim]
  int q_m = config.batch_size * config.seq_len_m * config.seq_len_r;
  int q_n = config.num_heads * config.head_dim;
  int q_k = config.q_dim;
  auto q_compute =
      fusion::AttnMatMul<T>(dev_ctx, false, false, q_m, q_n, q_k, false);
  q_compute.ComputeForward(query_weight, query, nullptr, query_out, nullptr);

  // k_out = GEMM(key, key_weight)
  // key: shape=[batch_size, seq_len_m, m_size, kv_dim]
  // key_weight: shape=[kv_dim, num_heads, head_dim]
  // key_out: shape=[batch_size, seq_len_m, m_size, num_heads, head_dim]
  int kv_m = config.batch_size * config.seq_len_m * config.m_size;
  int kv_n = config.num_heads * config.head_dim;
  int kv_k = config.kv_dim;
  auto kv_compute =
      fusion::AttnMatMul<T>(dev_ctx, false, false, kv_m, kv_n, kv_k, false);
  kv_compute.ComputeForward(key_weight, key, nullptr, key_out, nullptr);

  // value_out = GEMM(value, value_weight)
  kv_compute.ComputeForward(value_weight, key, nullptr, value_out, nullptr);
}

template <typename T>
void ComputeGatingLinearForward(const GPUContext &dev_ctx,
                                const funcs::GateAttentionConfig<T> &config,
                                const DenseTensor *query,
                                const DenseTensor *fmha_out,
                                DenseTensor *gate_bias_out,
                                bool use_fused_matmul_bias,
                                const DenseTensor &gate_weight_in,
                                const DenseTensor &gate_bias_in) {
  auto *gate_weight = &gate_weight_in;
  auto *gate_bias = &gate_bias_in;

  // The first gate_bias_out stores the result of the multiplication,
  // and the second gate_bias_out stores the result of the multiplication +
  // bias.
  //   gate_out = GEMM(query, gate_weight) + gate_bias
  int m = config.batch_size * config.seq_len_m * config.seq_len_r;
  int n = config.num_heads * config.head_dim;
  int k = config.q_dim;
  auto gate_linear =
      fusion::AttnMatMul<T>(dev_ctx, false, false, m, n, k, true);
  gate_linear.ComputeForward(gate_weight,
                             query,
                             gate_bias,
                             gate_bias_out,
                             gate_bias_out,
                             use_fused_matmul_bias);

  // gate_out = sigmoid(gate_out) * fmha_out
  std::vector<const DenseTensor *> ins = {gate_bias_out, fmha_out};
  std::vector<DenseTensor *> outs = {gate_bias_out};
  funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, SigmoidMultiplyFunctor<T>());
}

template <typename T>
void ComputeOutputLinearForward(const GPUContext &dev_ctx,
                                const funcs::GateAttentionConfig<T> &config,
                                const DenseTensor *fmha_or_gate_out,
                                DenseTensor *out,
                                bool use_fused_matmul_bias,
                                const DenseTensor &out_linear_weight_in,
                                const DenseTensor &out_linear_bias_in) {
  const auto *out_linear_weight = &out_linear_weight_in;
  const auto *out_linear_bias = &out_linear_bias_in;

  // out = GEMM(fmha_or_gate_out, out_linear_weight) + out_linear_bias
  int m = config.batch_size * config.seq_len_m * config.seq_len_r;
  int n = config.q_dim;
  int k = config.num_heads * config.head_dim;
  auto out_linear = fusion::AttnMatMul<T>(dev_ctx, false, false, m, n, k, true);
  out_linear.ComputeForward(out_linear_weight,
                            fmha_or_gate_out,
                            out_linear_bias,
                            out,
                            out,
                            use_fused_matmul_bias);
}

template <typename T, typename Context>
void FusedGateAttentionOpKernel(const Context &dev_ctx,
                                const DenseTensor &query_in,
                                const optional<DenseTensor> &key_in,
                                const optional<DenseTensor> &query_weight_in,
                                const optional<DenseTensor> &key_weight_in,
                                const optional<DenseTensor> &value_weight_in,
                                const optional<DenseTensor> &qkv_weight_in,
                                const optional<DenseTensor> &nonbatched_bias_in,
                                const DenseTensor &src_mask_in,
                                const optional<DenseTensor> &gate_weight_in,
                                const optional<DenseTensor> &gate_bias_in,
                                const DenseTensor &out_linear_weight_in,
                                const DenseTensor &out_linear_bias_in,
                                bool has_gating,
                                bool merge_qkv,
                                bool use_flash_attn,
                                DenseTensor *query_transpose_out,
                                DenseTensor *key_transpose_out,
                                DenseTensor *value_transpose_out,
                                DenseTensor *qkv_transpose_out,
                                DenseTensor *softmax_out,
                                DenseTensor *softmax_lse,
                                DenseTensor *fmha_out,
                                DenseTensor *gate_out,
                                DenseTensor *out) {
  const auto *query = &query_in;
  const auto *key = key_in.get_ptr();
  const auto *query_weight = query_weight_in.get_ptr();
  const auto *qkv_weight = qkv_weight_in.get_ptr();

  const auto *src_mask = &src_mask_in;
  const auto *nonbatched_bias = nonbatched_bias_in.get_ptr();

  auto *q_transpose_out = query_transpose_out;
  auto *k_transpose_out = key_transpose_out;
  auto *v_transpose_out = value_transpose_out;

  bool use_fused_matmul_bias = true;
  funcs::AllocWithDebugInfo<T>(dev_ctx, "fmha_out", fmha_out);
  if (has_gating) {
    funcs::AllocWithDebugInfo<T>(dev_ctx, "gate_out", gate_out);
  }
  funcs::AllocWithDebugInfo<T>(dev_ctx, "out", out);

  // When seq_len_r = m_size, q_dim = kv_dim, QKV matmul can be merged.
  funcs::GateAttentionConfig<T> config(dev_ctx,
                                       query,
                                       key,
                                       query_weight,
                                       qkv_weight,
                                       merge_qkv,
                                       has_gating,
                                       use_flash_attn);

  if (merge_qkv) {
    PADDLE_ENFORCE_EQ(
        !key || query == key || query->data<T>() == key->data<T>(),
        true,
        errors::InvalidArgument("key is expected to be nullptr or the same as "
                                "query, but received key=%p, query=%p.",
                                key,
                                query));

    // 1. Merged QKV Matmul: einsum(nbhqk,nbkhc -> nbqhc)
    DenseTensor *qkv_out = config.GetQKVOut();
    ComputeMergedQKVMatmulForward<T>(
        dev_ctx, config, query, qkv_out, qkv_weight_in.get());

    if (config.CanUseFlashAttn()) {
      qkv_transpose_out->Resize({3,
                                 config.batch_size,
                                 config.seq_len_m,
                                 config.seq_len_r,
                                 config.num_heads,
                                 config.head_dim});
    }
    funcs::AllocWithDebugInfo<T>(
        dev_ctx, "qkv_transpose_out", qkv_transpose_out);
  } else {
    // 1. Separated QKV Matmul
    DenseTensor *query_out = config.GetQueryOut();
    DenseTensor *key_out = config.GetKeyOut();
    DenseTensor *value_out = config.GetValueOut();
    ComputeSeparatedQKVMatmulForward<T>(dev_ctx,
                                        config,
                                        query,
                                        key,
                                        query_out,
                                        key_out,
                                        value_out,
                                        query_weight_in.get(),
                                        key_weight_in.get(),
                                        value_weight_in.get());

    funcs::AllocWithDebugInfo<T>(dev_ctx, "q_transpose_out", q_transpose_out);
    funcs::AllocWithDebugInfo<T>(dev_ctx, "k_transpose_out", k_transpose_out);
    funcs::AllocWithDebugInfo<T>(dev_ctx, "v_transpose_out", v_transpose_out);
  }

  // 2. FMHA
  if (config.CanUseFlashAttn()) {
    auto fmha_compute = funcs::FlashAttnWithGating<T>(dev_ctx, merge_qkv);
    fmha_compute.ComputeForward(nonbatched_bias,
                                src_mask,
                                qkv_transpose_out,
                                softmax_lse,
                                fmha_out,
                                &config);
  } else {
    funcs::AllocWithDebugInfo<T>(dev_ctx, "softmax_out", softmax_out);

    auto fmha_compute = funcs::FMHAGateRef<T>(dev_ctx, merge_qkv);
    fmha_compute.ComputeForward(nonbatched_bias,
                                src_mask,
                                q_transpose_out,
                                k_transpose_out,
                                v_transpose_out,
                                qkv_transpose_out,
                                softmax_out,
                                fmha_out,
                                gate_out,
                                &config);
  }

  // 3. Gating Linear
  if (has_gating) {
    ComputeGatingLinearForward<T>(dev_ctx,
                                  config,
                                  query,
                                  fmha_out,
                                  gate_out,
                                  use_fused_matmul_bias,
                                  gate_weight_in.get(),
                                  gate_bias_in.get());
  }

  // 4. Output Linear
  DenseTensor *fmha_or_gate_out = has_gating ? gate_out : fmha_out;
  ComputeOutputLinearForward<T>(dev_ctx,
                                config,
                                fmha_or_gate_out,
                                out,
                                use_fused_matmul_bias,
                                out_linear_weight_in,
                                out_linear_bias_in);
}
}  // namespace fusion
}  // namespace phi

#ifdef PADDLE_WITH_HIP
PD_REGISTER_KERNEL(fused_gate_attention,
                   GPU,
                   ALL_LAYOUT,
                   phi::fusion::FusedGateAttentionOpKernel,
                   float,
                   phi::float16,
                   phi::bfloat16) {}
#else
PD_REGISTER_KERNEL(fused_gate_attention,
                   GPU,
                   ALL_LAYOUT,
                   phi::fusion::FusedGateAttentionOpKernel,
                   float,
                   double,
                   phi::float16,
                   phi::bfloat16) {}
#endif