123456789101112131415161718192021222324252627282930313233343536373839 |
- #include "custom_cuda_layers.h"
- __global__ void param_update_kernel(const float* input, __half* output, int size)
- {
- int id = blockIdx.x * blockDim.x + threadIdx.x;
- if (id < size) { output[id] = (__half)input[id]; }
- }
- void launch_param_update(const float* input, __half* output, int size, cudaStream_t stream)
- {
- int threads = 1024;
- dim3 grid_dim((size - 1) / threads + 1);
- dim3 block_dim(threads);
- param_update_kernel<<<grid_dim, block_dim, 0, stream>>>(input, output, size);
- }
- __global__ void param_update_kernel_half(const float* input, __half* output, int size)
- {
- int id = blockIdx.x * blockDim.x + threadIdx.x;
- __half2* output_cast = reinterpret_cast<__half2*>(output);
- if (id < size) {
- float input_f = input[id];
- __half2* input_h = reinterpret_cast<__half2*>(&input_f);
- output_cast[id] = *input_h;
- }
- }
- void launch_param_update_half(const float* input, __half* output, int size, cudaStream_t stream)
- {
- int threads = 1024;
- size /= 2;
- dim3 grid_dim((size - 1) / threads + 1);
- dim3 block_dim(threads);
- param_update_kernel_half<<<grid_dim, block_dim, 0, stream>>>(input, output, size);
- }
|