diff --git a/kernels/src/binary_add.cu b/kernels/src/binary_add.cu index 16156bc2..979a429e 100644 --- a/kernels/src/binary_add.cu +++ b/kernels/src/binary_add.cu @@ -1,20 +1,8 @@ #include "binary_op_macros.cuh" -struct BinaryAddOp {}; - #if __CUDA_ARCH__ >= 530 -BINARY_OP(__half, badd_fwd_f16, badd_bwd_lhs_f16, badd_bwd_rhs_f16, BinaryAddOp, - x + y, - 1.0, - 1.0) +BINARY_OP(__half, badd_f16, x + y) #endif -BINARY_OP(float, badd_fwd_f32, badd_bwd_lhs_f32, badd_bwd_rhs_f32, BinaryAddOp, - x + y, - 1.0, - 1.0) - -BINARY_OP(double, badd_fwd_f64, badd_bwd_lhs_f64, badd_bwd_rhs_f64, BinaryAddOp, - x + y, - 1.0, - 1.0) +BINARY_OP(float, badd_f32, x + y) +BINARY_OP(double, badd_fwd_f64, x + y); diff --git a/kernels/src/binary_div.cu b/kernels/src/binary_div.cu index a36c58b8..410eed62 100644 --- a/kernels/src/binary_div.cu +++ b/kernels/src/binary_div.cu @@ -1,21 +1,8 @@ #include "binary_op_macros.cuh" -struct BinaryDivOp {}; - #if __CUDA_ARCH__ >= 530 -BINARY_OP(__half, bdiv_fwd_f16, bdiv_bwd_lhs_f16, bdiv_bwd_rhs_f16, BinaryDivOp, - x / y, - recipg(y), - -x / (y * y)) +BINARY_OP(__half, bdiv_f16, x / y) #endif -BINARY_OP(float, bdiv_fwd_f32, bdiv_bwd_lhs_f32, bdiv_bwd_rhs_f32, BinaryDivOp, - x / y, - recipg(y), - -x / (y * y)) - -BINARY_OP(double, bdiv_fwd_f64, bdiv_bwd_lhs_f64, bdiv_bwd_rhs_f64, BinaryDivOp, - x / y, - recipg(y), - -x / (y * y)) - +BINARY_OP(float, bdiv_f32, x / y) +BINARY_OP(double, bdiv_f64, x / y); diff --git a/kernels/src/binary_mul.cu b/kernels/src/binary_mul.cu index 9a6d1430..8adab74b 100644 --- a/kernels/src/binary_mul.cu +++ b/kernels/src/binary_mul.cu @@ -1,21 +1,8 @@ #include "binary_op_macros.cuh" -struct BinaryMulKernalOp {}; - #if __CUDA_ARCH__ >= 530 -BINARY_OP(__half, bmul_fwd_f16, bmul_bwd_lhs_f16, bmul_bwd_rhs_f16, BinaryMulKernalOp, - x * y, - y, - x) +BINARY_OP(__half, bmul_f16, x * y) #endif -BINARY_OP(float, bmul_fwd_f32, bmul_bwd_lhs_f32, bmul_bwd_rhs_f32, BinaryMulKernalOp, - x * y, - y, - x) - -BINARY_OP(double, bmul_fwd_f64, bmul_bwd_lhs_f64, bmul_bwd_rhs_f64, BinaryMulKernalOp, - x * y, - y, - x) - +BINARY_OP(float, bmul_f32, x * y) +BINARY_OP(double, bmul_f64, x * y); diff --git a/kernels/src/binary_op_macros.cuh b/kernels/src/binary_op_macros.cuh index b79a4d81..8d4954f5 100644 --- a/kernels/src/binary_op_macros.cuh +++ b/kernels/src/binary_op_macros.cuh @@ -1,18 +1,16 @@ #include "cuda_utils.cuh" -#define LONG_BINARY_OP(TYPENAME, FORWARD, BACKWARD_LHS, BACKWARD_RHS, OP_STRUCT, FUNC, DFDX, DFDY) \ -extern "C" __global__ void FORWARD( \ - const OP_STRUCT op, \ +#define BINARY_OP(TYPENAME, FN_NAME, FUNC) \ +extern "C" __global__ void FN_NAME( \ const size_t numel, \ const size_t num_dims, \ - const size_t *info, \ + const size_t *dims, \ + const size_t *lhs_strides, \ + const size_t *rhs_strides, \ const TYPENAME *lhs, \ const TYPENAME *rhs, \ TYPENAME *out \ ) { \ - const size_t *dims = info; \ - const size_t *lhs_strides = info + num_dims; \ - const size_t *rhs_strides = info + 2 * num_dims; \ for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ unsigned int tmp_i = i; \ unsigned int lhs_i = 0; \ @@ -25,77 +23,6 @@ extern "C" __global__ void FORWARD( \ } \ TYPENAME x = lhs ? lhs[lhs_i] : out[i]; \ TYPENAME y = rhs ? rhs[rhs_i] : out[i]; \ - TYPENAME fx; \ - FUNC\ - out[i] = fx; \ + out[i] = FUNC; \ } \ } \ -\ -extern "C" __global__ void BACKWARD_LHS( \ - const OP_STRUCT op, \ - const size_t numel, \ - const size_t num_dims, \ - const size_t *info, \ - const TYPENAME *lhs, \ - TYPENAME *grad_lhs, \ - const size_t chunk_len, \ - const TYPENAME *rhs, \ - const TYPENAME *grad_out \ -) { \ - const size_t *dims = info + 0 * num_dims; \ - const size_t *out_strides = info + 1 * num_dims; \ - const size_t *rhs_strides = info + 2 * num_dims; \ - TYPENAME zero = 0.0; \ - for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ - unsigned int tmp_i = i; \ - unsigned int out_i = 0; \ - unsigned int rhs_i = 0; \ - for (int d = num_dims - 1; d >= 0; d--) { \ - unsigned int i_dim = tmp_i % dims[d]; \ - out_i += i_dim * out_strides[d]; \ - rhs_i += i_dim * rhs_strides[d]; \ - tmp_i /= dims[d]; \ - } \ - TYPENAME x = lhs ? lhs[i / chunk_len] : zero; \ - TYPENAME y = rhs ? rhs[rhs_i] : zero; \ - TYPENAME go = grad_out[out_i]; \ - TYPENAME dfdx = (DFDX); \ - chunk_sum(chunk_len, dfdx * go, grad_lhs); \ - } \ -} \ -\ -extern "C" __global__ void BACKWARD_RHS( \ - const OP_STRUCT op, \ - const size_t numel, \ - const size_t num_dims, \ - const size_t *info, \ - const TYPENAME *lhs, \ - const TYPENAME *rhs, \ - TYPENAME *grad_rhs, \ - const size_t chunk_len, \ - const TYPENAME *grad_out \ -) { \ - const size_t *dims = info + 3 * num_dims; \ - const size_t *out_strides = info + 4 * num_dims; \ - const size_t *lhs_strides = info + 5 * num_dims; \ - TYPENAME zero = 0.0; \ - for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ - unsigned int tmp_i = i; \ - unsigned int lhs_i = 0; \ - unsigned int out_i = 0; \ - for (int d = num_dims - 1; d >= 0; d--) { \ - unsigned int i_dim = tmp_i % dims[d]; \ - lhs_i += i_dim * lhs_strides[d]; \ - out_i += i_dim * out_strides[d]; \ - tmp_i /= dims[d]; \ - } \ - TYPENAME x = lhs ? lhs[lhs_i] : zero; \ - TYPENAME y = rhs ? rhs[i / chunk_len] : zero; \ - TYPENAME go = grad_out[out_i]; \ - TYPENAME dfdy = (DFDY); \ - chunk_sum(chunk_len, dfdy * go, grad_rhs); \ - } \ -} - -#define BINARY_OP(TYPENAME, FORWARD, BACKWARD_LHS, BACKWARD_RHS, OP_STRUCT, FUNC, DFDX, DFDY) \ - LONG_BINARY_OP(TYPENAME, FORWARD, BACKWARD_LHS, BACKWARD_RHS, OP_STRUCT, fx = (FUNC);, DFDX, DFDY) diff --git a/kernels/src/binary_sub.cu b/kernels/src/binary_sub.cu index f4c93972..891932db 100644 --- a/kernels/src/binary_sub.cu +++ b/kernels/src/binary_sub.cu @@ -1,21 +1,8 @@ #include "binary_op_macros.cuh" -struct BinarySubKernelOp {}; - #if __CUDA_ARCH__ >= 530 -BINARY_OP(__half, bsub_fwd_f16, bsub_bwd_lhs_f16, bsub_bwd_rhs_f16, BinarySubKernelOp, - x - y, - 1.0, - -1.0) +BINARY_OP(__half, bsub_f16, x - y) #endif -BINARY_OP(float, bsub_fwd_f32, bsub_bwd_lhs_f32, bsub_bwd_rhs_f32, BinarySubKernelOp, - x - y, - 1.0, - -1.0) - -BINARY_OP(double, bsub_fwd_f64, bsub_bwd_lhs_f64, bsub_bwd_rhs_f64, BinarySubKernelOp, - x - y, - 1.0, - -1.0) - +BINARY_OP(float, bsub_f32, x - y) +BINARY_OP(double, bsub_f64, x - y);