wehub-resource-sync
587 lines
22 KiB
Plaintext
587 lines
22 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.
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
#include <hiprand.h>
|
|
#include <hiprand_kernel.h>
|
|
typedef hiprandState curandState;
|
|
#else
|
|
#include <curand.h>
|
|
#include <curand_kernel.h>
|
|
#endif
|
|
|
|
#include <iterator>
|
|
#include <random>
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
#include "paddle/phi/kernels/class_center_sample_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/cub.h"
|
|
|
|
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
|
#endif
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi {
|
|
#define CUDA_KERNEL_LOOP_TYPE(i, n, index_type) \
|
|
for (index_type i = blockIdx.x * blockDim.x + threadIdx.x, \
|
|
step = blockDim.x * gridDim.x; \
|
|
i < (n); \
|
|
i += step)
|
|
|
|
#define CUDA_KERNEL_LOOP(i, n) CUDA_KERNEL_LOOP_TYPE(i, n, int32_t)
|
|
|
|
static constexpr int kNumCUDAThreads = 512;
|
|
static constexpr int kNumMaximumNumBlocks = 4096;
|
|
|
|
inline int32_t NumBlocks(const int32_t n) {
|
|
return std::min((n + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
|
kNumMaximumNumBlocks);
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void RandomSampleClassCenter(const int64_t n,
|
|
int64_t seed,
|
|
int64_t increment,
|
|
const int64_t max_val,
|
|
T* buffer) {
|
|
const int id = blockIdx.x * blockDim.x + threadIdx.x;
|
|
curandState localState;
|
|
size_t local_seed =
|
|
(static_cast<size_t>(seed) + 0x9E3779B9U +
|
|
(static_cast<size_t>(id) << 6U) + (static_cast<size_t>(id) >> 2U));
|
|
#ifdef PADDLE_WITH_HIP
|
|
hiprand_init(local_seed, id, increment, &localState);
|
|
CUDA_KERNEL_LOOP(i, n) {
|
|
buffer[i] = static_cast<T>(hiprand(&localState) % max_val);
|
|
}
|
|
#else
|
|
curand_init(local_seed, id, increment, &localState);
|
|
CUDA_KERNEL_LOOP(i, n) {
|
|
buffer[i] = static_cast<T>(curand(&localState) % max_val);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void Range(const int64_t n, T* out) {
|
|
CUDA_KERNEL_LOOP(i, n) { out[i] = static_cast<T>(i); }
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void MarkPositiveClassCenter(const int64_t n,
|
|
const int64_t rank,
|
|
const T* class_interval_ptr,
|
|
const int num_classes,
|
|
const T* labels,
|
|
T* out) {
|
|
CUDA_KERNEL_LOOP(i, n) {
|
|
T label = labels[i] - class_interval_ptr[rank];
|
|
if (label >= 0 && label < num_classes) {
|
|
out[label] = label - num_classes;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__device__ void FindIntervalIndex(const T* class_interval_ptr,
|
|
const int64_t nranks,
|
|
const T value,
|
|
int64_t* find_index) {
|
|
int64_t start = 0;
|
|
int64_t end = nranks;
|
|
int64_t mid = ((end - start) >> 1) + start + 1;
|
|
while (start < end) {
|
|
if (class_interval_ptr[mid] == value) break;
|
|
if (class_interval_ptr[mid] > value)
|
|
end = mid - 1;
|
|
else
|
|
start = mid;
|
|
mid = ((end - start) >> 1) + start + 1;
|
|
}
|
|
*find_index = min(mid, end);
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GetClassCenterBound(const int64_t n,
|
|
const int64_t nranks,
|
|
const T* class_interval_ptr,
|
|
const T* key_ptr,
|
|
const T* value_ptr,
|
|
T* bound_index,
|
|
T* bound_value) {
|
|
CUDA_KERNEL_LOOP(i, n) {
|
|
if (i != 0) {
|
|
int64_t cur_index, pre_index;
|
|
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[i], &cur_index);
|
|
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[i - 1], &pre_index);
|
|
if (cur_index > pre_index) {
|
|
assert(cur_index < nranks);
|
|
#pragma unroll
|
|
for (int32_t j = pre_index + 1; j <= cur_index; ++j) {
|
|
bound_index[j] = static_cast<T>(i);
|
|
bound_value[j] = value_ptr[i];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
CUDA_KERNEL_LOOP(i, nranks + 1) {
|
|
int64_t first_index, last_index;
|
|
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[0], &first_index);
|
|
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[n - 1], &last_index);
|
|
if (i <= first_index) {
|
|
bound_index[i] = 0;
|
|
bound_value[i] = value_ptr[0];
|
|
} else if (i > last_index) {
|
|
bound_index[i] = n;
|
|
bound_value[i] = value_ptr[n - 1] + 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GetRemappedLabel(const int64_t n,
|
|
const int64_t nranks,
|
|
const T* sampled_class_interval_ptr,
|
|
const T* bound_index,
|
|
const T* bound_value,
|
|
const T* label_map_key,
|
|
T* label_map_value,
|
|
T* mapped_label) {
|
|
CUDA_KERNEL_LOOP(i, n) {
|
|
#pragma unroll
|
|
for (int64_t j = 0; j < nranks; j++) {
|
|
if (i >= bound_index[j] && i < bound_index[j + 1]) {
|
|
label_map_value[i] =
|
|
label_map_value[i] - bound_value[j] + sampled_class_interval_ptr[j];
|
|
}
|
|
}
|
|
mapped_label[label_map_key[i]] = label_map_value[i];
|
|
}
|
|
}
|
|
|
|
// aligned vector generates vectorized load/store on CUDA
|
|
template <typename T, int Size>
|
|
struct alignas(sizeof(T) * Size) AlignedVector {
|
|
T val[Size];
|
|
};
|
|
|
|
template <typename T>
|
|
inline int VectorizedSize(const T* pointer) {
|
|
uint64_t address = reinterpret_cast<uint64_t>(pointer);
|
|
constexpr int vec4 = std::alignment_of<AlignedVector<T, 4>>::value; // NOLINT
|
|
if (address % vec4 == 0) {
|
|
return 4;
|
|
}
|
|
return 1;
|
|
}
|
|
|
|
#undef CUDA_KERNEL_LOOP
|
|
|
|
template <typename T>
|
|
class NotEqualToPreviousAdjacentIterator {
|
|
public:
|
|
using self_type = NotEqualToPreviousAdjacentIterator;
|
|
using value_type = T;
|
|
using difference_type = std::ptrdiff_t;
|
|
using pointer = T*;
|
|
using reference = T;
|
|
using iterator_category = std::input_iterator_tag;
|
|
|
|
public:
|
|
__host__ __device__ __forceinline__
|
|
NotEqualToPreviousAdjacentIterator(const T* arr, int64_t offset)
|
|
: arr_(arr), offset_(offset) {}
|
|
|
|
__host__ __device__ __forceinline__ reference operator*() const {
|
|
return offset_ == 0 ? 0 : (arr_[offset_] == arr_[offset_ - 1] ? 0 : 1);
|
|
}
|
|
|
|
template <typename Distance>
|
|
__host__ __device__ __forceinline__ self_type operator+(Distance n) const {
|
|
self_type ret(arr_, offset_ + n);
|
|
return ret;
|
|
}
|
|
|
|
template <typename Distance>
|
|
__host__ __device__ __forceinline__ self_type operator-(Distance n) const {
|
|
self_type ret(arr_, offset_ - n);
|
|
return ret;
|
|
}
|
|
|
|
template <typename Distance>
|
|
__host__ __device__ __forceinline__ reference operator[](Distance n) const {
|
|
return *(*this + n);
|
|
}
|
|
|
|
private:
|
|
const T* arr_;
|
|
int64_t offset_;
|
|
};
|
|
|
|
template <typename T>
|
|
struct ActualNumSampledFunctor {
|
|
__host__ __device__ __forceinline__ T operator()(const T& a,
|
|
const T& b) const {
|
|
return max(num_samples, (b - a));
|
|
}
|
|
T num_samples;
|
|
explicit ActualNumSampledFunctor(const T num) : num_samples(num) {}
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
class MemoryBuffer {
|
|
public:
|
|
MemoryBuffer(const int num_buffer_ele,
|
|
const int num_temp_ele,
|
|
const int nranks,
|
|
const Context& dev_ctx) {
|
|
offset1 = 0;
|
|
offset2 = offset1 + num_buffer_ele;
|
|
offset3 = offset2 + num_buffer_ele;
|
|
offset4 = offset3 + num_buffer_ele;
|
|
offset5 = offset4 + num_buffer_ele;
|
|
offset6 = offset5 + (nranks + 1);
|
|
offset7 = offset6 + (nranks + 1);
|
|
offset8 = offset7 + (nranks + 1);
|
|
offset9 = offset8 + num_temp_ele;
|
|
|
|
buffer.Resize({4 * num_buffer_ele + 3 * (nranks + 1) + num_temp_ele});
|
|
buffer_ptr = dev_ctx.template Alloc<T>(&buffer);
|
|
}
|
|
|
|
T* cub_sort_keys_ptr() { return buffer_ptr + offset1; }
|
|
T* cub_sort_keys_out_ptr() { return buffer_ptr + offset2; }
|
|
T* cub_sort_values_ptr() { return buffer_ptr + offset3; }
|
|
T* cub_sort_values_out_ptr() { return buffer_ptr + offset4; }
|
|
T* bound_index_ptr() { return buffer_ptr + offset5; }
|
|
T* bound_value_ptr() { return buffer_ptr + offset6; }
|
|
T* class_interval_ptr() { return buffer_ptr + offset7; }
|
|
void* cub_temp_storage_ptr() {
|
|
return reinterpret_cast<void*>(buffer_ptr + offset8);
|
|
}
|
|
|
|
private:
|
|
DenseTensor buffer;
|
|
T* buffer_ptr;
|
|
int offset1;
|
|
int offset2;
|
|
int offset3;
|
|
int offset4;
|
|
int offset5;
|
|
int offset6;
|
|
int offset7;
|
|
int offset8;
|
|
int offset9;
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void ClassCenterSampleKernel(const Context& dev_ctx,
|
|
const DenseTensor& label,
|
|
int num_classes,
|
|
int num_samples,
|
|
int ring_id,
|
|
int rank,
|
|
int nranks,
|
|
bool fix_seed,
|
|
int seed,
|
|
DenseTensor* remapped_label,
|
|
DenseTensor* sampled_local_class_center) {
|
|
PADDLE_ENFORCE_GT(num_classes,
|
|
0,
|
|
errors::InvalidArgument(
|
|
"The value 'num_classes' for Op(class_center_sample) "
|
|
"must be greater than 0, "
|
|
"but the value given is %d.",
|
|
num_classes));
|
|
|
|
PADDLE_ENFORCE_GT(num_samples,
|
|
0,
|
|
errors::InvalidArgument(
|
|
"The value 'num_samples' for Op(class_center_sample) "
|
|
"must be greater than 0, "
|
|
"but the value given is %d.",
|
|
num_samples));
|
|
|
|
PADDLE_ENFORCE_LE(num_samples,
|
|
num_classes,
|
|
errors::InvalidArgument(
|
|
"The value 'num_samples' for Op(class_center_sample) "
|
|
"must be less than or equal to %d, "
|
|
"but the value given is %d.",
|
|
num_classes,
|
|
num_samples));
|
|
|
|
auto place = dev_ctx.GetPlace();
|
|
|
|
int64_t batch_size = label.numel();
|
|
// TODO(large-tensor): downstream functors may still use int
|
|
PADDLE_ENFORCE_LE_INT_MAX(label.numel(), "label.numel()");
|
|
// Algorithm:
|
|
// We first randomly generate a value in [0, num_classes) on each position
|
|
// in a array(shape[num_classes]). Then, we mark the element as negative
|
|
// value in the array according input label. Now, we can sort the array
|
|
// by ascending to ensure that the positive class center always in the
|
|
// front of the sorted array. So, we can get the sampled class center
|
|
// index by sorted keys. Finally, we can get the rempped label by remap
|
|
// the input label according sampled class center.
|
|
|
|
// step 1: Calculate num classes per device using nccl all reduce
|
|
std::vector<T> shard_dim_vec(nranks + 1, 0);
|
|
shard_dim_vec[rank + 1] = num_classes;
|
|
DenseTensor num_classes_per_device;
|
|
TensorFromVector(shard_dim_vec, dev_ctx, &num_classes_per_device);
|
|
T* num_classes_per_device_ptr = num_classes_per_device.data<T>();
|
|
|
|
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
|
if (nranks > 1) {
|
|
auto stream = dev_ctx.stream();
|
|
distributed::NCCLCommContext* comm_ctx = nullptr;
|
|
comm_ctx =
|
|
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
|
PADDLE_ENFORCE_NE(comm_ctx,
|
|
nullptr,
|
|
common::errors::Unavailable(
|
|
"NCCLCommContext is nullptr, collective op should "
|
|
"has ring_id attr."));
|
|
|
|
comm_ctx->AllReduce(
|
|
&num_classes_per_device, num_classes_per_device, ncclSum, stream);
|
|
backends::gpu::GpuStreamSync(stream);
|
|
}
|
|
#endif
|
|
|
|
// step 2: Determine temporary device storage requirements
|
|
int num_buffer_ele = std::max(static_cast<int>(batch_size), num_classes);
|
|
size_t cub_sort_temp_store_size = 0;
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceRadixSort::SortPairs<T, T>(nullptr,
|
|
cub_sort_temp_store_size,
|
|
nullptr,
|
|
nullptr,
|
|
nullptr,
|
|
nullptr,
|
|
num_buffer_ele,
|
|
0,
|
|
sizeof(T) * 8,
|
|
dev_ctx.stream())));
|
|
|
|
size_t cub_sum_temp_store_size = 0;
|
|
NotEqualToPreviousAdjacentIterator<T> unique_counting_iter_temp(nullptr, 0);
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceScan::InclusiveSum<NotEqualToPreviousAdjacentIterator<T>, T*>(
|
|
nullptr,
|
|
cub_sum_temp_store_size,
|
|
unique_counting_iter_temp,
|
|
nullptr,
|
|
batch_size,
|
|
dev_ctx.stream())));
|
|
|
|
size_t cub_scan_temp_store_size = 0;
|
|
ActualNumSampledFunctor<T> actual_num_sampled_op_temp(num_samples);
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceScan::InclusiveScan(nullptr,
|
|
cub_scan_temp_store_size,
|
|
num_classes_per_device_ptr,
|
|
num_classes_per_device_ptr,
|
|
actual_num_sampled_op_temp,
|
|
nranks + 1,
|
|
dev_ctx.stream())));
|
|
|
|
size_t cub_temp_storage_bytes =
|
|
std::max(std::max(cub_sort_temp_store_size, cub_scan_temp_store_size),
|
|
cub_sum_temp_store_size);
|
|
int num_temp_ele = cub_temp_storage_bytes / sizeof(T) + 1;
|
|
PADDLE_ENFORCE_GT(
|
|
(4 * num_buffer_ele + 3 * (nranks + 1) + num_temp_ele),
|
|
0,
|
|
errors::InvalidArgument(
|
|
"Illegal memory allocation, total allocated space must be greater "
|
|
"than 0, "
|
|
"but received %d."
|
|
"This is mainly caused by the size of 'label' being too large.",
|
|
(4 * num_buffer_ele + 3 * (nranks + 1) + num_temp_ele)));
|
|
|
|
// step 3: Alloc buffer memory so that we can reuse allocated memory
|
|
MemoryBuffer<T, Context> memory_buffer =
|
|
MemoryBuffer<T, Context>(num_buffer_ele, num_temp_ele, nranks, dev_ctx);
|
|
|
|
T* cub_sort_keys_ptr = memory_buffer.cub_sort_keys_ptr();
|
|
T* cub_sort_keys_out_ptr = memory_buffer.cub_sort_keys_out_ptr();
|
|
T* cub_sort_values_ptr = memory_buffer.cub_sort_values_ptr();
|
|
T* cub_sort_values_out_ptr = memory_buffer.cub_sort_values_out_ptr();
|
|
T* bound_index_ptr = memory_buffer.bound_index_ptr();
|
|
T* bound_value_ptr = memory_buffer.bound_value_ptr();
|
|
T* class_interval_ptr = memory_buffer.class_interval_ptr();
|
|
void* cub_temp_storage_ptr = memory_buffer.cub_temp_storage_ptr();
|
|
|
|
// step 4: Calculate class interval among nranks
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceScan::InclusiveSum(cub_temp_storage_ptr,
|
|
cub_temp_storage_bytes,
|
|
num_classes_per_device_ptr,
|
|
class_interval_ptr,
|
|
nranks + 1,
|
|
dev_ctx.stream())));
|
|
|
|
// step 5: random sample negative class center
|
|
uint64_t seed_data;
|
|
uint64_t increment;
|
|
int vec_size = VectorizedSize<T>(cub_sort_keys_ptr);
|
|
auto offset = ((num_classes - 1) /
|
|
(NumBlocks(num_classes) * kNumCUDAThreads * vec_size) +
|
|
1) *
|
|
vec_size;
|
|
// auto gen_cuda = DefaultCUDAGenerator(device_id);
|
|
auto gen_cuda = dev_ctx.GetGenerator();
|
|
if (!fix_seed) {
|
|
auto seed_offset = gen_cuda->IncrementOffset(offset);
|
|
seed_data = seed_offset.first;
|
|
increment = seed_offset.second;
|
|
} else {
|
|
seed_data = seed + rank;
|
|
increment = offset;
|
|
}
|
|
RandomSampleClassCenter<T>
|
|
<<<NumBlocks(num_classes), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
|
num_classes, seed_data, increment, num_classes, cub_sort_keys_ptr);
|
|
|
|
// step 6: mark positive class center as negative value
|
|
// fill the sort values to index 0, 1, ..., batch_size-1
|
|
MarkPositiveClassCenter<T>
|
|
<<<NumBlocks(batch_size), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
|
batch_size,
|
|
rank,
|
|
class_interval_ptr,
|
|
num_classes,
|
|
label.data<T>(),
|
|
cub_sort_keys_ptr);
|
|
Range<T><<<NumBlocks(num_buffer_ele), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
|
num_buffer_ele, cub_sort_values_ptr);
|
|
|
|
// step 7: sort class center by ascending, so that positive class center
|
|
// always be sampled.
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceRadixSort::SortPairs<T, T>(cub_temp_storage_ptr,
|
|
cub_temp_storage_bytes,
|
|
cub_sort_keys_ptr,
|
|
cub_sort_keys_out_ptr,
|
|
cub_sort_values_ptr,
|
|
cub_sort_values_out_ptr,
|
|
num_classes,
|
|
0,
|
|
sizeof(T) * 8,
|
|
dev_ctx.stream())));
|
|
|
|
// step 8: sort input label ascending
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceRadixSort::SortPairs<T, T>(cub_temp_storage_ptr,
|
|
cub_temp_storage_bytes,
|
|
label.data<T>(),
|
|
cub_sort_keys_out_ptr,
|
|
cub_sort_values_ptr,
|
|
cub_sort_keys_ptr,
|
|
batch_size,
|
|
0,
|
|
sizeof(T) * 8,
|
|
dev_ctx.stream())));
|
|
|
|
// step 9: Calculate new index using InclusiveSum on ascending sorted input
|
|
// label
|
|
NotEqualToPreviousAdjacentIterator<T> unique_counting_iter(
|
|
cub_sort_keys_out_ptr, 0);
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceScan::InclusiveSum<NotEqualToPreviousAdjacentIterator<T>, T*>(
|
|
cub_temp_storage_ptr,
|
|
cub_temp_storage_bytes,
|
|
unique_counting_iter,
|
|
cub_sort_values_ptr,
|
|
batch_size,
|
|
dev_ctx.stream())));
|
|
|
|
// step 10: Calculate new class center bound among ranks
|
|
GetClassCenterBound<T>
|
|
<<<NumBlocks(batch_size), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
|
batch_size,
|
|
nranks,
|
|
class_interval_ptr,
|
|
cub_sort_keys_out_ptr,
|
|
cub_sort_values_ptr,
|
|
bound_index_ptr,
|
|
bound_value_ptr);
|
|
|
|
// step 11: Calculate actual number of sampled class per device.
|
|
// Since maybe num_positive_class_center > num_samples,
|
|
// we need to ensure all positive class center per device are sampled.
|
|
ActualNumSampledFunctor<T> actual_num_sampled_op(num_samples);
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceScan::InclusiveScan(cub_temp_storage_ptr,
|
|
cub_temp_storage_bytes,
|
|
bound_value_ptr,
|
|
num_classes_per_device_ptr,
|
|
actual_num_sampled_op,
|
|
nranks + 1,
|
|
dev_ctx.stream())));
|
|
|
|
// step 12: Calculate actual sampled class interval among nranks
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
(cub::DeviceScan::InclusiveSum(cub_temp_storage_ptr,
|
|
cub_temp_storage_bytes,
|
|
num_classes_per_device_ptr,
|
|
class_interval_ptr,
|
|
nranks + 1,
|
|
dev_ctx.stream())));
|
|
|
|
// step 13: Get remapped label for output
|
|
GetRemappedLabel<T>
|
|
<<<NumBlocks(batch_size), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
|
batch_size,
|
|
nranks,
|
|
class_interval_ptr,
|
|
bound_index_ptr,
|
|
bound_value_ptr,
|
|
cub_sort_keys_ptr,
|
|
cub_sort_values_ptr,
|
|
dev_ctx.template Alloc<T>(remapped_label));
|
|
|
|
// step 14: Get sampled class center for output
|
|
Copy<Context>(dev_ctx,
|
|
num_classes_per_device,
|
|
CPUPlace(),
|
|
true,
|
|
&num_classes_per_device);
|
|
T actual_num_samples = num_classes_per_device.data<T>()[rank + 1];
|
|
sampled_local_class_center->Resize({actual_num_samples});
|
|
|
|
T* sampled_local_class_center_ptr =
|
|
dev_ctx.template Alloc<T>(sampled_local_class_center);
|
|
memory_utils::Copy(dev_ctx.GetPlace(),
|
|
sampled_local_class_center_ptr,
|
|
dev_ctx.GetPlace(),
|
|
cub_sort_values_out_ptr,
|
|
actual_num_samples * sizeof(T),
|
|
nullptr);
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(class_center_sample,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::ClassCenterSampleKernel,
|
|
int64_t,
|
|
int) {}
|