chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,453 @@
|
||||
/* 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
|
||||
Reference in New Issue
Block a user