// 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.

#include <algorithm>

#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/fused_softmax_mask_kernel.h"
#include "paddle/phi/kernels/fusion/gpu/fused_softmax_mask_utils.h"

namespace phi {
namespace fusion {

#define LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, MT, pow2_index) \
  SoftmaxMaskFuseV2GPUKernel<T, MT, pow2_index>                   \
      <<<blocks, threads, 0, stream>>>(x_data,                    \
                                       mask_data,                 \
                                       y_data,                    \
                                       batch_count,               \
                                       attn_heads,                \
                                       query_seqs,                \
                                       key_seq_len);

#define LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, MT, pow2_index) \
  SoftmaxMaskFuseV1GPUKernel<T, MT, pow2_index>                   \
      <<<blocks, threads, 0, stream>>>(                           \
          x_data, mask_data, y_data, batch_count, key_seq_len);

// T == fp16
// SoftmaxMaskFuseV1GPUKernel is only suitable for small-scale data, with
// limited block partitioning and a relatively small memory index. But it has
// good performance
template <typename T, typename MT, int pow2_index>
__global__ void SoftmaxMaskFuseV1GPUKernel(const T* x_data,
                                           const MT* mask_data,
                                           T* y_data,
                                           int batch_count,
                                           int key_seq_len) {
  // the forward gpu kernel
  constexpr int next_pow2 = 1 << pow2_index;
  constexpr int warp_size = (next_pow2 < WARP_SIZE) ? next_pow2 : WARP_SIZE;
  constexpr int kLocalIterations = std::max(next_pow2 / warp_size, 4);
  constexpr int kLocalBatchSize = (next_pow2 <= 128) ? 2 : 1;
  constexpr int kOneLoadingCounts = 4;

  int64_t data_first_idx = (static_cast<int64_t>(blockDim.y) *
                                (static_cast<int64_t>(blockIdx.x) +
                                 static_cast<int64_t>(gridDim.x) *
                                     (static_cast<int64_t>(blockIdx.y) +
                                      static_cast<int64_t>(gridDim.y) *
                                          static_cast<int64_t>(blockIdx.z))) +
                            static_cast<int64_t>(threadIdx.y)) *
                           kLocalBatchSize;

  int64_t mask_fist_idx = (static_cast<int64_t>(blockDim.y) *
                               (static_cast<int64_t>(blockIdx.x) +
                                static_cast<int64_t>(gridDim.x) *
                                    static_cast<int64_t>(blockIdx.z)) +
                           static_cast<int64_t>(threadIdx.y)) *
                          kLocalBatchSize;

  // batch_count might not be a multiple of kLocalBatchSize. Check how
  // many batches have to computed within this WARP.
  int local_batches = batch_count - data_first_idx;
  if (local_batches > kLocalBatchSize) local_batches = kLocalBatchSize;

  // might be many batches per warp. compute the index within the batch
  int local_idx = threadIdx.x;

  int64_t x_offset =
      data_first_idx * key_seq_len + kOneLoadingCounts * local_idx;
  int64_t mask_offset =
      mask_fist_idx * key_seq_len + kOneLoadingCounts * local_idx;
  x_data += x_offset;
  mask_data += mask_offset;
  y_data += x_offset;

  // using float for all inter compute
  float data[kLocalBatchSize][kLocalIterations];
  T temp_data[kOneLoadingCounts];
  MT temp_mask[kOneLoadingCounts];

#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
    int batch_data = (i >= local_batches) ? 0 : key_seq_len;

#pragma unroll
    for (int ii = 0; ii < kLocalIterations; ii += kOneLoadingCounts) {
      int data_index = kOneLoadingCounts * local_idx + ii * warp_size;

      if (data_index < batch_data) {
        int itr_idx = i * key_seq_len + ii * warp_size;

        // efficiently load data from global memory
        load_data(temp_data, x_data + itr_idx);
        load_data(temp_mask, mask_data + itr_idx);

#pragma unroll
        for (int counter = 0; counter < kOneLoadingCounts; ++counter) {
          data[i][ii + counter] = static_cast<float>(temp_data[counter]) +
                                  static_cast<float>(temp_mask[counter]);
        }
      } else {
#pragma unroll
        for (int counter = 0; counter < kOneLoadingCounts; ++counter) {
          data[i][ii + counter] = -std::numeric_limits<float>::infinity();
        }
      }
    }
  }

