/* Copyright (c) 2025 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 <iostream>
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/aligned_vector.h"
#include "paddle/phi/kernels/primitive/kernel_primitives.h"
#include "paddle/phi/kernels/strided_copy_kernel.h"

namespace phi {
bool VerifyStridedCopyThreadConfigurationParameters(const dim3& block,
                                                    const dim3& grid) {
  return block.x <= 1024 && block.y <= 1024 && block.z <= 64 &&
         block.x * block.y * block.z <= 1024 &&
         block.x * block.y * block.z >= 96 && grid.y < 65536 && grid.z < 65536;
}

__device__ bool is_aligned(const void* ptr, size_t alignment) {
  return (reinterpret_cast<uintptr_t>(ptr) % alignment) == 0;
}

template <typename T, size_t N>
__global__ void Contiguous2StridedCaseOneFunc(
    const T* input_data,
    T* out_data,
    phi::Array<int64_t, DDim::kMaxRank + 1> output_stride,
    phi::Array<int64_t, 6> dims,
    const int64_t x_max) {
  int64_t x = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
  if (x < x_max) {
    int64_t input_offset =
        (static_cast<int64_t>(blockIdx.z) * static_cast<int64_t>(gridDim.y) +
         static_cast<int64_t>(blockIdx.y)) *
            x_max +
        x;
    int64_t output_offset = 0;

    int64_t reg_dims[6] = {
        dims[0], dims[1], dims[2], dims[3], dims[4], dims[5]};
    int64_t coordinate[DDim::kMaxRank + 1];

    switch (N) {
      case 1:
        coordinate[0] = x % reg_dims[0];
        break;
      case 2:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        break;
      case 3:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        break;
      case 4:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        break;
      case 5:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        break;
      case 6:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        break;
      case 7:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        coordinate[6] = blockIdx.z % reg_dims[4];
        break;
      case 8:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        coordinate[6] = blockIdx.z % reg_dims[4];
        coordinate[7] = blockIdx.z / reg_dims[4] % reg_dims[5];
        break;
      case 9:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        coordinate[6] = blockIdx.z % reg_dims[4];
        coordinate[7] = blockIdx.z / reg_dims[4] % reg_dims[5];
        coordinate[8] = blockIdx.z / (reg_dims[4] * reg_dims[5]);
        break;
    }

#pragma unroll
    for (int dim = N - 1; dim >= 0; --dim) {
      output_offset += coordinate[N - 1 - dim] * output_stride[dim];
    }

    out_data[output_offset] = input_data[input_offset];
  }
}

template <typename T, size_t N>
__global__ void Contiguous2StridedCaseOneDiffDimFunc(
    const T* input_data,
    T* out_data,
    phi::Array<int64_t, DDim::kMaxRank + 1> output_stride,
    phi::Array<int64_t, 6> dims,
    const int64_t x_max) {
  int64_t x = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
  if (x < x_max) {
    int64_t output_offset = 0;

    int64_t reg_dims[6] = {
        dims[0], dims[1], dims[2], dims[3], dims[4], dims[5]};
    int64_t coordinate[DDim::kMaxRank + 1];

    switch (N) {
      case 1:
        coordinate[0] = x % reg_dims[0];
        break;
      case 2:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        break;
      case 3:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        break;
      case 4:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        break;
      case 5:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        break;
      case 6:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        break;
      case 7:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        coordinate[6] = blockIdx.z % reg_dims[4];
        break;
      case 8:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        coordinate[6] = blockIdx.z % reg_dims[4];
        coordinate[7] = blockIdx.z / reg_dims[4] % reg_dims[5];
        break;
      case 9:
        coordinate[0] = x % reg_dims[0];
        coordinate[1] = x / reg_dims[0] % reg_dims[1];
        coordinate[2] = x / (reg_dims[0] * reg_dims[1]);
        coordinate[3] = blockIdx.y % reg_dims[2];
        coordinate[4] = blockIdx.y / reg_dims[2] % reg_dims[3];
        coordinate[5] = blockIdx.y / (reg_dims[2] * reg_dims[3]);
        coordinate[6] = blockIdx.z % reg_dims[4];
        coordinate[7] = blockIdx.z / reg_dims[4] % reg_dims[5];
        coordinate[8] = blockIdx.z / (reg_dims[4] * reg_dims[5]);
        break;
    }

#pragma unroll
    for (int dim = N - 1; dim >= 0; --dim) {
      output_offset += coordinate[N - 1 - dim] * output_stride[dim];
    }

    out_data[output_offset] = input_data[0];
  }
}

// Check whether "out" is the output of the stride slice.
bool CheckStride(const phi::Array<int64_t, DDim::kMaxRank + 1>& output_stride,
                 const phi::Array<int64_t, DDim::kMaxRank + 1>& dims,
                 int rank,
                 int64_t output_numel) {
  if (output_numel == 0) return true;

  int64_t stride = output_numel;
  int64_t last_stride = 1;
  for (size_t i = 0; i < rank; i++) {
    if (output_stride[i] < last_stride) return true;
    last_stride = output_stride[i];
    stride = stride / dims[i];
    if (output_stride[i] > stride) return false;
  }
  return true;
}

template <typename T, typename Context>
bool LaunchContiguous2StridedCaseOneKernel(
    const Context& dev_ctx,
    const T* input_data,
    T* output_data,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& output_stride,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& dims,
    int rank,
    int64_t output_numel,
    bool diff_dims) {
  if (!CheckStride(output_stride, dims, rank, output_numel)) {
    return false;
  }
  dim3 grid(1, 1, 1), block(1, 1, 1);
  phi::Array<int64_t, 6> cur_dims;
  block.x = 512;

  if (rank >= 1) {
    grid.x = (output_numel + block.x - 1) / block.x;
    cur_dims[0] = dims[rank - 1];
  }

  if (rank >= 2) {
    cur_dims[1] = dims[rank - 2];
  }

  if (rank >= 4) {
    grid.x = (dims[rank - 1] * dims[rank - 2] * dims[rank - 3] + block.x - 1) /
             block.x;
    grid.y = dims[rank - 4];
    cur_dims[2] = dims[rank - 4];
  }

  if (rank >= 5) {
    grid.y = dims[rank - 4] * dims[rank - 5];
    cur_dims[2] = dims[rank - 4];
    cur_dims[3] = dims[rank - 5];
  }

  if (rank >= 6) {
    grid.y = dims[rank - 4] * dims[rank - 5] * dims[rank - 6];
  }

  if (rank >= 7) {
    grid.z = dims[rank - 7];
    cur_dims[4] = dims[rank - 7];
  }

  if (rank >= 8) {
    grid.z = dims[rank - 7] * dims[rank - 8];
    cur_dims[5] = dims[rank - 8];
  }

  if (rank >= 9) {
    grid.z = dims[rank - 7] * dims[rank - 8] * dims[rank - 9];
  }

  if (!VerifyStridedCopyThreadConfigurationParameters(block, grid)) {
    return false;
  }

  if (diff_dims) {
    switch (rank) {
      case 1:
        Contiguous2StridedCaseOneDiffDimFunc<T, 1>
            <<<grid, block, 0, dev_ctx.stream()>>>(input_data,
                                                   output_data,
                                                   output_stride,
                                                   cur_dims,
                                                   dims[rank - 1]);
        break;
      case 2:
        Contiguous2StridedCaseOneDiffDimFunc<T, 2>
            <<<grid, block, 0, dev_ctx.stream()>>>(
                input_data,
                output_data,
                output_stride,
                cur_dims,
                dims[rank - 1] * dims[rank - 2]);
        break;
#define CASE_RANK(__Rk)                                        \
  case __Rk:                                                   \
    Contiguous2StridedCaseOneDiffDimFunc<T, __Rk>              \
        <<<grid, block, 0, dev_ctx.stream()>>>(                \
            input_data,                                        \
            output_data,                                       \
            output_stride,                                     \
            cur_dims,                                          \
            dims[rank - 1] * dims[rank - 2] * dims[rank - 3]); \
    break;
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
      default:
        PADDLE_THROW(common::errors::InvalidArgument(
            "The rank of input should be less than 9, but received %d.", rank));
    }
  } else {
    switch (rank) {
      case 1:
        Contiguous2StridedCaseOneFunc<T, 1>
            <<<grid, block, 0, dev_ctx.stream()>>>(input_data,
                                                   output_data,
                                                   output_stride,
                                                   cur_dims,
                                                   dims[rank - 1]);
        break;
      case 2:
        Contiguous2StridedCaseOneFunc<T, 2>
            <<<grid, block, 0, dev_ctx.stream()>>>(
                input_data,
                output_data,
                output_stride,
                cur_dims,
                dims[rank - 1] * dims[rank - 2]);
        break;
#define CASE_RANK(__Rk)                                        \
  case __Rk:                                                   \
    Contiguous2StridedCaseOneFunc<T, __Rk>                     \
        <<<grid, block, 0, dev_ctx.stream()>>>(                \
            input_data,                                        \
            output_data,                                       \
            output_stride,                                     \
            cur_dims,                                          \
            dims[rank - 1] * dims[rank - 2] * dims[rank - 3]); \
    break;
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
      default:
        PADDLE_THROW(common::errors::InvalidArgument(
            "The rank of input should be less than 9, but received %d.", rank));
    }
  }

  return true;
}

template <typename T, size_t RANK>
__global__ void Contiguous2StridedCaseZeroFunc(
    const T* input_data,
    T* output_data,
    phi::Array<int64_t, DDim::kMaxRank + 1> output_stride) {
  int64_t input_offset =
      (static_cast<int64_t>(blockIdx.z) * static_cast<int64_t>(gridDim.y) *
           static_cast<int64_t>(gridDim.x) +
       static_cast<int64_t>(blockIdx.y) * static_cast<int64_t>(gridDim.x) +
       static_cast<int64_t>(blockIdx.x)) *
          static_cast<int64_t>(blockDim.z) * static_cast<int64_t>(blockDim.y) *
          static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.z) * static_cast<int64_t>(blockDim.y) *
          static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.y) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);
  int64_t output_offset = 0;

  int64_t coordinate[6] = {threadIdx.x,
                           threadIdx.y,
                           threadIdx.z,
                           blockIdx.x,
                           blockIdx.y,
                           blockIdx.z};

#pragma unroll
  for (int dim = RANK - 1; dim >= 0; --dim) {
    output_offset += coordinate[RANK - 1 - dim] * output_stride[dim];
  }

  output_data[output_offset] = input_data[input_offset];
}

template <typename T, size_t RANK>
__global__ void Contiguous2StridedCaseZeroDiffDimFunc(
    const T* input_data,
    T* output_data,
    phi::Array<int64_t, DDim::kMaxRank + 1> output_stride) {
  int64_t output_offset = 0;

  int64_t coordinate[6] = {threadIdx.x,
                           threadIdx.y,
                           threadIdx.z,
                           blockIdx.x,
                           blockIdx.y,
                           blockIdx.z};

#pragma unroll
  for (int dim = RANK - 1; dim >= 0; --dim) {
    output_offset += coordinate[RANK - 1 - dim] * output_stride[dim];
  }

  output_data[output_offset] = input_data[0];
}

template <typename T, typename Context>
bool LaunchContiguous2StridedCaseZeroKernel(
    const Context& dev_ctx,
    const T* input_data,
    T* output_data,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& output_stride,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& dims,
    int rank,
    bool diff_dims) {
  if (rank > 6) {
    return false;
  }

  dim3 grid(1, 1, 1), block(1, 1, 1);

  if (rank >= 1) {
    block.x = dims[rank - 1];
  }

  if (rank >= 2) {
    block.y = dims[rank - 2];
  }

  if (rank >= 3) {
    block.z = dims[rank - 3];
  }

  if (rank >= 4) {
    grid.x = dims[rank - 4];
  }

  if (rank >= 5) {
    grid.y = dims[rank - 5];
  }

  if (rank >= 6) {
    grid.z = dims[rank - 6];
  }

  if (!VerifyStridedCopyThreadConfigurationParameters(block, grid)) {
    return false;
  }

  if (diff_dims) {
    switch (rank) {
#define CASE_RANK(__Rk)                              \
  case __Rk:                                         \
    Contiguous2StridedCaseZeroDiffDimFunc<T, __Rk>   \
        <<<grid, block, 0, dev_ctx.stream()>>>(      \
            input_data, output_data, output_stride); \
    break;
      CASE_RANK(1);
      CASE_RANK(2);
      CASE_RANK(3);
      CASE_RANK(4);
      CASE_RANK(5);
      CASE_RANK(6);
#undef CASE_RANK
      default:
        PADDLE_THROW(common::errors::InvalidArgument(
            "The rank of input should be less than 9, but received %d.", rank));
    }
  } else {
    switch (rank) {
#define CASE_RANK(__Rk)                              \
  case __Rk:                                         \
    Contiguous2StridedCaseZeroFunc<T, __Rk>          \
        <<<grid, block, 0, dev_ctx.stream()>>>(      \
            input_data, output_data, output_stride); \
    break;
      CASE_RANK(1);
      CASE_RANK(2);
      CASE_RANK(3);
      CASE_RANK(4);
      CASE_RANK(5);
      CASE_RANK(6);
#undef CASE_RANK
      default:
        PADDLE_THROW(common::errors::InvalidArgument(
            "The rank of input should be less than 9, but received %d.", rank));
    }
  }

  return true;
}

template <typename T, int VecSize, size_t OUT_RANK>
__global__ void Contiguous2StridedDefaultDiffDimFunc(
    const T* input_data,
    T* output_data,
    Array<int64_t, DDim::kMaxRank + 1> output_stride,
    Array<int64_t, DDim::kMaxRank + 1> dims,
    const int64_t output_numel) {
  int MAX_LOAD_BYTES = VecSize * sizeof(T);
  int64_t gid =
      (static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
       static_cast<int64_t>(threadIdx.x)) *
      VecSize;
  T set_value[VecSize];
#pragma unroll
  for (int i = 0; i < VecSize; i++) {
    set_value[i] = input_data[0];
  }
#pragma unroll
  for (int64_t i = gid; i < output_numel;
       i += blockDim.x * gridDim.x * VecSize) {
    int64_t output_offset = 0;
    int64_t index_tmp = i;

    for (int dim = OUT_RANK - 1; dim >= 0; --dim) {
      output_offset += (index_tmp % dims[dim]) * output_stride[dim];
      index_tmp = index_tmp / dims[dim];
    }
    if (is_aligned(&output_data[output_offset], MAX_LOAD_BYTES)) {
      using VecType = kps::details::VectorType<T, VecSize>;
      const VecType* src = reinterpret_cast<const VecType*>(&set_value[0]);
      VecType* dst = reinterpret_cast<VecType*>(&output_data[output_offset]);
      *dst = *src;
    } else {
      for (int j = 0; j < VecSize; j++) {
        output_data[output_offset + j] = input_data[0];
      }
    }
  }
}

template <typename T, int VecSize, size_t OUT_RANK>
__global__ void Contiguous2StridedDefaultFunc(
    const T* input_data,
    T* output_data,
    Array<int64_t, DDim::kMaxRank + 1> output_stride,
    Array<int64_t, DDim::kMaxRank + 1> dims,
    const int64_t output_numel) {
  int MAX_LOAD_BYTES = VecSize * sizeof(T);
  int64_t gid =
      (static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
       static_cast<int64_t>(threadIdx.x)) *
      VecSize;
#pragma unroll
  for (int64_t i = gid; i < output_numel;
       i += blockDim.x * gridDim.x * VecSize) {
    int64_t output_offset = 0;
    int64_t index_tmp = i;
    for (int dim = OUT_RANK - 1; dim >= 0; --dim) {
      output_offset += (index_tmp % dims[dim]) * output_stride[dim];
      index_tmp = index_tmp / dims[dim];
    }
    if (is_aligned(&output_data[output_offset], MAX_LOAD_BYTES)) {
      using VecType = kps::details::VectorType<T, VecSize>;
      const VecType* src = reinterpret_cast<const VecType*>(&input_data[i]);
      VecType* dst = reinterpret_cast<VecType*>(&output_data[output_offset]);
      *dst = *src;
    } else {
      for (int j = 0; j < VecSize; j++) {
        output_data[output_offset + j] = input_data[i + j];
      }
    }
  }
}

template <typename T, int VecSize, size_t OUT_RANK>
__global__ void Contiguous2StridedExpandDefaultFunc(
    const T* input_data,
    T* output_data,
    Array<int64_t, DDim::kMaxRank + 1> output_stride,
    Array<int64_t, DDim::kMaxRank + 1> dims,
    const int64_t input_numel,
    const int64_t output_numel) {
  int MAX_LOAD_BYTES = VecSize * sizeof(T);
  int64_t gid =
      (static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
       static_cast<int64_t>(threadIdx.x)) *
      VecSize;
#pragma unroll
  for (int64_t i = gid; i < output_numel;
       i += blockDim.x * gridDim.x * VecSize) {
    int64_t output_offset = 0;
    int64_t index_tmp = i;
    int64_t input_idx = i % input_numel;
    for (int dim = OUT_RANK - 1; dim >= 0; --dim) {
      output_offset += (index_tmp % dims[dim]) * output_stride[dim];
      index_tmp = index_tmp / dims[dim];
    }
    if (is_aligned(&output_data[output_offset], MAX_LOAD_BYTES)) {
      using VecType = kps::details::VectorType<T, VecSize>;
      const VecType* src =
          reinterpret_cast<const VecType*>(&input_data[input_idx]);
      VecType* dst = reinterpret_cast<VecType*>(&output_data[output_offset]);
      *dst = *src;
    } else {
      for (int j = 0; j < VecSize; j++) {
        output_data[output_offset + j] = input_data[input_idx + j];
      }
    }
  }
}

template <typename T, typename Context, int VecSize>
void LaunchContiguous2StridedDefaultKernel(
    const Context& dev_ctx,
    const T* input_data,
    T* output_data,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& output_stride,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& dims,
    int rank,
    int64_t input_numel,
    int64_t output_numel,
    bool diff_dims) {
  constexpr int loop_count = 4;
  auto config = phi::backends::gpu::GetGpuLaunchConfig1D(
      dev_ctx, output_numel, VecSize * loop_count);
  auto& grid = config.block_per_grid;
  auto& block = config.thread_per_block;
  if (diff_dims) {
    if (VecSize == 8) {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultDiffDimFunc<T, 8, __Rk>                     \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break
        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK

        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    } else if (VecSize == 4) {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultDiffDimFunc<T, 4, __Rk>                     \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break
        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK

        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }

    } else if (VecSize == 2) {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultDiffDimFunc<T, 2, __Rk>                     \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break
        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK

        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    } else {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultDiffDimFunc<T, 1, __Rk>                     \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break
        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK

        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    }
  } else if (input_numel != output_numel) {
    if (VecSize == 8) {
      switch (rank) {
#define CASE_RANK(__Rk)                                       \
  case __Rk:                                                  \
    Contiguous2StridedExpandDefaultFunc<T, 8, __Rk>           \
        <<<grid, block, 0, dev_ctx.stream()>>>(input_data,    \
                                               output_data,   \
                                               output_stride, \
                                               dims,          \
                                               input_numel,   \
                                               output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }

    } else if (VecSize == 4) {
      switch (rank) {
#define CASE_RANK(__Rk)                                       \
  case __Rk:                                                  \
    Contiguous2StridedExpandDefaultFunc<T, 4, __Rk>           \
        <<<grid, block, 0, dev_ctx.stream()>>>(input_data,    \
                                               output_data,   \
                                               output_stride, \
                                               dims,          \
                                               input_numel,   \
                                               output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    } else if (VecSize == 2) {
      switch (rank) {
#define CASE_RANK(__Rk)                                       \
  case __Rk:                                                  \
    Contiguous2StridedExpandDefaultFunc<T, 2, __Rk>           \
        <<<grid, block, 0, dev_ctx.stream()>>>(input_data,    \
                                               output_data,   \
                                               output_stride, \
                                               dims,          \
                                               input_numel,   \
                                               output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    } else {
      switch (rank) {
#define CASE_RANK(__Rk)                                       \
  case __Rk:                                                  \
    Contiguous2StridedExpandDefaultFunc<T, 1, __Rk>           \
        <<<grid, block, 0, dev_ctx.stream()>>>(input_data,    \
                                               output_data,   \
                                               output_stride, \
                                               dims,          \
                                               input_numel,   \
                                               output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    }
  } else {
    if (VecSize == 8) {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultFunc<T, 8, __Rk>                            \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }

    } else if (VecSize == 4) {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultFunc<T, 4, __Rk>                            \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    } else if (VecSize == 2) {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultFunc<T, 2, __Rk>                            \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    } else {
      switch (rank) {
#define CASE_RANK(__Rk)                                                  \
  case __Rk:                                                             \
    Contiguous2StridedDefaultFunc<T, 1, __Rk>                            \
        <<<grid, block, 0, dev_ctx.stream()>>>(                          \
            input_data, output_data, output_stride, dims, output_numel); \
    break

        CASE_RANK(1);
        CASE_RANK(2);
        CASE_RANK(3);
        CASE_RANK(4);
        CASE_RANK(5);
        CASE_RANK(6);
        CASE_RANK(7);
        CASE_RANK(8);
        CASE_RANK(9);
#undef CASE_RANK
        default:
          PADDLE_THROW(common::errors::InvalidArgument(
              "The rank of input should be less than 9, but received %d.",
              rank));
      }
    }
  }
}

template <typename T, typename Context, int VecSize>
void StrideCopyDiffDimKernel(
    const Context& dev_ctx,
    const T* input_data,
    T* output_data,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& output_stride,
    const phi::Array<int64_t, DDim::kMaxRank + 1>& output_dims,
    int rank,
    int64_t input_numel,
    int64_t output_numel) {
  if (LaunchContiguous2StridedCaseZeroKernel<T, Context>(dev_ctx,
                                                         input_data,
                                                         output_data,
                                                         output_stride,
                                                         output_dims,
                                                         rank,
                                                         true)) {
  } else if (LaunchContiguous2StridedCaseOneKernel<T, Context>(dev_ctx,
                                                               input_data,
                                                               output_data,
                                                               output_stride,
                                                               output_dims,
                                                               rank,
                                                               output_numel,
                                                               true)) {
  } else {
    switch (VecSize) {
#define CASE_VECSIZE(__Sz)                                                 \
  case __Sz:                                                               \
    LaunchContiguous2StridedDefaultKernel<T, Context, __Sz>(dev_ctx,       \
                                                            input_data,    \
                                                            output_data,   \
                                                            output_stride, \
                                                            output_dims,   \
                                                            rank,          \
                                                            input_numel,   \
                                                            output_numel,  \
                                                            true);         \
    break;
      CASE_VECSIZE(1);
      CASE_VECSIZE(2);
      CASE_VECSIZE(4);
      CASE_VECSIZE(8);
#undef CASE_VECSIZE
      default:
        PADDLE_THROW(common::errors::InvalidArgument(
            "unsurport vecsize %d for LaunchContiguous2StridedDefaultKernel",
            VecSize));
    }
  }
}

}  // namespace phi