wehub-resource-sync
110 lines
3.6 KiB
Plaintext
110 lines
3.6 KiB
Plaintext
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include <algorithm>
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/cub.h"
|
|
#include "paddle/phi/kernels/funcs/math.h"
|
|
#include "paddle/phi/kernels/impl/sequence_softmax_kernel_impl.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, int BlockDim>
|
|
using BlockReduce = cub::BlockReduce<T, BlockDim>;
|
|
|
|
template <typename T, int BlockDim>
|
|
using BlockReduceTempStorage = typename BlockReduce<T, BlockDim>::TempStorage;
|
|
|
|
template <typename T, int BlockDim>
|
|
__global__ void sequence_softmax_kernel(const T *in_data,
|
|
const size_t *ref_lod,
|
|
const size_t src_height,
|
|
T *out_data) {
|
|
__shared__ BlockReduceTempStorage<T, BlockDim> temp_storage;
|
|
__shared__ T shared_max_data;
|
|
__shared__ T shared_sum_data;
|
|
|
|
for (size_t i = blockIdx.x; i < src_height; i += gridDim.x) {
|
|
size_t start = ref_lod[i];
|
|
size_t span = ref_lod[i + 1] - start;
|
|
|
|
// Find the max ele
|
|
T max_ele = -FLT_MAX;
|
|
for (size_t tid = threadIdx.x; tid < span; tid += blockDim.x) {
|
|
T ele = in_data[start + tid];
|
|
max_ele = max_ele > ele ? max_ele : ele;
|
|
}
|
|
max_ele =
|
|
BlockReduce<T, BlockDim>(temp_storage).Reduce(max_ele, cub::Max());
|
|
if (threadIdx.x == 0) {
|
|
shared_max_data = max_ele;
|
|
}
|
|
__syncthreads();
|
|
|
|
// sum
|
|
T sum_data = 0;
|
|
for (size_t tid = threadIdx.x; tid < span; tid += blockDim.x) {
|
|
T ele = in_data[start + tid];
|
|
sum_data += funcs::real_exp(ele - shared_max_data);
|
|
}
|
|
sum_data =
|
|
BlockReduce<T, BlockDim>(temp_storage).Reduce(sum_data, cub::Sum());
|
|
if (threadIdx.x == 0) {
|
|
shared_sum_data = sum_data;
|
|
}
|
|
__syncthreads();
|
|
|
|
// get final resit
|
|
for (size_t tid = threadIdx.x; tid < span; tid += blockDim.x) {
|
|
T ele = in_data[start + tid];
|
|
ele = funcs::real_exp(ele - shared_max_data) / shared_sum_data;
|
|
out_data[start + tid] = ele;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
struct SequenceSoftmaxFunctor<GPUContext, T> {
|
|
void operator()(const GPUContext &dev_ctx,
|
|
const DenseTensor &x,
|
|
const Vector<size_t> &ref_lod, /*referenced lod*/
|
|
DenseTensor *out) {
|
|
int height = ref_lod.size() - 1;
|
|
|
|
const int kThreadsPerBlock = 32;
|
|
int thread_x = kThreadsPerBlock;
|
|
int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
|
int max_blocks = std::max(max_threads / kThreadsPerBlock, 1);
|
|
|
|
dim3 block_size(thread_x);
|
|
dim3 grid_size(max_blocks);
|
|
MixVector<size_t> mixv_ref_lod(&ref_lod);
|
|
sequence_softmax_kernel<T, kThreadsPerBlock>
|
|
<<<grid_size, block_size, 0, dev_ctx.stream()>>>(
|
|
x.data<T>(),
|
|
mixv_ref_lod.CUDAData(dev_ctx.GetPlace()),
|
|
height,
|
|
dev_ctx.Alloc<T>(out));
|
|
}
|
|
};
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(sequence_softmax,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::SequenceSoftmaxKernel,
|
|
float,
|
|
double) {}
|