  // compute max_value
  // max value for each batch for current warp
  float samples_max_value[kLocalBatchSize];
#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
    samples_max_value[i] = data[i][0];
#pragma unroll
    for (int ii = 1; ii < kLocalIterations; ++ii) {
      samples_max_value[i] = (samples_max_value[i] > data[i][ii])
                                 ? samples_max_value[i]
                                 : data[i][ii];
    }
  }
  // max value for each batch for all warp
  warp_reduce<float, kLocalBatchSize, warp_size, MaxOP>(samples_max_value);

  // compute the sum for each batch for current warp
  float samples_sum[kLocalBatchSize]{0.0f};
#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
#pragma unroll
    for (int ii = 0; ii < kLocalIterations; ++ii) {
      data[i][ii] = std::exp((data[i][ii] - samples_max_value[i]));
      samples_sum[i] += data[i][ii];
    }
  }
  // samples_sum for each batch for all warp
  warp_reduce<float, kLocalBatchSize, warp_size, AddOP>(samples_sum);

  // load the result from device back to host
  T samples_out[kOneLoadingCounts];
#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
    if (i >= local_batches) break;
#pragma unroll
    for (int ii = 0; ii < kLocalIterations; ii += kOneLoadingCounts) {
      int idx = kOneLoadingCounts * local_idx + ii * warp_size;
      if (idx < key_seq_len) {
#pragma unroll
        for (int counter = 0; counter < kOneLoadingCounts; ++counter) {
          samples_out[counter] = data[i][ii + counter] / samples_sum[i];
        }
        load_data(y_data + i * key_seq_len + ii * warp_size, samples_out);
      } else {
        break;
      }
    }
  }
}

// T == fp16
// SoftmaxMaskFuseV2GPUKernel is used for larger scale data and has stronger
// universality, but its performance is slightly lower than
// SoftmaxMaskFuseV1GPUKernel

template <typename T, typename MT, int pow2_index>
__global__ void SoftmaxMaskFuseV2GPUKernel(const T* x_data,
                                           const MT* mask_data,
                                           T* y_data,
                                           int64_t batch_count,
                                           int64_t attn_heads,
                                           int64_t query_seqs,
                                           int key_seq_len) {
  // the forward gpu kernel
  constexpr int next_pow2 = 1 << pow2_index;
  constexpr int warp_size = (next_pow2 < WARP_SIZE) ? next_pow2 : WARP_SIZE;
  constexpr int kLocalIterations = std::max(next_pow2 / warp_size, 4);
  constexpr int kLocalBatchSize = (next_pow2 <= 128) ? 2 : 1;
  constexpr int kOneLoadingCounts = 4;

  int64_t blockInGrid = static_cast<int64_t>(blockIdx.x);
  int64_t indexInMaskDim0 = blockInGrid / (attn_heads * query_seqs);
  int64_t indexInMaskDim2 = blockInGrid % query_seqs;

  int64_t data_first_idx = (static_cast<int64_t>(blockDim.y) * blockInGrid +
                            static_cast<int64_t>(threadIdx.y)) *
                           kLocalBatchSize;

  // The original implementation was like this
  // int64_t mask_fist_idx =
  //     (blockDim.y * (blockIdx.x + gridDim.x * blockIdx.z) + threadIdx.y) *
  //     kLocalBatchSize;
  // The mapping relationship is as follows：
  // query_seqs <-> gridDim.x
  // attn_heads <-> gridDim.y
  // indexInMaskDim0 <-> blockIdx.z
  // indexInMaskDim2 <-> blockIdx.x

  int64_t mask_fist_idx =
      (blockDim.y * (indexInMaskDim2 +
                     static_cast<int64_t>(query_seqs) * indexInMaskDim0) +
       threadIdx.y) *
      kLocalBatchSize;

  // batch_count might not be a multiple of kLocalBatchSize. Check how
  // many batches have to computed within this WARP.
  int64_t local_batches = batch_count - data_first_idx;
  if (local_batches > kLocalBatchSize) local_batches = kLocalBatchSize;

  // might be many batches per warp. compute the index within the batch
  int local_idx = threadIdx.x;

  int64_t x_offset =
      data_first_idx * key_seq_len + kOneLoadingCounts * local_idx;
  int64_t mask_offset =
      mask_fist_idx * key_seq_len + kOneLoadingCounts * local_idx;
  x_data += x_offset;
  mask_data += mask_offset;
  y_data += x_offset;

  // using float for all inter compute
  float data[kLocalBatchSize][kLocalIterations];
  T temp_data[kOneLoadingCounts];
  MT temp_mask[kOneLoadingCounts];

#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
    int batch_data = (i >= local_batches) ? 0 : key_seq_len;

#pragma unroll
    for (int ii = 0; ii < kLocalIterations; ii += kOneLoadingCounts) {
      int data_index = kOneLoadingCounts * local_idx + ii * warp_size;

      if (data_index < batch_data) {
        int itr_idx = i * key_seq_len + ii * warp_size;

        // efficiently load data from global memory
        load_data(temp_data, x_data + itr_idx);
        load_data(temp_mask, mask_data + itr_idx);

#pragma unroll
        for (int counter = 0; counter < kOneLoadingCounts; ++counter) {
          data[i][ii + counter] = static_cast<float>(temp_data[counter]) +
                                  static_cast<float>(temp_mask[counter]);
        }
      } else {
#pragma unroll
        for (int counter = 0; counter < kOneLoadingCounts; ++counter) {
          data[i][ii + counter] = -std::numeric_limits<float>::infinity();
        }
      }
    }
  }

  // compute max_value
  // max value for each batch for current warp
  float samples_max_value[kLocalBatchSize];
