math_ext_kernel.h

Go to the documentation of this file.

 
#ifndef __NEKO_HIP_MATH_EXT_KERNELS__
#define __NEKO_HIP_MATH_EXT_KERNELS__
 
template <typename T>
__global__ void cadd_mask_kernel(
    T* __restrict__ a, const T c, const int size, int* __restrict__ mask,
    const int mask_size) {
 
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int str = blockDim.x * gridDim.x;
 
    for (int i = idx; i < mask_size; i += str) { a[mask[i]] = a[mask[i]] + c; }
}
 
template <typename T>
__global__ void invcol1_mask_kernel(
    T* __restrict__ a, const int size, int* __restrict__ mask,
    const int mask_size) {
 
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int str = blockDim.x * gridDim.x;
 
    for (int i = idx; i < mask_size; i += str) {
        a[mask[i]] = 1.0 / a[mask[i]];
    }
}
 
template <typename T>
__global__ void col2_mask_kernel(
    T* __restrict__ a, T* __restrict__ b, const int size,
    int* __restrict__ mask, const int mask_size) {
 
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int str = blockDim.x * gridDim.x;
 
    for (int i = idx; i < mask_size; i += str) {
        a[mask[i]] = a[mask[i]] * b[mask[i]];
    }
}
 
template <typename T>
__global__ void col3_mask_kernel(
    T* __restrict__ a, T* __restrict__ b, T* __restrict__ c, const int size,
    int* __restrict__ mask, const int mask_size) {
 
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int str = blockDim.x * gridDim.x;
 
    for (int i = idx; i < mask_size; i += str) {
        a[mask[i]] = b[mask[i]] * c[mask[i]];
    }
}
 
template <typename T>
__global__ void sub3_mask_kernel(
    T* __restrict__ a, T* __restrict__ b, T* __restrict__ c, const int size,
    int* __restrict__ mask, const int mask_size) {
 
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int str = blockDim.x * gridDim.x;
 
    for (int i = idx; i < mask_size; i += str) {
        a[mask[i]] = b[mask[i]] - c[mask[i]];
    }
}
 
#endif // __NEKO_CUDA_MATH_EXT_KERNELS__