wehub-resource-sync
289 lines
11 KiB
Plaintext
289 lines
11 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 "paddle/phi/kernels/gpu/collect_fpn_proposals_kernel.h"
|
|
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
|
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/allocator.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/mixed_vector.h"
|
|
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
|
|
#include "paddle/phi/kernels/funcs/cub.h"
|
|
#include "paddle/phi/kernels/funcs/detection/bbox_util.h"
|
|
#include "paddle/phi/kernels/funcs/for_range.h"
|
|
#include "paddle/phi/kernels/funcs/gather.cu.h"
|
|
#include "paddle/phi/kernels/funcs/strided_memcpy.h"
|
|
#include "paddle/phi/kernels/impl/collect_fpn_proposals_kernel_impl.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
namespace phi {
|
|
|
|
static constexpr int kNumCUDAThreads = 64;
|
|
static constexpr int kNumMaximumNumBlocks = 4096;
|
|
|
|
const int kBBoxSize = 4;
|
|
|
|
static inline int NumBlocks(const int N) {
|
|
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
|
kNumMaximumNumBlocks);
|
|
}
|
|
|
|
static __global__ void GetLengthLoD(const int nthreads,
|
|
const int* batch_ids,
|
|
int* length_lod) {
|
|
CUDA_KERNEL_LOOP(i, nthreads) { CudaAtomicAdd(length_lod + batch_ids[i], 1); }
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void GPUCollectFpnProposalsOpKernel(
|
|
const Context& dev_ctx,
|
|
const std::vector<const DenseTensor*>& multi_level_rois,
|
|
const std::vector<const DenseTensor*>& multi_level_scores,
|
|
const optional<std::vector<const DenseTensor*>>& multi_level_rois_num,
|
|
int post_nms_topn,
|
|
DenseTensor* fpn_rois_out,
|
|
DenseTensor* rois_num_out) {
|
|
const auto roi_ins = multi_level_rois;
|
|
const auto score_ins = multi_level_scores;
|
|
auto fpn_rois = fpn_rois_out;
|
|
|
|
const int post_nms_topN = post_nms_topn;
|
|
|
|
// concat inputs along axis = 0
|
|
int roi_offset = 0;
|
|
int score_offset = 0;
|
|
int total_roi_num = 0;
|
|
for (size_t i = 0; i < roi_ins.size(); ++i) {
|
|
total_roi_num += roi_ins[i]->dims()[0];
|
|
}
|
|
|
|
int real_post_num = min(post_nms_topN, total_roi_num);
|
|
fpn_rois->Resize({real_post_num, kBBoxSize});
|
|
dev_ctx.template Alloc<T>(fpn_rois);
|
|
DenseTensor concat_rois;
|
|
DenseTensor concat_scores;
|
|
concat_rois.Resize({total_roi_num, kBBoxSize});
|
|
T* concat_rois_data = dev_ctx.template Alloc<T>(&concat_rois);
|
|
concat_scores.Resize({total_roi_num, 1});
|
|
T* concat_scores_data = dev_ctx.template Alloc<T>(&concat_scores);
|
|
DenseTensor roi_batch_id_list;
|
|
roi_batch_id_list.Resize({total_roi_num});
|
|
int* roi_batch_id_data = dev_ctx.template HostAlloc<int>(&roi_batch_id_list);
|
|
int index = 0;
|
|
int lod_size;
|
|
auto place = dev_ctx.GetPlace();
|
|
|
|
auto multi_rois_num = multi_level_rois_num
|
|
? multi_level_rois_num.get()
|
|
: std::vector<const DenseTensor*>();
|
|
for (size_t i = 0; i < roi_ins.size(); ++i) {
|
|
auto roi_in = roi_ins[i];
|
|
auto score_in = score_ins[i];
|
|
if (multi_rois_num.size() > 0) {
|
|
DenseTensor temp;
|
|
Copy(dev_ctx, *multi_rois_num[i], CPUPlace(), true, &temp);
|
|
const int* length_in = temp.data<int>();
|
|
lod_size = multi_rois_num[i]->numel();
|
|
for (size_t n = 0; n < lod_size; ++n) {
|
|
for (size_t j = 0; j < length_in[n]; ++j) {
|
|
roi_batch_id_data[index++] = n;
|
|
}
|
|
}
|
|
} else {
|
|
auto length_in = roi_in->lod().back();
|
|
lod_size = length_in.size() - 1;
|
|
for (size_t n = 0; n < lod_size; ++n) {
|
|
for (size_t j = length_in[n]; j < length_in[n + 1]; ++j) {
|
|
roi_batch_id_data[index++] = n;
|
|
}
|
|
}
|
|
}
|
|
|
|
memory_utils::Copy(place,
|
|
concat_rois_data + roi_offset,
|
|
place,
|
|
roi_in->data<T>(),
|
|
roi_in->numel() * sizeof(T),
|
|
dev_ctx.stream());
|
|
memory_utils::Copy(place,
|
|
concat_scores_data + score_offset,
|
|
place,
|
|
score_in->data<T>(),
|
|
score_in->numel() * sizeof(T),
|
|
dev_ctx.stream());
|
|
roi_offset += roi_in->numel();
|
|
score_offset += score_in->numel();
|
|
}
|
|
|
|
// copy batch id list to GPU
|
|
DenseTensor roi_batch_id_list_gpu;
|
|
Copy(dev_ctx,
|
|
roi_batch_id_list,
|
|
dev_ctx.GetPlace(),
|
|
false,
|
|
&roi_batch_id_list_gpu);
|
|
|
|
DenseTensor index_in_t;
|
|
index_in_t.Resize({total_roi_num});
|
|
int* idx_in = dev_ctx.template Alloc<int>(&index_in_t);
|
|
funcs::ForRange<GPUContext> for_range_total(dev_ctx, total_roi_num);
|
|
for_range_total(funcs::RangeInitFunctor{0, 1, idx_in});
|
|
|
|
DenseTensor keys_out_t;
|
|
keys_out_t.Resize({total_roi_num});
|
|
T* keys_out = dev_ctx.template Alloc<T>(&keys_out_t);
|
|
DenseTensor index_out_t;
|
|
index_out_t.Resize({total_roi_num});
|
|
int* idx_out = dev_ctx.template Alloc<int>(&index_out_t);
|
|
|
|
// Determine temporary device storage requirements
|
|
size_t temp_storage_bytes = 0;
|
|
cub::DeviceRadixSort::SortPairsDescending<T, int>(nullptr,
|
|
temp_storage_bytes,
|
|
concat_scores.data<T>(),
|
|
keys_out,
|
|
idx_in,
|
|
idx_out,
|
|
total_roi_num,
|
|
0,
|
|
sizeof(T) * 8,
|
|
dev_ctx.stream());
|
|
// Allocate temporary storage
|
|
auto d_temp_storage = memory_utils::Alloc(place, temp_storage_bytes);
|
|
|
|
// Run sorting operation
|
|
// sort score to get corresponding index
|
|
cub::DeviceRadixSort::SortPairsDescending<T, int>(d_temp_storage->ptr(),
|
|
temp_storage_bytes,
|
|
concat_scores.data<T>(),
|
|
keys_out,
|
|
idx_in,
|
|
idx_out,
|
|
total_roi_num,
|
|
0,
|
|
sizeof(T) * 8,
|
|
dev_ctx.stream());
|
|
index_out_t.Resize({real_post_num});
|
|
DenseTensor sorted_rois;
|
|
sorted_rois.Resize({real_post_num, kBBoxSize});
|
|
dev_ctx.template Alloc<T>(&sorted_rois);
|
|
DenseTensor sorted_batch_id;
|
|
sorted_batch_id.Resize({real_post_num});
|
|
dev_ctx.template Alloc<int>(&sorted_batch_id);
|
|
funcs::GPUGather<T>(dev_ctx, concat_rois, index_out_t, &sorted_rois);
|
|
funcs::GPUGather<int>(
|
|
dev_ctx, roi_batch_id_list_gpu, index_out_t, &sorted_batch_id);
|
|
|
|
DenseTensor batch_index_t;
|
|
batch_index_t.Resize({real_post_num});
|
|
int* batch_idx_in = dev_ctx.template Alloc<int>(&batch_index_t);
|
|
funcs::ForRange<GPUContext> for_range_post(dev_ctx, real_post_num);
|
|
for_range_post(funcs::RangeInitFunctor{0, 1, batch_idx_in});
|
|
|
|
DenseTensor out_id_t;
|
|
out_id_t.Resize({real_post_num});
|
|
int* out_id_data = dev_ctx.template Alloc<int>(&out_id_t);
|
|
// Determine temporary device storage requirements
|
|
temp_storage_bytes = 0;
|
|
cub::DeviceRadixSort::SortPairs<int, int>(nullptr,
|
|
temp_storage_bytes,
|
|
sorted_batch_id.data<int>(),
|
|
out_id_data,
|
|
batch_idx_in,
|
|
index_out_t.data<int>(),
|
|
real_post_num,
|
|
0,
|
|
sizeof(int) * 8,
|
|
dev_ctx.stream());
|
|
// Allocate temporary storage
|
|
d_temp_storage = memory_utils::Alloc(place, temp_storage_bytes);
|
|
|
|
// Run sorting operation
|
|
// sort batch_id to get corresponding index
|
|
cub::DeviceRadixSort::SortPairs<int, int>(d_temp_storage->ptr(),
|
|
temp_storage_bytes,
|
|
sorted_batch_id.data<int>(),
|
|
out_id_data,
|
|
batch_idx_in,
|
|
index_out_t.data<int>(),
|
|
real_post_num,
|
|
0,
|
|
sizeof(int) * 8,
|
|
dev_ctx.stream());
|
|
|
|
funcs::GPUGather<T>(dev_ctx, sorted_rois, index_out_t, fpn_rois);
|
|
|
|
DenseTensor length_lod;
|
|
length_lod.Resize({lod_size});
|
|
int* length_lod_data = dev_ctx.template Alloc<int>(&length_lod);
|
|
funcs::SetConstant<GPUContext, int> set_zero;
|
|
set_zero(dev_ctx, &length_lod, static_cast<int>(0));
|
|
|
|
int blocks = NumBlocks(real_post_num);
|
|
int threads = kNumCUDAThreads;
|
|
|
|
// get length-based lod by batch ids
|
|
GetLengthLoD<<<blocks, threads, 0, dev_ctx.stream()>>>(
|
|
real_post_num, out_id_data, length_lod_data);
|
|
PADDLE_ENFORCE_EQ(
|
|
backends::gpu::IsCUDAGraphCapturing(),
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"CollectFpnProposals does not support CUDA Graph capture: async D2H "
|
|
"copy to local vector 'length_lod_cpu' will bake the destination "
|
|
"address into the graph; on replay the vector is re-created at a "
|
|
"different address, causing a dangling-pointer write."));
|
|
std::vector<int> length_lod_cpu(lod_size);
|
|
memory_utils::Copy(CPUPlace(),
|
|
length_lod_cpu.data(),
|
|
place,
|
|
length_lod_data,
|
|
sizeof(int) * lod_size,
|
|
dev_ctx.stream());
|
|
dev_ctx.Wait();
|
|
|
|
std::vector<size_t> offset(1, 0);
|
|
for (int i = 0; i < lod_size; ++i) {
|
|
offset.emplace_back(offset.back() + length_lod_cpu[i]);
|
|
}
|
|
|
|
if (rois_num_out != nullptr) {
|
|
auto* rois_num = rois_num_out;
|
|
rois_num->Resize({lod_size});
|
|
int* rois_num_data = dev_ctx.template Alloc<int>(rois_num);
|
|
memory_utils::Copy(place,
|
|
rois_num_data,
|
|
place,
|
|
length_lod_data,
|
|
lod_size * sizeof(int),
|
|
dev_ctx.stream());
|
|
}
|
|
|
|
LegacyLoD lod;
|
|
lod.emplace_back(offset);
|
|
fpn_rois->set_lod(lod);
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(collect_fpn_proposals,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::GPUCollectFpnProposalsOpKernel,
|
|
float,
|
|
double) {
|
|
kernel->InputAt(2).SetDataType(phi::DataType::INT32);
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
|
|
}
|