#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
    samples_max_value[i] = data[i][0];
#pragma unroll
    for (int ii = 1; ii < kLocalIterations; ++ii) {
      samples_max_value[i] = (samples_max_value[i] > data[i][ii])
                                 ? samples_max_value[i]
                                 : data[i][ii];
    }
  }
  // max value for each batch for all warp
  warp_reduce<float, kLocalBatchSize, warp_size, MaxOP>(samples_max_value);

  // compute the sum for each batch for current warp
  float samples_sum[kLocalBatchSize]{0.0f};
#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
#pragma unroll
    for (int ii = 0; ii < kLocalIterations; ++ii) {
      data[i][ii] = std::exp((data[i][ii] - samples_max_value[i]));
      samples_sum[i] += data[i][ii];
    }
  }
  // samples_sum for each batch for all warp
  warp_reduce<float, kLocalBatchSize, warp_size, AddOP>(samples_sum);

  // load the result from device back to host
  T samples_out[kOneLoadingCounts];
#pragma unroll
  for (int i = 0; i < kLocalBatchSize; ++i) {
    if (i >= local_batches) break;
#pragma unroll
    for (int ii = 0; ii < kLocalIterations; ii += kOneLoadingCounts) {
      int idx = kOneLoadingCounts * local_idx + ii * warp_size;
      if (idx < key_seq_len) {
#pragma unroll
        for (int counter = 0; counter < kOneLoadingCounts; ++counter) {
          samples_out[counter] = data[i][ii + counter] / samples_sum[i];
        }
        load_data(y_data + i * key_seq_len + ii * warp_size, samples_out);
      } else {
        break;
      }
    }
  }
}
template <typename T, typename Context>
static void CallSoftmaxMaskGPUKernelV1(const Context& dev_ctx,
                                       const DenseTensor& x,
                                       const DenseTensor& mask,
                                       DenseTensor* out,
                                       int batch_count,
                                       int key_seq_len,
                                       int pow2_index,
                                       const dim3& blocks,
                                       const dim3& threads) {
  auto* x_data = x.data<T>();
  auto* y_data = out->data<T>();
  auto stream = dev_ctx.stream();
  if (mask.dtype() == x.dtype()) {
    auto* mask_data = mask.data<T>();
    switch (pow2_index) {
      case 5:  // 32
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 5)
        break;
      case 6:  // 64
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 6)
        break;
      case 7:  // 128
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 7)
        break;
      case 8:  // 256
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 8)
        break;
      case 9:  // 512
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 9)
        break;
      case 10:  // 1024
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 10)
        break;
      case 11:  // 2048
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 11)
        break;
      case 12:  // 4096
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 12)
        break;
      case 13:  // 8192
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, T, 13)
        break;
      default:
        break;
    }
  } else if (mask.dtype() == phi::DataType::FLOAT32) {
    auto* mask_data = mask.data<float>();
    switch (pow2_index) {
      case 5:  // 32
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 5)
        break;
      case 6:  // 64
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 6)
        break;
      case 7:  // 128
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 7)
        break;
      case 8:  // 256
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 8)
        break;
      case 9:  // 512
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 9)
        break;
      case 10:  // 1024
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 10)
        break;
      case 11:  // 2048
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 11)
        break;
      case 12:  // 4096
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 12)
        break;
      case 13:  // 8192
        LAUNCH_SOFTMAX_MASK_FUSE_V1_GPU_KERNEL(T, float, 13)
        break;
      default:
        break;
    }
  }
}

