diff options
Diffstat (limited to 'candle-kernels/src')
| -rw-r--r-- | candle-kernels/src/cast.cu | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/candle-kernels/src/cast.cu b/candle-kernels/src/cast.cu index ea611eba..0c4ddbc6 100644 --- a/candle-kernels/src/cast.cu +++ b/candle-kernels/src/cast.cu @@ -13,13 +13,13 @@ extern "C" __global__ void FN_NAME( \ const size_t *strides = info + num_dims; \ if (is_contiguous(num_dims, dims, strides)) { \ for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ - out[i] = inp[i]; \ + out[i] = (DST_TYPENAME) inp[i]; \ } \ } \ else { \ for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ - out[i] = inp[strided_i]; \ + out[i] = (DST_TYPENAME) inp[strided_i]; \ } \ } \ } \ @@ -29,12 +29,12 @@ CAST_OP(__nv_bfloat16, __nv_bfloat16, cast_bf16_bf16) // CAST_OP(__nv_bfloat16, uint8_t, cast_bf16_u8) CAST_OP(__nv_bfloat16, uint32_t, cast_bf16_u32) -// CAST_OP(__nv_bfloat16, __half, cast_bf16_f16) +CAST_OP(__nv_bfloat16, __half, cast_bf16_f16) CAST_OP(__nv_bfloat16, float, cast_bf16_f32) CAST_OP(__nv_bfloat16, double, cast_bf16_f64) CAST_OP(uint8_t, __nv_bfloat16, cast_u8_bf16) CAST_OP(uint32_t, __nv_bfloat16, cast_u32_bf16) -// CAST_OP(__half, __nv_bfloat16, cast_f16_bf16) +CAST_OP(__half, __nv_bfloat16, cast_f16_bf16) CAST_OP(float, __nv_bfloat16, cast_f32_bf16) CAST_OP(double, __nv_bfloat16, cast_f64_bf16) #endif |