wehub-resource-sync
454 lines
17 KiB
Plaintext
454 lines
17 KiB
Plaintext
/* Copyright (c) 2016 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 <vector>
|
|
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/funcs/math_function_impl.h"
|
|
#ifndef PADDLE_WITH_CUSTOM_DEVICE
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/math_function_blas_impl.h"
|
|
#else
|
|
#include "paddle/phi/backends/gpu/gpu_info.h"
|
|
#endif
|
|
|
|
namespace phi {
|
|
namespace funcs {
|
|
|
|
// The following part of the code refers to NVIDIA-cutlass
|
|
// https://github.com/NVIDIA/cutlass/blob/master/tools/util/include/cutlass/util/device_nchw_to_nhwc.h
|
|
// Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
|
|
// reserved. SPDX-License-Identifier: BSD-3-Clause
|
|
template <typename T>
|
|
__global__ void batch_transpose_kernel(T* output,
|
|
const T* input,
|
|
const int batch,
|
|
const int M,
|
|
const int N,
|
|
int swizzle) {
|
|
const int64_t num = static_cast<int64_t>(M) * N;
|
|
// "+1" to avoid smem bank conflict
|
|
__shared__ T shbuf[32 * (32 + 1)];
|
|
const int32_t tid = threadIdx.y * blockDim.x + threadIdx.x;
|
|
const int32_t wid = tid / 32;
|
|
const int32_t lid = tid % 32;
|
|
const int32_t batch_i = blockIdx.z;
|
|
const int32_t mi0 = (blockIdx.y * swizzle + blockIdx.x % swizzle) * 32;
|
|
const int32_t ni0 = blockIdx.x / swizzle * 32;
|
|
|
|
const size_t input_idx = batch_i * num + (mi0 + wid) * N + ni0;
|
|
const T* A = input + input_idx;
|
|
if (ni0 + lid < N) {
|
|
const int lid_x_33 = lid * 33;
|
|
if ((mi0 + 32) <= M) {
|
|
int mi = wid; // between 0 and 7
|
|
#pragma unroll
|
|
for (int mLoopIdx = 0; mLoopIdx < 4; mLoopIdx++) {
|
|
shbuf[lid_x_33 + mi] = A[lid];
|
|
A = &A[8 * N];
|
|
mi += 8;
|
|
}
|
|
} else {
|
|
for (int mi = wid; mi < 32; mi += 8) {
|
|
if ((mi + mi0) < M) {
|
|
shbuf[lid_x_33 + mi] = A[lid];
|
|
}
|
|
A = &A[8 * N];
|
|
}
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
const int32_t miOut = mi0 + lid;
|
|
output = &output[batch_i * num + miOut];
|
|
if (miOut < M) {
|
|
if (ni0 + 32 < N) {
|
|
int nI = wid;
|
|
#pragma unroll
|
|
for (int nLoopIdx = 0; nLoopIdx < 4; ++nLoopIdx) {
|
|
output[(ni0 + nI) * M] = shbuf[(nI)*33 + lid];
|
|
nI += 8;
|
|
}
|
|
} else {
|
|
for (int nI = wid; nI < 32; nI += 8) {
|
|
if (ni0 + nI < N) {
|
|
output[(ni0 + nI) * M] = shbuf[(nI)*33 + lid];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void BatchTranspose(T* output,
|
|
const T* input,
|
|
int64_t batch,
|
|
int64_t m,
|
|
int64_t n,
|
|
const GPUContext* dev_ctx) {
|
|
int64_t device_id = dev_ctx->GetPlace().GetDeviceId();
|
|
const auto& prop = phi::backends::gpu::GetDeviceProperties(device_id);
|
|
int max_grid_y = prop.maxGridSize[1];
|
|
int64_t input_num = batch * m * n;
|
|
|
|
if (input_num >= std::numeric_limits<int>::max()) {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Unsupported input size, batch: %ld,m: %ld, n: %ld", batch, m, n));
|
|
}
|
|
|
|
dim3 logical_grid((n + 31) / 32, (m + 31) / 32, batch);
|
|
dim3 block(32, 8);
|
|
// we set swizzle to 2 default.
|
|
int swizzle = (logical_grid.y + max_grid_y - 1) / max_grid_y;
|
|
swizzle = std::max(swizzle, 2);
|
|
dim3 physical_grid(logical_grid.x * swizzle,
|
|
(logical_grid.y + swizzle - 1) / swizzle,
|
|
batch);
|
|
batch_transpose_kernel<<<physical_grid, block>>>(
|
|
output, input, batch, m, n, swizzle);
|
|
}
|
|
|
|
template void BatchTranspose(float16* output,
|
|
const float16* input,
|
|
int64_t batch,
|
|
int64_t m,
|
|
int64_t n,
|
|
const GPUContext* dev_ctx);
|
|
template void BatchTranspose(float* output,
|
|
const float* input,
|
|
int64_t batch,
|
|
int64_t m,
|
|
int64_t n,
|
|
const GPUContext* dev_ctx);
|
|
template void BatchTranspose(bfloat16* output,
|
|
const bfloat16* input,
|
|
int64_t batch,
|
|
int64_t m,
|
|
int64_t n,
|
|
const GPUContext* dev_ctx);
|
|
|
|
template struct SetConstant<GPUContext, float8_e4m3fn>;
|
|
template struct SetConstant<GPUContext, float8_e5m2>;
|
|
template struct SetConstant<GPUContext, float16>;
|
|
template struct SetConstant<GPUContext, bfloat16>;
|
|
template struct SetConstant<GPUContext, float>;
|
|
template struct SetConstant<GPUContext, double>;
|
|
template struct SetConstant<GPUContext, uint8_t>;
|
|
template struct SetConstant<GPUContext, uint16_t>;
|
|
template struct SetConstant<GPUContext, uint32_t>;
|
|
template struct SetConstant<GPUContext, uint64_t>;
|
|
template struct SetConstant<GPUContext, int8_t>;
|
|
template struct SetConstant<GPUContext, int>;
|
|
template struct SetConstant<GPUContext, int16_t>;
|
|
template struct SetConstant<GPUContext, int64_t>;
|
|
template struct SetConstant<GPUContext, bool>;
|
|
template struct SetConstant<GPUContext, phi::complex64>;
|
|
template struct SetConstant<GPUContext, phi::complex128>;
|
|
|
|
#ifndef PADDLE_WITH_CUSTOM_DEVICE
|
|
template struct SetConstant<phi::GPUPinnedContext, float16>;
|
|
template struct SetConstant<phi::GPUPinnedContext, bfloat16>;
|
|
template struct SetConstant<phi::GPUPinnedContext, float>;
|
|
template struct SetConstant<phi::GPUPinnedContext, double>;
|
|
template struct SetConstant<phi::GPUPinnedContext, uint8_t>;
|
|
template struct SetConstant<phi::GPUPinnedContext, uint16_t>;
|
|
template struct SetConstant<phi::GPUPinnedContext, uint32_t>;
|
|
template struct SetConstant<phi::GPUPinnedContext, uint64_t>;
|
|
template struct SetConstant<phi::GPUPinnedContext, int8_t>;
|
|
template struct SetConstant<phi::GPUPinnedContext, int>;
|
|
template struct SetConstant<phi::GPUPinnedContext, int16_t>;
|
|
template struct SetConstant<phi::GPUPinnedContext, int64_t>;
|
|
template struct SetConstant<phi::GPUPinnedContext, bool>;
|
|
template struct SetConstant<phi::GPUPinnedContext, phi::complex64>;
|
|
template struct SetConstant<phi::GPUPinnedContext, phi::complex128>;
|
|
#endif
|
|
|
|
#define DEFINE_GPU_TRANS(RANK) \
|
|
template struct Transpose<GPUContext, bool, RANK>; \
|
|
template struct Transpose<GPUContext, uint8_t, RANK>; \
|
|
template struct Transpose<GPUContext, uint16_t, RANK>; \
|
|
template struct Transpose<GPUContext, uint32_t, RANK>; \
|
|
template struct Transpose<GPUContext, uint64_t, RANK>; \
|
|
template struct Transpose<GPUContext, float, RANK>; \
|
|
template struct Transpose<GPUContext, double, RANK>; \
|
|
template struct Transpose<GPUContext, float8_e4m3fn, RANK>; \
|
|
template struct Transpose<GPUContext, float8_e5m2, RANK>; \
|
|
template struct Transpose<GPUContext, float16, RANK>; \
|
|
template struct Transpose<GPUContext, bfloat16, RANK>; \
|
|
template struct Transpose<GPUContext, int8_t, RANK>; \
|
|
template struct Transpose<GPUContext, int16_t, RANK>; \
|
|
template struct Transpose<GPUContext, int32_t, RANK>; \
|
|
template struct Transpose<GPUContext, int64_t, RANK>; \
|
|
template struct Transpose<GPUContext, phi::complex64, RANK>; \
|
|
template struct Transpose<GPUContext, phi::complex128, RANK>;
|
|
|
|
DEFINE_GPU_TRANS(1);
|
|
DEFINE_GPU_TRANS(2);
|
|
DEFINE_GPU_TRANS(3);
|
|
DEFINE_GPU_TRANS(4);
|
|
DEFINE_GPU_TRANS(5);
|
|
DEFINE_GPU_TRANS(6);
|
|
|
|
template <typename T>
|
|
__global__ void FillConstantKernel(const int N, T* a, const T val) {
|
|
for (int64_t i =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
i < N;
|
|
i += blockDim.x * gridDim.x) {
|
|
a[i] = val;
|
|
}
|
|
}
|
|
|
|
#define REINTERPRET(T, DST_PTR, SRC_PTR) \
|
|
T* DST_PTR = reinterpret_cast<T*>(SRC_PTR)
|
|
|
|
template <typename T>
|
|
__global__ void TransposeNormalKernel(const T* in_ptr,
|
|
T* out_ptr,
|
|
int64_t element,
|
|
const int64_t* in_stride_ptr,
|
|
const int64_t* out_stride_ptr,
|
|
const int64_t* axis_ptr,
|
|
int rank) {
|
|
CUDA_KERNEL_LOOP(out_idx, element) {
|
|
int64_t in_idx = 0;
|
|
int64_t tmp_idx = out_idx;
|
|
for (int i = 0; i < rank; ++i) {
|
|
const int64_t coordinate = tmp_idx / out_stride_ptr[i];
|
|
tmp_idx -= coordinate * out_stride_ptr[i];
|
|
in_idx += coordinate * in_stride_ptr[axis_ptr[i]];
|
|
}
|
|
out_ptr[out_idx] = in_ptr[in_idx];
|
|
}
|
|
}
|
|
|
|
template <typename DeviceContext, typename T>
|
|
void TransposeNormal<DeviceContext, T>::operator()(
|
|
const DeviceContext& dev_ctx,
|
|
const DenseTensor& in,
|
|
DenseTensor* out,
|
|
const std::vector<int>& axis) {
|
|
const int rank = axis.size();
|
|
auto in_stride = common::stride(in.dims());
|
|
auto out_stride = common::stride(out->dims());
|
|
auto* in_ptr = in.data<T>();
|
|
auto* out_ptr = out->data<T>();
|
|
|
|
// copy in_stride, out_stride, axis to gpu device
|
|
const phi::Place& cuda_place = dev_ctx.GetPlace();
|
|
CPUPlace cpu_place = CPUPlace();
|
|
size_t size = 3 * rank * sizeof(int64_t);
|
|
auto cpu_buf_holder = phi::memory_utils::Alloc(cpu_place, size);
|
|
auto cuda_buf_holder = phi::memory_utils::Alloc(cuda_place, size);
|
|
REINTERPRET(int64_t, cpu_buf, cpu_buf_holder->ptr());
|
|
REINTERPRET(int64_t, cuda_buf, cuda_buf_holder->ptr());
|
|
for (int i = 0; i < rank; ++i) {
|
|
cpu_buf[i] = in_stride[i];
|
|
cpu_buf[rank + i] = out_stride[i];
|
|
cpu_buf[2 * rank + i] = axis[i];
|
|
}
|
|
memory_utils::Copy(
|
|
cuda_place, cuda_buf, cpu_place, cpu_buf, size, dev_ctx.stream());
|
|
REINTERPRET(const int64_t, in_stride_ptr, cuda_buf);
|
|
REINTERPRET(const int64_t, out_stride_ptr, cuda_buf + rank);
|
|
REINTERPRET(const int64_t, axis_ptr, cuda_buf + 2 * rank);
|
|
|
|
const int MAX_BLOCK_DIM = dev_ctx.GetMaxThreadsPerBlock();
|
|
const int MAX_GRID_DIM = dev_ctx.GetMaxPhysicalThreadCount() / MAX_BLOCK_DIM;
|
|
int64_t elements = in.numel();
|
|
int block_size = (elements >= MAX_BLOCK_DIM)
|
|
? MAX_BLOCK_DIM
|
|
: (1 << static_cast<int>(std::log2(elements)));
|
|
const int64_t grid_size_64 =
|
|
std::min(elements / block_size, static_cast<int64_t>(MAX_GRID_DIM));
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_size_64, "grid_size");
|
|
const uint32_t grid_size = static_cast<uint32_t>(grid_size_64);
|
|
TransposeNormalKernel<T><<<grid_size, block_size, 0, dev_ctx.stream()>>>(
|
|
in_ptr, out_ptr, elements, in_stride_ptr, out_stride_ptr, axis_ptr, rank);
|
|
}
|
|
|
|
template <typename T>
|
|
struct TransposeNormal<GPUContext, T> {
|
|
void operator()(const GPUContext& dev_ctx,
|
|
const DenseTensor& in,
|
|
DenseTensor* out,
|
|
const std::vector<int>& axis) {
|
|
const int rank = axis.size();
|
|
auto in_stride = stride(in.dims());
|
|
auto out_stride = stride(out->dims());
|
|
auto* in_ptr = in.data<T>();
|
|
auto* out_ptr = out->data<T>();
|
|
|
|
// copy in_stride, out_stride, axis to gpu device
|
|
const phi::Place& cuda_place = dev_ctx.GetPlace();
|
|
CPUPlace cpu_place = CPUPlace();
|
|
size_t size = 3 * rank * sizeof(int64_t);
|
|
auto cpu_buf_holder = phi::memory_utils::Alloc(cpu_place, size);
|
|
auto cuda_buf_holder = phi::memory_utils::Alloc(cuda_place, size);
|
|
REINTERPRET(int64_t, cpu_buf, cpu_buf_holder->ptr());
|
|
REINTERPRET(int64_t, cuda_buf, cuda_buf_holder->ptr());
|
|
for (int i = 0; i < rank; ++i) {
|
|
cpu_buf[i] = in_stride[i];
|
|
cpu_buf[rank + i] = out_stride[i];
|
|
cpu_buf[2 * rank + i] = axis[i];
|
|
}
|
|
memory_utils::Copy(
|
|
cuda_place, cuda_buf, cpu_place, cpu_buf, size, dev_ctx.stream());
|
|
REINTERPRET(const int64_t, in_stride_ptr, cuda_buf);
|
|
REINTERPRET(const int64_t, out_stride_ptr, cuda_buf + rank);
|
|
REINTERPRET(const int64_t, axis_ptr, cuda_buf + 2 * rank);
|
|
|
|
const int MAX_BLOCK_DIM = dev_ctx.GetMaxThreadsPerBlock();
|
|
const int MAX_GRID_DIM =
|
|
dev_ctx.GetMaxPhysicalThreadCount() / MAX_BLOCK_DIM;
|
|
int64_t elements = in.numel();
|
|
int block_size = (elements >= MAX_BLOCK_DIM)
|
|
? MAX_BLOCK_DIM
|
|
: (1 << static_cast<int>(std::log2(elements)));
|
|
const int64_t grid_size_64 =
|
|
std::min(elements / block_size, static_cast<int64_t>(MAX_GRID_DIM));
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_size_64, "grid_size");
|
|
const uint32_t grid_size = static_cast<uint32_t>(grid_size_64);
|
|
TransposeNormalKernel<T>
|
|
<<<grid_size, block_size, 0, dev_ctx.stream()>>>(in_ptr,
|
|
out_ptr,
|
|
elements,
|
|
in_stride_ptr,
|
|
out_stride_ptr,
|
|
axis_ptr,
|
|
rank);
|
|
}
|
|
};
|
|
|
|
// define transpose normal
|
|
#define DEFINE_GPU_TRANS_NORMAL(TYPE) \
|
|
template struct TransposeNormal<GPUContext, TYPE>
|
|
|
|
DEFINE_GPU_TRANS_NORMAL(phi::float8_e4m3fn);
|
|
DEFINE_GPU_TRANS_NORMAL(phi::float8_e5m2);
|
|
DEFINE_GPU_TRANS_NORMAL(float16);
|
|
DEFINE_GPU_TRANS_NORMAL(bfloat16);
|
|
DEFINE_GPU_TRANS_NORMAL(float);
|
|
DEFINE_GPU_TRANS_NORMAL(double);
|
|
DEFINE_GPU_TRANS_NORMAL(int);
|
|
DEFINE_GPU_TRANS_NORMAL(int64_t);
|
|
DEFINE_GPU_TRANS_NORMAL(bool);
|
|
DEFINE_GPU_TRANS_NORMAL(int16_t);
|
|
DEFINE_GPU_TRANS_NORMAL(uint8_t);
|
|
DEFINE_GPU_TRANS_NORMAL(uint16_t);
|
|
DEFINE_GPU_TRANS_NORMAL(uint32_t);
|
|
DEFINE_GPU_TRANS_NORMAL(uint64_t);
|
|
DEFINE_GPU_TRANS_NORMAL(int8_t);
|
|
DEFINE_GPU_TRANS_NORMAL(phi::complex64);
|
|
DEFINE_GPU_TRANS_NORMAL(phi::complex128);
|
|
|
|
struct TensorSetConstantGPU {
|
|
TensorSetConstantGPU(const DeviceContext& dev_ctx,
|
|
DenseTensor* tensor,
|
|
float value)
|
|
: dev_ctx_(dev_ctx), tensor_(tensor), value_(value) {}
|
|
|
|
template <typename T>
|
|
void apply() const {
|
|
SetConstant<GPUContext, T> functor;
|
|
functor(reinterpret_cast<const GPUContext&>(dev_ctx_),
|
|
tensor_,
|
|
static_cast<T>(value_));
|
|
}
|
|
|
|
const DeviceContext& dev_ctx_;
|
|
DenseTensor* tensor_;
|
|
float value_;
|
|
};
|
|
|
|
template <>
|
|
void set_constant_with_place<GPUPlace>(const DeviceContext& dev_ctx,
|
|
DenseTensor* tensor,
|
|
float value) {
|
|
phi::VisitDataType(tensor->dtype(),
|
|
TensorSetConstantGPU(dev_ctx, tensor, value));
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void RowwiseAddKernel(
|
|
const T* a, const T* b, T* c, int64_t width, int64_t num) {
|
|
T tmp = 1.0 / width;
|
|
CUDA_KERNEL_LOOP_TYPE(i, num, int64_t) {
|
|
int64_t h = i * tmp;
|
|
int64_t w = i - h * width;
|
|
c[i] = a[i] + b[w];
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
struct RowwiseAdd<GPUContext, T> {
|
|
void operator()(const GPUContext& dev_ctx,
|
|
const DenseTensor& input,
|
|
const DenseTensor& vector,
|
|
DenseTensor* output) {
|
|
auto in_dims = input.dims();
|
|
auto out_dims = output->dims();
|
|
auto size = input.numel() / in_dims[0];
|
|
PADDLE_ENFORCE_EQ(
|
|
vector.numel(),
|
|
size,
|
|
common::errors::InvalidArgument(
|
|
"The input vector size"
|
|
" should be equal to the size of each row of input tensor."
|
|
" Expected vector size=%d, but received %d",
|
|
size,
|
|
vector.numel()));
|
|
const char* in_dims_cstr = in_dims.to_str().c_str();
|
|
const char* out_dims_cstr = out_dims.to_str().c_str();
|
|
PADDLE_ENFORCE_EQ(
|
|
out_dims,
|
|
in_dims,
|
|
common::errors::InvalidArgument(
|
|
"The output tensor shape should be same as the input tensor"
|
|
" shape. Expected output tensor shape: %s,"
|
|
" but received %s",
|
|
in_dims_cstr,
|
|
out_dims_cstr));
|
|
int blocks = 512;
|
|
int64_t max_grids = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
|
int grids = std::min((input.numel() + blocks - 1) / blocks, max_grids);
|
|
RowwiseAddKernel<T>
|
|
<<<grids, blocks, 0, dev_ctx.stream()>>>(input.data<T>(),
|
|
vector.data<T>(),
|
|
output->data<T>(),
|
|
in_dims[1],
|
|
input.numel());
|
|
}
|
|
};
|
|
|
|
template struct RowwiseAdd<GPUContext, float>;
|
|
template struct RowwiseAdd<GPUContext, double>;
|
|
template struct ColwiseSum<GPUContext, float>;
|
|
template struct ColwiseSum<GPUContext, int>;
|
|
template struct ColwiseSum<GPUContext, int64_t>;
|
|
|
|
template struct RowwiseSum<GPUContext, float>;
|
|
|
|
template struct RowwiseMean<GPUContext, float>;
|
|
template struct RowwiseMean<GPUContext, double>;
|
|
|
|
} // namespace funcs
|
|
} // namespace phi
|