template <typename T, typename Context>
static void CallSoftmaxMaskGPUKernelV2(const Context& dev_ctx,
                                       const DenseTensor& x,
                                       const DenseTensor& mask,
                                       DenseTensor* out,
                                       int64_t batch_count,
                                       int64_t attn_heads,
                                       int64_t query_seqs,
                                       int key_seq_len,
                                       int pow2_index,
                                       const dim3& blocks,
                                       const dim3& threads) {
  auto* x_data = x.data<T>();
  auto* y_data = out->data<T>();
  auto stream = dev_ctx.stream();
  if (mask.dtype() == x.dtype()) {
    auto* mask_data = mask.data<T>();
    switch (pow2_index) {
      case 5:  // 32
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 5)
        break;
      case 6:  // 64
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 6)
        break;
      case 7:  // 128
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 7)
        break;
      case 8:  // 256
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 8)
        break;
      case 9:  // 512
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 9)
        break;
      case 10:  // 1024
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 10)
        break;
      case 11:  // 2048
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 11)
        break;
      case 12:  // 4096
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 12)
        break;
      case 13:  // 8192
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, T, 13)
        break;
      default:
        break;
    }
  } else if (mask.dtype() == phi::DataType::FLOAT32) {
    auto* mask_data = mask.data<float>();
    switch (pow2_index) {
      case 5:  // 32
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 5)
        break;
      case 6:  // 64
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 6)
        break;
      case 7:  // 128
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 7)
        break;
      case 8:  // 256
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 8)
        break;
      case 9:  // 512
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 9)
        break;
      case 10:  // 1024
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 10)
        break;
      case 11:  // 2048
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 11)
        break;
      case 12:  // 4096
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 12)
        break;
      case 13:  // 8192
        LAUNCH_SOFTMAX_MASK_FUSE_V2_GPU_KERNEL(T, float, 13)
        break;
      default:
        break;
    }
  }
}
// T only supports fp16
// leave as template only for future update
template <typename T, typename Context>
void FusedSoftmaxMaskKernel(const Context& dev_ctx,
                            const DenseTensor& x,
                            const DenseTensor& mask,
                            DenseTensor* out) {
  auto* x_data = x.data<T>();
  auto* y_data = dev_ctx.template Alloc<T>(out);
  if (out && out->numel() == 0) return;

  auto x_dim = x.dims();
  auto mask_dim = mask.dims();
  auto batches = x_dim[0];
  auto attn_heads = x_dim[1];
  auto query_seq_len = x_dim[2];
  auto key_seq_len = x_dim[3];

  PADDLE_ENFORCE_GT(query_seq_len,
                    1,
                    common::errors::InvalidArgument(
                        "Input x's second last dim must be large than 1 but "
                        "received the second last dimension of x is %d",
                        query_seq_len));

  PADDLE_ENFORCE_EQ(key_seq_len >= 32 && key_seq_len < 8192,
                    true,
                    common::errors::InvalidArgument(
                        "Input x's last dim must be between [32, 8192) "
                        "received the last dimension of x is %d",
                        key_seq_len));

  PADDLE_ENFORCE_EQ(mask_dim[1],
                    1,
                    common::errors::InvalidArgument(
                        "Input mask's second dim must be 1 "
                        "received the second dimension of mask is %d",
                        mask_dim[1]));

  // dim of x and mask must be equal
  for (size_t idx = 0; idx < 4; ++idx) {
    if (idx == 1) continue;
    PADDLE_ENFORCE_EQ(
        x_dim[idx],
        mask_dim[idx],
        common::errors::InvalidArgument(
            "Input x's %dth dim should be equal with input mask's %dth dim "
            "but "
            "received the %dth dimension of x and mask are not equal "
            "the %dth dim of x is %d, while the %dth dim of mask is %d.",
            idx,
            idx,
            idx,
            idx,
            x_dim[idx],
            idx,
            mask_dim[idx]));
  }

  int pow2_index = get_pow2(key_seq_len);
  const int next_pow2 = 1 << pow2_index;
  int64_t batch_count = batches * attn_heads * query_seq_len;
  int warp_size = (next_pow2 < WARP_SIZE) ? next_pow2 : WARP_SIZE;
  int batches_per_warp = (next_pow2 <= 128) ? 2 : 1;
  // use 128 threads per block to maximum gpu utilization
  constexpr int threads_per_block = 128;

  int warps_per_block = (threads_per_block / warp_size);
  int batches_per_block = warps_per_block * batches_per_warp;
  PADDLE_ENFORCE_EQ(
      query_seq_len % batches_per_block,
      0,
      common::errors::InvalidArgument(
          "The query seq len (third dim of input X) must can divide the "
          "number of batches per block. The query seq len is %d, while "
          "the number of batches per block is %d.",
          query_seq_len,
          batches_per_block));

  // The original implementation was like this:
  // dim3 blocks(query_seq_len / batches_per_block, attn_heads, batches);
  // If attn_heads or batches beyond 65535, it will cause CUDA error 9
  int64_t total_elements = batch_count * key_seq_len;

  dim3 threads(warp_size, warps_per_block, 1);
  if (total_elements > std::numeric_limits<int>::max() ||
      attn_heads > dev_ctx.GetCUDAMaxGridDimSize()[1] ||
      batches > dev_ctx.GetCUDAMaxGridDimSize()[2]) {
    int64_t total_blocks = batch_count / batches_per_block;
    PADDLE_ENFORCE_LE(total_blocks,
                      dev_ctx.GetCUDAMaxGridDimSize()[0],
                      common::errors::InvalidArgument(
                          "The grid.x of fused_softmax_mask CUDA kernel must "
                          "not exceed the device limit. Expected total_blocks "
                          "<= %d, but received total_blocks = %ld.",
                          dev_ctx.GetCUDAMaxGridDimSize()[0],
                          total_blocks));
    PADDLE_ENFORCE_LE_UINT32_MAX(total_blocks,
                                 "fused_softmax_mask CUDA launch grid.x");
    dim3 blocks(static_cast<uint32_t>(total_blocks));
    int64_t query_seqs = query_seq_len / batches_per_block;
    CallSoftmaxMaskGPUKernelV2<T, Context>(dev_ctx,
                                           x,
                                           mask,
                                           out,
                                           batch_count,
                                           attn_heads,
                                           query_seqs,
                                           static_cast<int>(key_seq_len),
                                           pow2_index,
                                           blocks,
                                           threads);
  } else {
    const int64_t query_seq_blocks = query_seq_len / batches_per_block;
    PADDLE_ENFORCE_LE_UINT32_MAX(query_seq_blocks,
                                 "fused_softmax_mask CUDA launch grid.x");
    PADDLE_ENFORCE_LE_UINT32_MAX(attn_heads,
                                 "fused_softmax_mask CUDA launch grid.y");
    PADDLE_ENFORCE_LE_UINT32_MAX(batches,
                                 "fused_softmax_mask CUDA launch grid.z");
    dim3 blocks(static_cast<uint32_t>(query_seq_blocks),
                static_cast<uint32_t>(attn_heads),
                static_cast<uint32_t>(batches));
    CallSoftmaxMaskGPUKernelV1<T, Context>(dev_ctx,
                                           x,
                                           mask,
                                           out,
                                           static_cast<int>(batch_count),
                                           static_cast<int>(key_seq_len),
                                           pow2_index,
                                           blocks,
                                           threads);
  }
}

}  // namespace fusion
}  // namespace phi

PD_REGISTER_KERNEL(fused_softmax_mask,
                   GPU,
                   ALL_LAYOUT,
                   phi::fusion::FusedSoftmaxMaskKernel,
                   float,
                   phi::float16) {}