aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add broadcasted cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-16 17:23:38 -03:00
parent ddf04363d1
commit 54a5f81b3e
3 changed files with 83 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

63
norch/csrc/cuda.cu
View file

@ -58,6 +58,69 @@ __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
cudaDeviceSynchronize();
}
__global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* shape1, int* shape2, int* broadcasted_shape, int ndim1, int ndim2, int size) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= size) return;
int idx1 = 0, idx2 = 0;
int stride1 = 1, stride2 = 1;
int linear_idx = i;
for (int j = max(ndim1, ndim2) - 1; j >= 0; j--) {
int dim1 = j < ndim1 ? shape1[ndim1 - 1 - j] : 1;
int dim2 = j < ndim2 ? shape2[ndim2 - 1 - j] : 1;
int broadcasted_dim = broadcasted_shape[j];
int pos = linear_idx % broadcasted_dim;
linear_idx /= broadcasted_dim;
if (dim1 > 1) {
idx1 += pos * stride1;
}
if (dim2 > 1) {
idx2 += pos * stride2;
}
stride1 *= dim1;
stride2 *= dim2;
}
result_data[i] = data1[idx1] + data2[idx2];
}
__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
int size = tensor1->size;
// Copy the shapes to device memory
int* d_shape1;
int* d_shape2;
int* d_broadcasted_shape;
int ndim1 = tensor1->ndim;
int ndim2 = tensor2->ndim;
cudaMalloc((void**)&d_shape1, ndim1 * sizeof(int));
cudaMalloc((void**)&d_shape2, ndim2 * sizeof(int));
cudaMalloc((void**)&d_broadcasted_shape, max(ndim1, ndim2) * sizeof(int));
cudaMemcpy(d_shape1, tensor1->shape, ndim1 * sizeof(int), cudaMemcpyHostToDevice);
cudaMemcpy(d_shape2, tensor2->shape, ndim2 * sizeof(int), cudaMemcpyHostToDevice);
cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max(ndim1, ndim2) * sizeof(int), cudaMemcpyHostToDevice);
int number_of_blocks = (size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
add_broadcasted_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, d_shape1, d_shape2, d_broadcasted_shape, ndim1, ndim2, size);
cudaError_t error = cudaGetLastError();
if (error != cudaSuccess) {
printf("CUDA error: %s\n", cudaGetErrorString(error));
exit(-1);
}
cudaDeviceSynchronize();
cudaFree(d_shape1);
cudaFree(d_shape2);
cudaFree(d_broadcasted_shape);
}
__global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size) {
__shared__ float partial_sum[THREADS_PER_BLOCK * sizeof(float)];

5
norch/csrc/cuda.h
View file

@ -6,7 +6,10 @@
__global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
__host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
__global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* shape1, int* shape2, int* broadcasted_shape, int ndim1, int ndim2, int size);
__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);
__global__ void sum_tensor_cuda_kernel(float* data, float* result_data);
__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data);

25
norch/csrc/tensor.cpp
View file

@ -149,16 +149,23 @@ extern "C" {
broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
}
// Allocate memory for result tensor
float* result_data = (float*)malloc(tensor1->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
if (strcmp(tensor1->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
add_broadcasted_tensor_cuda(tensor1, tensor2, result_data);
return create_tensor(result_data, shape, ndim, device);
}
else {
float* result_data = (float*)malloc(tensor1->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
Tensor* sum_tensor(Tensor* tensor, int axis) {