Optimize the unary cuda kernels for the contiguous case.

This commit is contained in:
laurent
2023-06-23 18:40:15 +01:00
parent 4f9f14a06b
commit 5ca309ecb0
3 changed files with 43 additions and 13 deletions

View File

@ -1,6 +1,22 @@
#include "cuda_fp16.h"
#include "compatibility.cuh"
__device__ bool is_contiguous(
const size_t num_dims,
const size_t *dims,
const size_t *strides
) {
size_t acc = 1;
for (unsigned int d = 0; d < num_dims; d++) {
unsigned int dim_idx = num_dims - 1 - d;
if (acc != strides[dim_idx]) {
return false;
}
acc *= dims[dim_idx];
}
return true;
}
__device__ unsigned int get_strided_index(
unsigned int idx,
const size_t num_dims,