1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81
|
#define FILENAME(line) FILENAME_FOR_EXCEPTIONS_CUDA("src/cuda-kernels/manual_awkward_IndexedOptionArray_rpad_and_clip_mask_axis1.cu", line)
#include "awkward/kernels.h"
#include "standard_parallel_algorithms.h"
__global__ void
awkward_IndexedOptionArray_rpad_and_clip_mask_axis1_filter_mask(
const int8_t* frommask,
int8_t* filtered_mask,
int64_t length) {
int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
if(thread_id < length) {
if (!frommask[thread_id]) {
filtered_mask[thread_id] = 1;
}
}
}
template <typename T>
__global__ void
awkward_IndexedOptionArray_rpad_and_clip_mask_axis1_kernel(
T* toindex,
const int8_t* frommask,
int64_t* prefixedsum_mask,
int64_t length) {
int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
if(thread_id < length) {
if (frommask[thread_id]) {
toindex[thread_id] = -1;
}
else {
toindex[thread_id] = prefixedsum_mask[thread_id] - 1;
}
}
}
template <typename T>
ERROR awkward_IndexedOptionArray_rpad_and_clip_mask_axis1(
T* toindex,
const int8_t* frommask,
int64_t length) {
dim3 blocks_per_grid = blocks(length);
dim3 threads_per_block = threads(length);
int8_t* filtered_mask;
int64_t* res_temp;
HANDLE_ERROR(cudaMalloc((void**)&filtered_mask, sizeof(int8_t) * length));
HANDLE_ERROR(cudaMalloc((void**)&res_temp, sizeof(int64_t) * length));
HANDLE_ERROR(cudaMemset(filtered_mask, 0, sizeof(int8_t) * length));
awkward_IndexedOptionArray_rpad_and_clip_mask_axis1_filter_mask<<<blocks_per_grid, threads_per_block>>>(
frommask,
filtered_mask,
length);
exclusive_scan(res_temp, filtered_mask, length);
awkward_IndexedOptionArray_rpad_and_clip_mask_axis1_kernel<<<blocks_per_grid, threads_per_block>>>(
toindex,
frommask,
res_temp,
length);
cudaDeviceSynchronize();
return success();
}
ERROR awkward_IndexedOptionArray_rpad_and_clip_mask_axis1_64(
int64_t* toindex,
const int8_t* frommask,
int64_t length) {
return awkward_IndexedOptionArray_rpad_and_clip_mask_axis1<int64_t>(
toindex,
frommask,
length);
}
|