wehub-resource-sync
1051 lines
38 KiB
Plaintext
1051 lines
38 KiB
Plaintext
/* Copyright (c) 2023 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/kernels/strided_copy_kernel.h"
|
|
|
|
#include <cstdint>
|
|
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/empty_kernel.h"
|
|
#include "paddle/phi/kernels/expand_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/slice_utils.h"
|
|
#include "paddle/phi/kernels/funcs/strided_copy_kernel.cu.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, size_t RANK>
|
|
__global__ void StridedCopyCaseZeroFunc(
|
|
const T* input_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> input_stride,
|
|
T* output_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> output_stride) {
|
|
int64_t input_offset = 0;
|
|
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) {
|
|
input_offset += coordinate[RANK - 1 - dim] * input_stride[dim];
|
|
output_offset += coordinate[RANK - 1 - dim] * output_stride[dim];
|
|
}
|
|
|
|
output_data[output_offset] = input_data[input_offset];
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
bool LaunchStridedCopyCaseZeroKernel(
|
|
const Context& dev_ctx,
|
|
const T* input_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& input_stride,
|
|
T* output_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& output_stride,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& dims,
|
|
int rank) {
|
|
if (rank > 6) {
|
|
return false;
|
|
}
|
|
|
|
dim3 grid(1, 1, 1), block(1, 1, 1);
|
|
|
|
if (rank >= 1) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 1], "strided_copy block.x");
|
|
block.x = static_cast<uint32_t>(dims[rank - 1]);
|
|
}
|
|
|
|
if (rank >= 2) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 2], "strided_copy block.y");
|
|
block.y = static_cast<uint32_t>(dims[rank - 2]);
|
|
}
|
|
|
|
if (rank >= 3) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 3], "strided_copy block.z");
|
|
block.z = static_cast<uint32_t>(dims[rank - 3]);
|
|
}
|
|
|
|
if (rank >= 4) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 4], "strided_copy grid.x");
|
|
grid.x = static_cast<uint32_t>(dims[rank - 4]);
|
|
}
|
|
|
|
if (rank >= 5) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 5], "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(dims[rank - 5]);
|
|
}
|
|
|
|
if (rank >= 6) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 6], "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(dims[rank - 6]);
|
|
}
|
|
|
|
if (!VerifyStridedCopyThreadConfigurationParameters(block, grid)) {
|
|
return false;
|
|
}
|
|
|
|
switch (rank) {
|
|
case 1:
|
|
StridedCopyCaseZeroFunc<T, 1><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data, input_stride, output_data, output_stride);
|
|
break;
|
|
case 2:
|
|
StridedCopyCaseZeroFunc<T, 2><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data, input_stride, output_data, output_stride);
|
|
break;
|
|
case 3:
|
|
StridedCopyCaseZeroFunc<T, 3><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data, input_stride, output_data, output_stride);
|
|
break;
|
|
case 4:
|
|
StridedCopyCaseZeroFunc<T, 4><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data, input_stride, output_data, output_stride);
|
|
break;
|
|
case 5:
|
|
StridedCopyCaseZeroFunc<T, 5><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data, input_stride, output_data, output_stride);
|
|
break;
|
|
case 6:
|
|
StridedCopyCaseZeroFunc<T, 6><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data, input_stride, output_data, output_stride);
|
|
break;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
template <typename T, size_t N>
|
|
__global__ void StridedCopyCaseOneFunc(
|
|
const T* input_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> input_stride,
|
|
T* out_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> output_stride,
|
|
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 = 0;
|
|
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) {
|
|
input_offset += coordinate[N - 1 - dim] * input_stride[dim];
|
|
output_offset += coordinate[N - 1 - dim] * output_stride[dim];
|
|
}
|
|
|
|
out_data[output_offset] = input_data[input_offset];
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
bool LaunchStridedCopyCaseOneKernel(
|
|
const Context& dev_ctx,
|
|
const T* input_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& input_stride,
|
|
T* output_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& output_stride,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& dims,
|
|
int rank,
|
|
int64_t numel) {
|
|
dim3 grid(1, 1, 1), block(1, 1, 1);
|
|
Array<int64_t, 6> cur_dims;
|
|
block.x = 512;
|
|
|
|
if (rank >= 1) {
|
|
int64_t grid_x = (numel + static_cast<int64_t>(block.x) - 1) /
|
|
static_cast<int64_t>(block.x);
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_x, "strided_copy grid.x");
|
|
grid.x = static_cast<uint32_t>(grid_x);
|
|
cur_dims[0] = dims[rank - 1];
|
|
}
|
|
|
|
if (rank >= 2) {
|
|
cur_dims[1] = dims[rank - 2];
|
|
}
|
|
|
|
if (rank >= 4) {
|
|
int64_t grid_x = (dims[rank - 1] * dims[rank - 2] * dims[rank - 3] +
|
|
static_cast<int64_t>(block.x) - 1) /
|
|
static_cast<int64_t>(block.x);
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_x, "strided_copy grid.x");
|
|
grid.x = static_cast<uint32_t>(grid_x);
|
|
}
|
|
|
|
if (rank >= 6) {
|
|
int64_t grid_y = dims[rank - 4] * dims[rank - 5] * dims[rank - 6];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_y, "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(grid_y);
|
|
cur_dims[2] = dims[rank - 4];
|
|
cur_dims[3] = dims[rank - 5];
|
|
} else if (rank >= 5) {
|
|
int64_t grid_y = dims[rank - 4] * dims[rank - 5];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_y, "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(grid_y);
|
|
cur_dims[2] = dims[rank - 4];
|
|
cur_dims[3] = dims[rank - 5];
|
|
} else if (rank >= 4) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 4], "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(dims[rank - 4]);
|
|
cur_dims[2] = dims[rank - 4];
|
|
}
|
|
|
|
if (rank >= 9) {
|
|
int64_t grid_z = dims[rank - 7] * dims[rank - 8] * dims[rank - 9];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_z, "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(grid_z);
|
|
cur_dims[4] = dims[rank - 7];
|
|
cur_dims[5] = dims[rank - 8];
|
|
} else if (rank >= 8) {
|
|
int64_t grid_z = dims[rank - 7] * dims[rank - 8];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_z, "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(grid_z);
|
|
cur_dims[4] = dims[rank - 7];
|
|
cur_dims[5] = dims[rank - 8];
|
|
} else if (rank >= 7) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 7], "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(dims[rank - 7]);
|
|
cur_dims[4] = dims[rank - 7];
|
|
}
|
|
|
|
if (!VerifyStridedCopyThreadConfigurationParameters(block, grid)) {
|
|
return false;
|
|
}
|
|
|
|
switch (rank) {
|
|
case 1:
|
|
StridedCopyCaseOneFunc<T, 1>
|
|
<<<grid, block, 0, dev_ctx.stream()>>>(input_data,
|
|
input_stride,
|
|
output_data,
|
|
output_stride,
|
|
cur_dims,
|
|
dims[rank - 1]);
|
|
break;
|
|
case 2:
|
|
StridedCopyCaseOneFunc<T, 2><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data,
|
|
input_stride,
|
|
output_data,
|
|
output_stride,
|
|
cur_dims,
|
|
dims[rank - 1] * dims[rank - 2]);
|
|
break;
|
|
#define CASE_RANK(__Rk) \
|
|
case __Rk: \
|
|
StridedCopyCaseOneFunc<T, __Rk><<<grid, block, 0, dev_ctx.stream()>>>( \
|
|
input_data, \
|
|
input_stride, \
|
|
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 StridedCopyDefaultFunc(
|
|
const T* input_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> input_stride,
|
|
T* output_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> output_stride,
|
|
Array<int64_t, DDim::kMaxRank + 1> dims,
|
|
const int64_t numel) {
|
|
int64_t gid =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
#pragma unroll
|
|
for (int64_t i = gid; i < numel; i += static_cast<int64_t>(blockDim.x) *
|
|
static_cast<int64_t>(gridDim.x)) {
|
|
int64_t input_offset = 0;
|
|
int64_t index_tmp = i;
|
|
#pragma unroll
|
|
for (int dim = RANK - 1; dim >= 0; --dim) {
|
|
input_offset += (index_tmp % dims[dim]) * input_stride[dim];
|
|
index_tmp = index_tmp / dims[dim];
|
|
}
|
|
int64_t output_offset = 0;
|
|
index_tmp = i;
|
|
#pragma unroll
|
|
for (int dim = RANK - 1; dim >= 0; --dim) {
|
|
output_offset += (index_tmp % dims[dim]) * output_stride[dim];
|
|
index_tmp = index_tmp / dims[dim];
|
|
}
|
|
output_data[output_offset] = input_data[input_offset];
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void LaunchStridedCopyDefaultKernel(
|
|
const Context& dev_ctx,
|
|
const T* input_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& input_stride,
|
|
T* output_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& output_stride,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& dims,
|
|
int rank,
|
|
int64_t numel) {
|
|
constexpr uint32_t block = 512;
|
|
const int64_t grid_x = (numel + block - 1) / block;
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_x, "strided_copy grid.x");
|
|
const dim3 grid(static_cast<uint32_t>(grid_x));
|
|
const dim3 block_dim(block);
|
|
|
|
switch (rank) {
|
|
#define CASE_RANK(__Rk) \
|
|
case __Rk: \
|
|
StridedCopyDefaultFunc<T, __Rk><<<grid, block_dim, 0, dev_ctx.stream()>>>( \
|
|
input_data, input_stride, output_data, output_stride, dims, 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, size_t RANK>
|
|
__global__ void Strided2ContiguousCaseZeroFunc(
|
|
const T* input_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> input_stride,
|
|
T* output_data) {
|
|
int64_t input_offset = 0;
|
|
int64_t output_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 coordinate[6] = {threadIdx.x,
|
|
threadIdx.y,
|
|
threadIdx.z,
|
|
blockIdx.x,
|
|
blockIdx.y,
|
|
blockIdx.z};
|
|
|
|
#pragma unroll
|
|
for (int dim = RANK - 1; dim >= 0; --dim) {
|
|
input_offset += coordinate[RANK - 1 - dim] * input_stride[dim];
|
|
}
|
|
|
|
output_data[output_offset] = input_data[input_offset];
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
bool LaunchStrided2ContiguousCaseZeroKernel(
|
|
const Context& dev_ctx,
|
|
const T* input_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& input_stride,
|
|
T* output_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& dims,
|
|
int rank) {
|
|
if (rank > 6) {
|
|
return false;
|
|
}
|
|
|
|
dim3 grid(1, 1, 1), block(1, 1, 1);
|
|
|
|
if (rank >= 1) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 1], "strided_copy block.x");
|
|
block.x = static_cast<uint32_t>(dims[rank - 1]);
|
|
}
|
|
|
|
if (rank >= 2) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 2], "strided_copy block.y");
|
|
block.y = static_cast<uint32_t>(dims[rank - 2]);
|
|
}
|
|
|
|
if (rank >= 3) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 3], "strided_copy block.z");
|
|
block.z = static_cast<uint32_t>(dims[rank - 3]);
|
|
}
|
|
|
|
if (rank >= 4) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 4], "strided_copy grid.x");
|
|
grid.x = static_cast<uint32_t>(dims[rank - 4]);
|
|
}
|
|
|
|
if (rank >= 5) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 5], "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(dims[rank - 5]);
|
|
}
|
|
|
|
if (rank >= 6) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 6], "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(dims[rank - 6]);
|
|
}
|
|
|
|
if (!VerifyStridedCopyThreadConfigurationParameters(block, grid)) {
|
|
return false;
|
|
}
|
|
|
|
switch (rank) {
|
|
#define CASE_RANK(__Rk) \
|
|
case __Rk: \
|
|
Strided2ContiguousCaseZeroFunc<T, __Rk> \
|
|
<<<grid, block, 0, dev_ctx.stream()>>>( \
|
|
input_data, input_stride, output_data); \
|
|
break
|
|
CASE_RANK(1);
|
|
CASE_RANK(2);
|
|
CASE_RANK(3);
|
|
CASE_RANK(4);
|
|
CASE_RANK(5);
|
|
CASE_RANK(6);
|
|
#undef CASE_RANK
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
template <typename T, size_t N>
|
|
__global__ void Strided2ContiguousCaseOneFunc(
|
|
const T* input_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> input_stride,
|
|
T* out_data,
|
|
Array<int64_t, 6> dims,
|
|
const int64_t x_max) {
|
|
int64_t x =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
if (x < x_max) {
|
|
int64_t input_offset = 0;
|
|
int64_t output_offset =
|
|
(static_cast<int64_t>(blockIdx.z) * static_cast<int64_t>(gridDim.y) +
|
|
static_cast<int64_t>(blockIdx.y)) *
|
|
x_max +
|
|
x;
|
|
|
|
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) {
|
|
input_offset += coordinate[N - 1 - dim] * input_stride[dim];
|
|
}
|
|
|
|
out_data[output_offset] = input_data[input_offset];
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
bool LaunchStrided2ContiguousCaseOneKernel(
|
|
const Context& dev_ctx,
|
|
const T* input_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& input_stride,
|
|
T* output_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& dims,
|
|
int rank,
|
|
int64_t numel) {
|
|
dim3 grid(1, 1, 1), block(1, 1, 1);
|
|
Array<int64_t, 6> cur_dims;
|
|
block.x = 512;
|
|
|
|
if (rank >= 1) {
|
|
int64_t grid_x = (numel + static_cast<int64_t>(block.x) - 1) /
|
|
static_cast<int64_t>(block.x);
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_x, "strided_copy grid.x");
|
|
grid.x = static_cast<uint32_t>(grid_x);
|
|
cur_dims[0] = dims[rank - 1];
|
|
}
|
|
|
|
if (rank >= 2) {
|
|
cur_dims[1] = dims[rank - 2];
|
|
}
|
|
|
|
if (rank >= 4) {
|
|
int64_t grid_x = (dims[rank - 1] * dims[rank - 2] * dims[rank - 3] +
|
|
static_cast<int64_t>(block.x) - 1) /
|
|
static_cast<int64_t>(block.x);
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_x, "strided_copy grid.x");
|
|
grid.x = static_cast<uint32_t>(grid_x);
|
|
}
|
|
|
|
if (rank >= 6) {
|
|
int64_t grid_y = dims[rank - 4] * dims[rank - 5] * dims[rank - 6];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_y, "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(grid_y);
|
|
cur_dims[2] = dims[rank - 4];
|
|
cur_dims[3] = dims[rank - 5];
|
|
} else if (rank >= 5) {
|
|
int64_t grid_y = dims[rank - 4] * dims[rank - 5];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_y, "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(grid_y);
|
|
cur_dims[2] = dims[rank - 4];
|
|
cur_dims[3] = dims[rank - 5];
|
|
} else if (rank >= 4) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 4], "strided_copy grid.y");
|
|
grid.y = static_cast<uint32_t>(dims[rank - 4]);
|
|
cur_dims[2] = dims[rank - 4];
|
|
}
|
|
|
|
if (rank >= 9) {
|
|
int64_t grid_z = dims[rank - 7] * dims[rank - 8] * dims[rank - 9];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_z, "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(grid_z);
|
|
cur_dims[4] = dims[rank - 7];
|
|
cur_dims[5] = dims[rank - 8];
|
|
} else if (rank >= 8) {
|
|
int64_t grid_z = dims[rank - 7] * dims[rank - 8];
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_z, "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(grid_z);
|
|
cur_dims[4] = dims[rank - 7];
|
|
cur_dims[5] = dims[rank - 8];
|
|
} else if (rank >= 7) {
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(dims[rank - 7], "strided_copy grid.z");
|
|
grid.z = static_cast<uint32_t>(dims[rank - 7]);
|
|
cur_dims[4] = dims[rank - 7];
|
|
}
|
|
|
|
if (!VerifyStridedCopyThreadConfigurationParameters(block, grid)) {
|
|
return false;
|
|
}
|
|
|
|
switch (rank) {
|
|
case 1:
|
|
Strided2ContiguousCaseOneFunc<T, 1><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data, input_stride, output_data, cur_dims, dims[rank - 1]);
|
|
break;
|
|
case 2:
|
|
Strided2ContiguousCaseOneFunc<T, 2><<<grid, block, 0, dev_ctx.stream()>>>(
|
|
input_data,
|
|
input_stride,
|
|
output_data,
|
|
cur_dims,
|
|
dims[rank - 1] * dims[rank - 2]);
|
|
break;
|
|
#define CASE_RANK(__Rk) \
|
|
case __Rk: \
|
|
Strided2ContiguousCaseOneFunc<T, __Rk> \
|
|
<<<grid, block, 0, dev_ctx.stream()>>>( \
|
|
input_data, \
|
|
input_stride, \
|
|
output_data, \
|
|
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 IN_RANK>
|
|
__global__ void Strided2ContiguousDefaultFunc(
|
|
const T* input_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> input_stride,
|
|
T* output_data,
|
|
Array<int64_t, DDim::kMaxRank + 1> dims,
|
|
const int64_t numel) {
|
|
int64_t gid =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
#pragma unroll
|
|
for (int64_t i = gid; i < numel; i += static_cast<int64_t>(blockDim.x) *
|
|
static_cast<int64_t>(gridDim.x)) {
|
|
int64_t input_offset = 0;
|
|
int64_t index_tmp = i;
|
|
#pragma unroll
|
|
for (int dim = IN_RANK - 1; dim >= 0; --dim) {
|
|
input_offset += (index_tmp % dims[dim]) * input_stride[dim];
|
|
index_tmp = index_tmp / dims[dim];
|
|
}
|
|
output_data[i] = input_data[input_offset];
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void LaunchStrided2ContiguousDefaultKernel(
|
|
const Context& dev_ctx,
|
|
const T* input_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& input_stride,
|
|
T* output_data,
|
|
const Array<int64_t, DDim::kMaxRank + 1>& dims,
|
|
int rank,
|
|
int64_t numel) {
|
|
constexpr uint32_t block = 512;
|
|
const int64_t grid_x = (numel + block - 1) / block;
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_x, "strided_copy grid.x");
|
|
const dim3 grid(static_cast<uint32_t>(grid_x));
|
|
const dim3 block_dim(block);
|
|
|
|
switch (rank) {
|
|
#define CASE_RANK(__Rk) \
|
|
case __Rk: \
|
|
Strided2ContiguousDefaultFunc<T, __Rk> \
|
|
<<<grid, block_dim, 0, dev_ctx.stream()>>>( \
|
|
input_data, input_stride, output_data, dims, 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>
|
|
void StridedCopyKernel(const Context& dev_ctx,
|
|
const DenseTensor& input,
|
|
const std::vector<int64_t>& dims,
|
|
const std::vector<int64_t>& out_stride,
|
|
int64_t offset,
|
|
DenseTensor* out) {
|
|
DenseTensorMeta meta = input.meta();
|
|
meta.strides = make_ddim(out_stride);
|
|
meta.dims = make_ddim(dims);
|
|
meta.offset = offset;
|
|
out->set_meta(meta);
|
|
int rank = out->dims().size();
|
|
int64_t input_numel = input.numel();
|
|
int64_t output_numel = out->numel();
|
|
T* output_data = out->data<T>();
|
|
PADDLE_ENFORCE_NOT_NULL(output_data,
|
|
common::errors::InvalidArgument(
|
|
"StridedCopyKernel's out tensor must complete "
|
|
"mutable data before call kernel."));
|
|
Array<int64_t, DDim::kMaxRank + 1> output_dims;
|
|
Array<int64_t, DDim::kMaxRank + 1> output_stride;
|
|
for (int i = 0; i < meta.dims.size(); i++) {
|
|
output_dims[i] = meta.dims[i];
|
|
output_stride[i] = meta.strides[i];
|
|
}
|
|
|
|
const T* input_data = input.data<T>();
|
|
|
|
// count vecsize
|
|
int VecSize = 8;
|
|
VecSize = std::min(GetVectorizedSize<T>(input_data), VecSize);
|
|
VecSize = std::min(GetVectorizedSize<T>(output_data), VecSize);
|
|
while (VecSize > 1 && output_numel % VecSize != 0) {
|
|
VecSize /= 2;
|
|
}
|
|
if (input_numel != 1 && input_numel != output_numel) {
|
|
while (VecSize > 1 && input_numel % VecSize != 0) {
|
|
VecSize /= 2;
|
|
}
|
|
}
|
|
while (VecSize > 1 && output_dims[meta.dims.size() - 1] % VecSize != 0) {
|
|
VecSize /= 2;
|
|
}
|
|
if (output_stride[meta.dims.size() - 1] != 1) {
|
|
VecSize = 1;
|
|
}
|
|
if (input.dims() != out->dims()) {
|
|
if (input_numel == 1) {
|
|
switch (VecSize) {
|
|
#define CASE_VECSIZE(__Sz) \
|
|
case __Sz: \
|
|
StrideCopyDiffDimKernel<T, Context, __Sz>(dev_ctx, \
|
|
input_data, \
|
|
output_data, \
|
|
output_stride, \
|
|
output_dims, \
|
|
rank, \
|
|
input_numel, \
|
|
output_numel); \
|
|
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 StrideCopyDiffDimKernel", VecSize));
|
|
}
|
|
return;
|
|
} else {
|
|
bool can_expand = funcs::CheckIsLastDimsMatch(input.dims(), out->dims());
|
|
if (can_expand && input.meta().is_contiguous()) {
|
|
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, \
|
|
false); \
|
|
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));
|
|
}
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(input.dims(),
|
|
out->dims(),
|
|
common::errors::InvalidArgument(
|
|
"Input shape(%s) must be equal with out shape(%s).",
|
|
input.dims(),
|
|
out->dims()));
|
|
|
|
PADDLE_ENFORCE_EQ(input_numel,
|
|
output_numel,
|
|
common::errors::InvalidArgument(
|
|
"Input numel(%d) must be equal with out numel(%d).",
|
|
input_numel,
|
|
output_numel));
|
|
|
|
Array<int64_t, DDim::kMaxRank + 1> input_dims;
|
|
Array<int64_t, DDim::kMaxRank + 1> input_stride;
|
|
for (int i = 0; i < input.dims().size(); i++) {
|
|
input_dims[i] = input.dims()[i];
|
|
input_stride[i] = input.strides()[i];
|
|
}
|
|
|
|
if (output_numel == 1) {
|
|
#ifdef PADDLE_WITH_HIP
|
|
hipMemcpy(output_data,
|
|
input_data,
|
|
SizeOf(input.dtype()),
|
|
hipMemcpyDeviceToDevice);
|
|
#else
|
|
cudaMemcpy(output_data,
|
|
input_data,
|
|
SizeOf(input.dtype()),
|
|
cudaMemcpyDeviceToDevice);
|
|
#endif
|
|
|
|
return;
|
|
}
|
|
|
|
if (input.meta().is_contiguous()) {
|
|
if (LaunchContiguous2StridedCaseZeroKernel<T, Context>(dev_ctx,
|
|
input_data,
|
|
output_data,
|
|
output_stride,
|
|
output_dims,
|
|
rank,
|
|
false)) {
|
|
} else if (LaunchContiguous2StridedCaseOneKernel<T, Context>(dev_ctx,
|
|
input_data,
|
|
output_data,
|
|
output_stride,
|
|
output_dims,
|
|
rank,
|
|
output_numel,
|
|
false)) {
|
|
} 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, \
|
|
false); \
|
|
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 StrideCopyKernel", VecSize));
|
|
}
|
|
}
|
|
} else if (out->meta().is_contiguous()) {
|
|
if (LaunchStrided2ContiguousCaseZeroKernel<T, Context>(
|
|
dev_ctx, input_data, input_stride, output_data, input_dims, rank)) {
|
|
} else if (LaunchStrided2ContiguousCaseOneKernel<T, Context>(
|
|
dev_ctx,
|
|
input_data,
|
|
input_stride,
|
|
output_data,
|
|
input_dims,
|
|
rank,
|
|
output_numel)) {
|
|
} else {
|
|
LaunchStrided2ContiguousDefaultKernel<T, Context>(dev_ctx,
|
|
input_data,
|
|
input_stride,
|
|
output_data,
|
|
input_dims,
|
|
rank,
|
|
output_numel);
|
|
}
|
|
} else {
|
|
if (LaunchStridedCopyCaseZeroKernel<T, Context>(dev_ctx,
|
|
input_data,
|
|
input_stride,
|
|
output_data,
|
|
output_stride,
|
|
input_dims,
|
|
rank)) {
|
|
} else if (LaunchStridedCopyCaseOneKernel<T, Context>(dev_ctx,
|
|
input_data,
|
|
input_stride,
|
|
output_data,
|
|
output_stride,
|
|
input_dims,
|
|
rank,
|
|
output_numel)) {
|
|
} else {
|
|
LaunchStridedCopyDefaultKernel<T, Context>(dev_ctx,
|
|
input_data,
|
|
input_stride,
|
|
output_data,
|
|
output_stride,
|
|
input_dims,
|
|
rank,
|
|
output_numel);
|
|
}
|
|
}
|
|
}
|
|
#ifdef _WIN32
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(bool, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(uint8_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(uint16_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(uint32_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(uint64_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(int8_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(int16_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(int32_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(int64_t, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(float, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(double, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(dtype::float16, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(dtype::bfloat16, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(dtype::complex<float>, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(dtype::complex<double>, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(dtype::float8_e4m3fn, GPUContext)
|
|
INSTANTIATE_STRIDEDCOPY_KERNEL(dtype::float8_e5m2, GPUContext)
|
|
#endif
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(strided_copy,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::StridedCopyKernel,
|
|
bool,
|
|
uint8_t,
|
|
uint16_t,
|
|
uint32_t,
|
|
uint64_t,
|
|
int8_t,
|
|
int16_t,
|
|
int32_t,
|
|
int64_t,
|
|
float,
|
|
double,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
phi::complex64,
|
|
phi::complex128,
|
|
phi::float8_e4m3fn,
|
|
phi::float8_e5m2) {}
|