wehub-resource-sync
601 lines
22 KiB
C++
601 lines
22 KiB
C++
// Copyright (c) 2022 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
|
|
|
|
// CUDA and HIP use same api
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
#include <algorithm>
|
|
#include "paddle/common/ddim.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/kernels/empty_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/cub.h"
|
|
#include "paddle/phi/kernels/primitive/kernel_primitives.h"
|
|
|
|
namespace phi {
|
|
namespace funcs {
|
|
using Mode = kps::details::ReduceMode;
|
|
|
|
/*
|
|
* Count how many of the data being processed by the current block are true
|
|
* 1. Load data from global memory and cast from bool to int64_t
|
|
* 2. Get result of this thread according to thread reduce
|
|
* 3. Get result of this block according to block reduce
|
|
* 4. first block store 0 and current result
|
|
*/
|
|
template <typename T>
|
|
struct NonZeroFunctor {
|
|
HOSTDEVICE NonZeroFunctor() {}
|
|
HOSTDEVICE inline T operator()(const T in) {
|
|
if (in) {
|
|
return static_cast<T>(1);
|
|
} else {
|
|
return static_cast<T>(0);
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename InT, typename OutT, int VecSize, int IsBoundary>
|
|
__device__ void GetBlockCountImpl(const InT *in,
|
|
OutT *out,
|
|
int64_t num,
|
|
int64_t repeat) {
|
|
InT in_data[VecSize];
|
|
OutT temp[VecSize];
|
|
OutT result = static_cast<OutT>(0.0f);
|
|
using Add = kps::AddFunctor<OutT>;
|
|
using Cast = NonZeroFunctor<InT>;
|
|
int64_t store_fix = BLOCK_ID_X + repeat * GRID_NUM_X;
|
|
|
|
kps::Init<InT, VecSize>(&in_data[0], static_cast<InT>(0.0f));
|
|
kps::ReadData<InT, VecSize, 1, IsBoundary>(&in_data[0], in, num);
|
|
kps::ElementwiseUnary<InT, OutT, VecSize, 1, Cast>(
|
|
&temp[0], &in_data[0], Cast());
|
|
kps::Reduce<OutT, VecSize, 1, Add, Mode::kLocalMode>(
|
|
&result, &temp[0], Add(), true);
|
|
kps::Reduce<OutT, 1, 1, Add, Mode::kGlobalMode>(
|
|
&result, &result, Add(), true);
|
|
if (store_fix == 0) {
|
|
// first block's fix_size = 0;
|
|
OutT tmp = static_cast<OutT>(0.0f);
|
|
kps::WriteData<OutT, 1, 1, true>(out + store_fix, &tmp, 1);
|
|
}
|
|
|
|
// store num of this block
|
|
kps::WriteData<OutT, 1, 1, true>(out + store_fix + 1, &result, 1);
|
|
}
|
|
|
|
// Count how many data is not zero in current block
|
|
template <typename InT, typename OutT, int VecSize>
|
|
__global__ void GetBlockCountKernel(const InT *in,
|
|
OutT *out,
|
|
int64_t numel,
|
|
int64_t main_offset) {
|
|
int64_t size = static_cast<int64_t>(BLOCK_NUM_X) * VecSize;
|
|
int64_t data_offset = size * BLOCK_ID_X;
|
|
int64_t stride = size * GRID_NUM_X;
|
|
int64_t repeat = 0;
|
|
for (; data_offset < main_offset; data_offset += stride) {
|
|
GetBlockCountImpl<InT, OutT, VecSize, false>(
|
|
in + data_offset, out, size, repeat);
|
|
repeat++; // to get the real blockIdx
|
|
}
|
|
|
|
int64_t num = numel - data_offset;
|
|
if (num > 0) {
|
|
GetBlockCountImpl<InT, OutT, VecSize, true>(
|
|
in + data_offset, out, num, repeat);
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Get block num prefix us one block, VecSize must be 2
|
|
* 1. Each thread load 2 data : threadIdx.x and threadIdx.x + blockDimx.x
|
|
* 2. Cumsum limitation is blockDim.x must be less than 512
|
|
*/
|
|
|
|
template <typename InT,
|
|
typename OutT,
|
|
typename Functor,
|
|
int VecSize,
|
|
bool IsBoundary>
|
|
__device__ void CumsumImpl(
|
|
const InT *in, OutT *out, OutT *pre_cumsum, int num, Functor func) {
|
|
__shared__ OutT max_thread_data;
|
|
OutT temp[VecSize];
|
|
InT arg[VecSize];
|
|
OutT result[VecSize];
|
|
// init data_pr
|
|
kps::Init<InT, VecSize>(&arg[0], static_cast<InT>(0.0f));
|
|
// set pre_cumsum
|
|
kps::Init<OutT, VecSize>(&temp[0], *pre_cumsum);
|
|
// load data to arg
|
|
kps::ReadData<InT, InT, VecSize, 1, IsBoundary>(
|
|
&arg[0], in, num, 1, BLOCK_NUM_X, 1);
|
|
// block cumsum
|
|
kps::Cumsum<InT, OutT, Functor>(&result[0], &arg[0], func);
|
|
// result = cumsum_result + pre_cumsum
|
|
kps::ElementwiseBinary<OutT, OutT, VecSize, 1, Functor>(
|
|
&result[0], &result[0], &temp[0], func);
|
|
// get the last prefix sum
|
|
if ((THREAD_ID_X == BLOCK_NUM_X - 1) && !IsBoundary) {
|
|
max_thread_data = result[VecSize - 1];
|
|
}
|
|
__syncthreads();
|
|
// update pre_cumsum
|
|
*pre_cumsum = max_thread_data;
|
|
kps::WriteData<OutT, OutT, VecSize, 1, IsBoundary>(
|
|
out, &result[0], num, 1, BLOCK_NUM_X, 1);
|
|
}
|
|
|
|
// Compute this store_offset of this block
|
|
template <typename InT, typename OutT, typename Functor, int VecSize>
|
|
__global__ void CumsumOneBlock(const InT *in,
|
|
OutT *out,
|
|
int64_t numel,
|
|
int64_t main_offset,
|
|
Functor func) {
|
|
int64_t stride = static_cast<int64_t>(BLOCK_NUM_X) * VecSize;
|
|
int64_t offset = 0;
|
|
OutT pre_cumsum = static_cast<OutT>(0);
|
|
for (; offset < main_offset; offset += stride) {
|
|
CumsumImpl<InT, OutT, Functor, VecSize, false>(
|
|
in + offset, out + offset, &pre_cumsum, stride, func);
|
|
}
|
|
|
|
int64_t num = numel - offset;
|
|
if (num > 0) {
|
|
CumsumImpl<InT, OutT, Functor, VecSize, true>(
|
|
in + offset, out + offset, &pre_cumsum, num, func);
|
|
}
|
|
}
|
|
|
|
// where_index
|
|
template <typename OutT,
|
|
typename MT,
|
|
typename InT,
|
|
typename Functor,
|
|
int VecSize,
|
|
int IsBoundary,
|
|
int MaskData>
|
|
struct SelectCaller {
|
|
__device__ void inline operator()(OutT *out,
|
|
const MT *mask_data,
|
|
const InT *in,
|
|
Functor func,
|
|
int64_t data_offset,
|
|
int64_t store_num,
|
|
int64_t thread_fix,
|
|
int64_t num) {
|
|
int64_t in_data[VecSize];
|
|
OutT store_data[VecSize * DDim::kMaxRank];
|
|
// set index
|
|
kps::InitWithDataIndex<int64_t, VecSize, 1>(&in_data[0], data_offset);
|
|
// Get store data according to mask_idt
|
|
kps::OperatorTernary<MT, int64_t, OutT, Functor>(
|
|
store_data, mask_data, &in_data[0], func, VecSize);
|
|
kps::details::WriteData<OutT>(out + thread_fix, &store_data[0], store_num);
|
|
}
|
|
};
|
|
|
|
// masked_select
|
|
template <typename OutT,
|
|
typename MT,
|
|
typename InT,
|
|
typename Functor,
|
|
int VecSize,
|
|
int IsBoundary>
|
|
struct SelectCaller<OutT, MT, InT, Functor, VecSize, IsBoundary, 1> {
|
|
__device__ void inline operator()(OutT *out,
|
|
const MT *mask_data,
|
|
const InT *in,
|
|
Functor func,
|
|
int data_offset,
|
|
int store_num,
|
|
int thread_fix,
|
|
int num) {
|
|
InT in_data[VecSize];
|
|
OutT store_data[VecSize * DDim::kMaxRank];
|
|
kps::ReadData<InT, VecSize, 1, IsBoundary>(&in_data[0], in, num);
|
|
// Get store data according to mask_idt
|
|
kps::OperatorTernary<MT, InT, OutT, Functor>(
|
|
store_data, mask_data, &in_data[0], func, VecSize);
|
|
kps::details::WriteData<OutT>(out + thread_fix, &store_data[0], store_num);
|
|
}
|
|
};
|
|
|
|
// masked_select_grad
|
|
template <typename OutT,
|
|
typename MT,
|
|
typename InT,
|
|
typename Functor,
|
|
int VecSize,
|
|
int IsBoundary>
|
|
struct SelectCaller<OutT, MT, InT, Functor, VecSize, IsBoundary, 2> {
|
|
__device__ void inline operator()(OutT *out,
|
|
const MT *mask_data,
|
|
const InT *in,
|
|
Functor func,
|
|
int data_offset,
|
|
int store_num,
|
|
int thread_fix,
|
|
int num) {
|
|
InT in_data[VecSize];
|
|
OutT store_data[VecSize * DDim::kMaxRank];
|
|
kps::details::ReadData<InT>(&in_data[0], in + thread_fix, store_num);
|
|
kps::OperatorTernary<MT, InT, OutT, Functor>(
|
|
store_data, mask_data, &in_data[0], func, VecSize);
|
|
kps::WriteData<OutT, VecSize, 1, IsBoundary>(out, &store_data[0], num);
|
|
}
|
|
};
|
|
|
|
/**
|
|
* Get mask's index if mask == true
|
|
*/
|
|
template <typename InT,
|
|
typename MT,
|
|
typename OutT,
|
|
typename Functor,
|
|
int VecSize,
|
|
int MaskData,
|
|
int IsBoundary> // SelectType = 1 Mask_select else where_index
|
|
__device__ void SelectKernelImpl(OutT *out,
|
|
const MT *mask,
|
|
const InT *in,
|
|
Functor func,
|
|
int64_t num,
|
|
int64_t data_offset,
|
|
int64_t store_rank) {
|
|
const int kCVecSize = 2;
|
|
// each thread cumsum 2 data
|
|
using IdT = int64_t;
|
|
// Set index data type
|
|
using Add = kps::AddFunctor<IdT>; // for cumsum
|
|
using Cast = NonZeroFunctor<MT>; // for mask
|
|
|
|
IdT init_idx = static_cast<IdT>(0.0f);
|
|
MT init_mask = static_cast<MT>(0.0f);
|
|
|
|
IdT num_thread[kCVecSize];
|
|
IdT cumsum_thread[kCVecSize];
|
|
|
|
MT mask_data[VecSize];
|
|
IdT mask_idt[VecSize];
|
|
// init data_pr
|
|
kps::Init<IdT, kCVecSize>(&cumsum_thread[0], init_idx);
|
|
kps::Init<IdT, kCVecSize>(&num_thread[0], init_idx);
|
|
kps::Init<MT, VecSize>(&mask_data[0], init_mask);
|
|
// Load mask
|
|
kps::ReadData<MT, VecSize, 1, IsBoundary>(&mask_data[0], mask, num);
|
|
// Cast from MT to int
|
|
kps::ElementwiseUnary<MT, IdT, VecSize, 1, Cast>(
|
|
&mask_idt[0], &mask_data[0], Cast());
|
|
// Get the num of thread only num_thread[1] has data
|
|
kps::Reduce<IdT, VecSize, 1, Add, Mode::kLocalMode>(
|
|
&num_thread[0], &mask_idt[0], Add(), true);
|
|
// Get cumsum_thread cumsum from 0 to num_thread cumsum_thread[0] is the
|
|
// thread_fix
|
|
kps::Cumsum<IdT, IdT, Add>(&cumsum_thread[0], &num_thread[0], Add());
|
|
// get thread_fix
|
|
IdT thread_fix = (cumsum_thread[0] - num_thread[0]) * store_rank;
|
|
// get how many data need to store
|
|
IdT store_num = num_thread[0] * store_rank;
|
|
// thread store num data, each thread may has different num
|
|
// Get store data(index) according to mask_idt
|
|
SelectCaller<OutT, MT, InT, Functor, VecSize, IsBoundary, MaskData> select;
|
|
select(out, mask_data, in, func, data_offset, store_num, thread_fix, num);
|
|
}
|
|
|
|
template <typename MT,
|
|
typename InT,
|
|
typename CT,
|
|
typename OutT,
|
|
typename Functor,
|
|
int VecSize,
|
|
int MaskData>
|
|
__global__ void SelectKernel(OutT *out,
|
|
const MT *mask,
|
|
const InT *in,
|
|
CT *cumsum,
|
|
Functor func,
|
|
const int64_t numel,
|
|
int64_t main_offset,
|
|
int64_t store_rank) {
|
|
int64_t size = static_cast<int64_t>(BLOCK_ID_X) * VecSize;
|
|
int64_t data_offset = size * BLOCK_NUM_X;
|
|
int64_t stride = static_cast<int64_t>(BLOCK_NUM_X) * GRID_NUM_X * VecSize;
|
|
int64_t repeat = 0;
|
|
CT block_store_offset = 0;
|
|
for (; data_offset < main_offset; data_offset += stride) {
|
|
// Cumsum index
|
|
int64_t idx_cumsum = repeat * GRID_NUM_X + BLOCK_ID_X;
|
|
kps::details::ReadData<CT>(&block_store_offset, cumsum + idx_cumsum, 1);
|
|
int64_t out_fix =
|
|
MaskData < 2 ? block_store_offset * store_rank : data_offset;
|
|
int64_t in_fix =
|
|
MaskData < 2 ? data_offset : block_store_offset * store_rank;
|
|
SelectKernelImpl<InT, MT, OutT, Functor, VecSize, MaskData, false>(
|
|
out + out_fix,
|
|
mask + data_offset,
|
|
in + in_fix,
|
|
func,
|
|
size,
|
|
data_offset,
|
|
store_rank);
|
|
repeat++;
|
|
}
|
|
|
|
int64_t num = numel - data_offset;
|
|
if (num > 0) {
|
|
// Cumsum index
|
|
int64_t idx_cumsum = repeat * GRID_NUM_X + BLOCK_ID_X;
|
|
kps::details::ReadData<CT>(&block_store_offset, cumsum + idx_cumsum, 1);
|
|
int64_t out_fix =
|
|
MaskData < 2 ? block_store_offset * store_rank : data_offset;
|
|
int64_t in_fix =
|
|
MaskData < 2 ? data_offset : block_store_offset * store_rank;
|
|
SelectKernelImpl<InT, MT, OutT, Functor, VecSize, MaskData, true>(
|
|
out + out_fix,
|
|
mask + data_offset,
|
|
in + in_fix,
|
|
func,
|
|
num,
|
|
data_offset,
|
|
store_rank);
|
|
}
|
|
}
|
|
|
|
inline int64_t Floor(int64_t in, int64_t div) { return in / div * div; }
|
|
|
|
// SelectData = 1 then masked_select; SelectData = 0 then where_index
|
|
template <typename MT,
|
|
typename InT,
|
|
typename OutT,
|
|
int SelectData,
|
|
typename Functor>
|
|
void SelectKernel(const KPDevice &dev_ctx,
|
|
const DenseTensor &condition,
|
|
const DenseTensor &in_data,
|
|
DenseTensor *out,
|
|
Functor func) {
|
|
const MT *cond_data = condition.data<MT>();
|
|
const int64_t numel = condition.numel();
|
|
auto dims = condition.dims();
|
|
int rank = SelectData ? 1 : dims.size();
|
|
const InT *in_data_ptr = SelectData ? in_data.data<InT>() : nullptr;
|
|
// calculate the inclusive prefix sum of "true_num_array"
|
|
// to get the index of "out" tensor,
|
|
// and the total number of cond_data[i]==true.
|
|
// Example:
|
|
// condition: F T T F F F T T
|
|
// before: 0 1 1 0 0 0 1 1
|
|
// after: 0 1 2 2 2 2 3 4
|
|
// out: 1 2 6 7
|
|
// alloc for cpu
|
|
using CT = int64_t; // set Count_data Type
|
|
const int t_size = sizeof(CT);
|
|
|
|
const phi::Place &cuda_place = dev_ctx.GetPlace();
|
|
CPUPlace cpu_place = CPUPlace();
|
|
|
|
// 1.1 get stored data num of per block
|
|
int kVecSize = 4;
|
|
kVecSize = std::min(phi::GetVectorizedSize(&condition), kVecSize);
|
|
if (in_data.numel() > 0) {
|
|
kVecSize = std::min(phi::GetVectorizedSize(&in_data), kVecSize);
|
|
} else {
|
|
kVecSize = 1;
|
|
}
|
|
while (kVecSize > 1 && numel % kVecSize != 0) {
|
|
kVecSize /= 2;
|
|
}
|
|
#define CALL_GET_BLOCK_COUNT_KERNEL(kVecSize) \
|
|
case kVecSize: \
|
|
GetBlockCountKernel<MT, CT, kVecSize><<<grid, block, 0, stream>>>( \
|
|
cond_data, count_data, numel, main_offset); \
|
|
break;
|
|
|
|
#define CALL_SELECT_KERNEL(kVecSize) \
|
|
case kVecSize: \
|
|
SelectKernel<MT, InT, CT, OutT, Functor, kVecSize, SelectData> \
|
|
<<<grid, block, 0, stream>>>(out_data, \
|
|
cond_data, \
|
|
in_data_ptr, \
|
|
cumsum_data, \
|
|
func, \
|
|
numel, \
|
|
main_offset, \
|
|
rank); \
|
|
break;
|
|
|
|
#ifdef PADDLE_WITH_XPU_KP
|
|
int block = 64;
|
|
auto stream = dev_ctx.x_context()->xpu_stream;
|
|
const int num_per_block = kVecSize * block;
|
|
const int64_t need_grids = (numel + num_per_block - 1) / num_per_block;
|
|
const int64_t grid = std::min(need_grids, static_cast<int64_t>(8));
|
|
#else
|
|
const int block = 256;
|
|
const int num_per_block = kVecSize * block;
|
|
const int64_t need_grids = (numel + num_per_block - 1) / num_per_block;
|
|
const int64_t grid = std::min(need_grids, static_cast<int64_t>(256));
|
|
auto stream = dev_ctx.stream();
|
|
#endif
|
|
const int64_t main_offset = Floor(numel, num_per_block);
|
|
// 1.2 alloc tmp data for CoutBlock
|
|
const int64_t size_count_block = need_grids + 1;
|
|
std::vector<int64_t> dims_vec = {size_count_block * 2};
|
|
IntArray dims_array(dims_vec);
|
|
DenseTensor count_mem = Empty<CT, KPDevice>(dev_ctx, dims_array);
|
|
CT *count_data = count_mem.data<CT>();
|
|
// 1.3 launch CountKernel
|
|
switch (kVecSize) {
|
|
CALL_GET_BLOCK_COUNT_KERNEL(4)
|
|
CALL_GET_BLOCK_COUNT_KERNEL(2)
|
|
CALL_GET_BLOCK_COUNT_KERNEL(1)
|
|
default:
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Unsupported vectorized size: %d", kVecSize));
|
|
break;
|
|
}
|
|
|
|
// 2.1 alloc cumsum data for CoutBlock prefix
|
|
DenseTensor cumsum_mem = Empty<CT, KPDevice>(dev_ctx, dims_array);
|
|
CT *cumsum_data = cumsum_mem.data<CT>();
|
|
// 2.2 get prefix of count_data for real out_index
|
|
CT total_true_num = static_cast<CT>(0); // init
|
|
const int kCumVesize = 2;
|
|
const int block_c = 256;
|
|
const int64_t main_offset_c = Floor(size_count_block, (kCumVesize * block_c));
|
|
|
|
using Add = kps::AddFunctor<CT>;
|
|
CumsumOneBlock<CT, CT, Add, kCumVesize><<<1, block_c, 0, stream>>>(
|
|
count_data, cumsum_data, size_count_block, main_offset_c, Add());
|
|
// 3.1 set temp ptr for in;
|
|
// 3.1 alloc for out
|
|
// 3.1.1 get true_num for gpu place the last cumsum is the true_num
|
|
memory_utils::Copy(cpu_place,
|
|
&total_true_num,
|
|
cuda_place,
|
|
cumsum_data + need_grids,
|
|
t_size,
|
|
dev_ctx.stream());
|
|
|
|
dev_ctx.Wait();
|
|
// 3.1.2 allock for out with total_true_num
|
|
std::vector<int64_t> out_dim = {static_cast<int64_t>(total_true_num)};
|
|
|
|
if (SelectData == 1) {
|
|
out->Resize(out_dim);
|
|
} else if (SelectData == 0) { // == 0 where_index
|
|
out_dim.push_back(static_cast<int64_t>(rank));
|
|
out->Resize(out_dim);
|
|
}
|
|
auto out_data = dev_ctx.template Alloc<OutT>(out);
|
|
// 3.2 get true data's index according to cond_data and cumsum_data
|
|
if (total_true_num <= 0) return;
|
|
switch (kVecSize) {
|
|
CALL_SELECT_KERNEL(4)
|
|
CALL_SELECT_KERNEL(2)
|
|
CALL_SELECT_KERNEL(1)
|
|
default:
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Unsupported vectorized size: %d", kVecSize));
|
|
break;
|
|
}
|
|
#undef CALL_GET_BLOCK_COUNT_KERNEL
|
|
#undef CALL_SELECT_KERNEL
|
|
}
|
|
|
|
// SelectData = 1 then masked_select; SelectData = 0 then where_index
|
|
template <typename MT,
|
|
typename InT,
|
|
typename OutT,
|
|
int SelectData,
|
|
typename Functor>
|
|
void RestrictSelectKernel(const KPDevice &dev_ctx,
|
|
const DenseTensor &condition,
|
|
const DenseTensor &in_data,
|
|
const int64_t total_true_num,
|
|
DenseTensor *out,
|
|
Functor func) {
|
|
const MT *cond_data = condition.data<MT>();
|
|
const int64_t numel = condition.numel();
|
|
auto dims = condition.dims();
|
|
int rank = SelectData ? 1 : dims.size();
|
|
const InT *in_data_ptr = SelectData ? in_data.data<InT>() : nullptr;
|
|
// calculate the inclusive prefix sum of "true_num_array"
|
|
// to get the index of "out" tensor,
|
|
// and the total number of cond_data[i]==true.
|
|
// Example:
|
|
// condition: F T T F F F T T
|
|
// before: 0 1 1 0 0 0 1 1
|
|
// after: 0 1 2 2 2 2 3 4
|
|
// out: 1 2 6 7
|
|
// alloc for cpu
|
|
using CT = int64_t; // set Count_data Type
|
|
const int t_size = sizeof(CT);
|
|
|
|
const phi::Place &cuda_place = dev_ctx.GetPlace();
|
|
CPUPlace cpu_place = CPUPlace();
|
|
|
|
// 1.1 get stored data num of per block
|
|
const int kVecSize = 4;
|
|
#ifdef PADDLE_WITH_XPU_KP
|
|
int block = 64;
|
|
auto stream = dev_ctx.x_context()->xpu_stream;
|
|
const int num_per_block = kVecSize * block;
|
|
const int64_t need_grids = (numel + num_per_block - 1) / num_per_block;
|
|
const int grid = std::min(need_grids, static_cast<int64_t>(8));
|
|
#else
|
|
const int block = 256;
|
|
const int num_per_block = kVecSize * block;
|
|
const int64_t need_grids = (numel + num_per_block - 1) / num_per_block;
|
|
const int grid = std::min(need_grids, static_cast<int64_t>(256));
|
|
auto stream = dev_ctx.stream();
|
|
#endif
|
|
const int64_t main_offset = Floor(numel, num_per_block);
|
|
// 1.2 alloc tmp data for CoutBlock
|
|
const int size_count_block = need_grids + 1;
|
|
std::vector<int> dims_vec = {size_count_block * 2};
|
|
IntArray dims_array(dims_vec);
|
|
DenseTensor count_mem = Empty<CT, KPDevice>(dev_ctx, dims_array);
|
|
CT *count_data = count_mem.data<CT>();
|
|
// 1.3 launch CountKernel
|
|
GetBlockCountKernel<MT, CT, kVecSize>
|
|
<<<grid, block, 0, stream>>>(cond_data, count_data, numel, main_offset);
|
|
// 2.1 alloc cumsum data for CoutBlock prefix
|
|
DenseTensor cumsum_mem = Empty<CT, KPDevice>(dev_ctx, dims_array);
|
|
CT *cumsum_data = cumsum_mem.data<CT>();
|
|
// 2.2 get prefix of count_data for real out_index
|
|
// CT total_true_num = static_cast<CT>(0); // init
|
|
const int kCumVesize = 2;
|
|
const int block_c = 256;
|
|
const int main_offset_c = Floor(size_count_block, (kCumVesize * block_c));
|
|
|
|
using Add = kps::AddFunctor<CT>;
|
|
CumsumOneBlock<CT, CT, Add, kCumVesize><<<1, block_c, 0, stream>>>(
|
|
count_data, cumsum_data, size_count_block, main_offset_c, Add());
|
|
// 3.1 set temp ptr for in;
|
|
// 3.1 alloc for out
|
|
// 3.1.1 get true_num for gpu place the last cumsum is the true_num
|
|
// 3.1.2 allock for out with total_true_num
|
|
std::vector<int64_t> out_dim = {static_cast<int64_t>(total_true_num)};
|
|
|
|
if (SelectData == 1) {
|
|
out->Resize(out_dim);
|
|
} else if (SelectData == 0) { // == 0 where_index
|
|
out_dim.push_back(static_cast<int64_t>(rank));
|
|
out->Resize(out_dim);
|
|
}
|
|
auto out_data = dev_ctx.template Alloc<OutT>(out);
|
|
// 3.2 get true data's index according to cond_data and cumsum_data
|
|
if (total_true_num <= 0) return;
|
|
SelectKernel<MT, InT, CT, OutT, Functor, kVecSize, SelectData>
|
|
<<<grid, block, 0, stream>>>(out_data,
|
|
cond_data,
|
|
in_data_ptr,
|
|
cumsum_data,
|
|
func,
|
|
numel,
|
|
main_offset,
|
|
rank);
|
|
}
|
|
|
|
} // namespace funcs
|
|
} // namespace phi
|
|
|
|
#endif
|