mirror of
https://github.com/huggingface/candle.git
synced 2025-06-17 02:58:50 +00:00
70 lines
2.6 KiB
Plaintext
70 lines
2.6 KiB
Plaintext
#include "cuda_utils.cuh"
|
|
|
|
#define BINARY_OP_OUT(TYPENAME, OUT_TYPENAME, FN_NAME, FUNC) \
|
|
extern "C" __global__ void FN_NAME( \
|
|
const size_t numel, \
|
|
const size_t num_dims, \
|
|
const size_t *dims_and_strides, \
|
|
const TYPENAME *lhs, \
|
|
const TYPENAME *rhs, \
|
|
OUT_TYPENAME *out \
|
|
) { \
|
|
const size_t *dims = dims_and_strides; \
|
|
const size_t *lhs_strides = dims_and_strides + 1 * num_dims; \
|
|
const size_t *rhs_strides = dims_and_strides + 2 * num_dims; \
|
|
bool lhs_cont = is_contiguous(num_dims, dims, lhs_strides); \
|
|
bool rhs_cont = is_contiguous(num_dims, dims, rhs_strides); \
|
|
if (lhs_cont && rhs_cont) { \
|
|
for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
|
|
TYPENAME x = lhs[i]; \
|
|
TYPENAME y = rhs[i]; \
|
|
out[i] = FUNC; \
|
|
} \
|
|
} else if (lhs_cont) { \
|
|
for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
|
|
unsigned int tmp_i = i; \
|
|
unsigned int rhs_i = 0; \
|
|
for (int d = num_dims - 1; d >= 0; d--) { \
|
|
unsigned int i_dim = tmp_i % dims[d]; \
|
|
rhs_i += i_dim * rhs_strides[d]; \
|
|
tmp_i /= dims[d]; \
|
|
} \
|
|
TYPENAME x = lhs[i]; \
|
|
TYPENAME y = rhs[rhs_i]; \
|
|
out[i] = FUNC; \
|
|
} \
|
|
} else if (rhs_cont) { \
|
|
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; \
|
|
for (int d = num_dims - 1; d >= 0; d--) { \
|
|
unsigned int i_dim = tmp_i % dims[d]; \
|
|
lhs_i += i_dim * lhs_strides[d]; \
|
|
tmp_i /= dims[d]; \
|
|
} \
|
|
TYPENAME x = lhs[lhs_i]; \
|
|
TYPENAME y = rhs[i]; \
|
|
out[i] = FUNC; \
|
|
} \
|
|
} else { \
|
|
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 rhs_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]; \
|
|
rhs_i += i_dim * rhs_strides[d]; \
|
|
tmp_i /= dims[d]; \
|
|
} \
|
|
TYPENAME x = lhs[lhs_i]; \
|
|
TYPENAME y = rhs[rhs_i]; \
|
|
out[i] = FUNC; \
|
|
} \
|
|
} \
|
|
} \
|
|
|
|
|
|
#define BINARY_OP(TYPENAME, FN_NAME, FUNC) \
|
|
BINARY_OP_OUT(TYPENAME, TYPENAME, FN_NAME, FUNC)
|