heaviside_mapping_kernel.h

Go to the documentation of this file.

 
#ifndef __NEKO_CUDA_HEAVISIDE_MAPPING_KERNELS__
#define __NEKO_CUDA_HEAVISIDE_MAPPING_KERNELS__
 
#include <math.h>
 
template <typename T>
__global__ void heaviside_mapping_apply_kernel(
    const T beta, const T eta, T* __restrict__ X_out_d,
    T* __restrict__ X_in_d, const int n) {
 
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int str = blockDim.x * gridDim.x;
    const T den = tanh(beta * eta) + tanh(beta * (1.0 - eta));
    const T tanh_beta_eta = tanh(beta * eta);
 
    for (int i = idx; i < n; i += str) {
        X_out_d[i] = (tanh_beta_eta + tanh(beta * (X_in_d[i] - eta))) / den;
    }
}
 
template <typename T>
__global__ void heaviside_mapping_apply_backward_kernel(
    const T beta, const T eta, T* __restrict__ sens_out_d,
    T* __restrict__ sens_in_d, T* __restrict__ X_in_d, const int n) {
 
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int str = blockDim.x * gridDim.x;
    const T den = tanh(beta * eta) + tanh(beta * (1.0 - eta));
 
    for (int i = idx; i < n; i += str) {
        const T arg = beta * (X_in_d[i] - eta);
        const T tanh_arg = tanh(arg);
        const T dproj = beta * (1.0 - tanh_arg * tanh_arg) / den;
        sens_out_d[i] = dproj * sens_in_d[i];
    }
}
 
#endif // __NEKO_CUDA_HEAVISIDE_MAPPING_KERNELS__