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

sub broadcasted cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-17 13:15:33 -03:00
parent 64fb3f9edb
commit de72cd67ba
14 changed files with 128 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
build/cpu.o
View file

Binary file not shown.

BIN
build/cuda.cu.o
View file

Binary file not shown.

BIN
build/tensor.o
View file

Binary file not shown.

BIN
norch/__pycache__/tensor.cpython-38.pyc
View file

Binary file not shown.

BIN
norch/autograd/__pycache__/functions.cpython-38.pyc
View file

Binary file not shown.

14
norch/autograd/functions.py
View file

@ -29,7 +29,13 @@ class AddBroadcastedBackward:
return gradient
class SubBackward:
def __init__(self, x, y):
self.input = [x, y]
def backward(self, gradient):
return [gradient, -gradient]
class SubBroadcastedBackward:
def __init__(self, x, y):
self.input = [x, y]
@ -49,15 +55,7 @@ class SubBroadcastedBackward:
for i in range(len(shape)):
if shape[i] == 1:
gradient = gradient.sum(axis=i)
return gradient
class SubBackward:
def __init__(self, x, y):
self.input = [x, y]
def backward(self, gradient):
return [gradient, -gradient]
class ScalarMulBackward:
def __init__(self, x, scalar):

8
norch/csrc/cpu.cpp
View file

@ -58,7 +58,7 @@ void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
}
}
void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
// Calculate strides for broadcasting
@ -75,12 +75,12 @@ void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_
int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
stride1 *= broadcasted_shape[i];
stride2 *= broadcasted_shape[i];
stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
}
// Perform element-wise addition with broadcasting
for (int i = 0; i < tensor1->size; i++) {
for (int i = 0; i < broadcasted_size; i++) {
int index1 = 0, index2 = 0;
int linear_index = i;
for (int j = max_ndim - 1; j >= 0; j--) {

2
norch/csrc/cpu.h
View file

@ -7,7 +7,7 @@ void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis);
void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);
void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
void scalar_div_tensor_cpu(float scalar, Tensor* tensor, float* result_data);
void tensor_div_scalar_cpu(Tensor* tensor, float scalar, float* result_data);

71
norch/csrc/cuda.cu
View file

@ -71,7 +71,7 @@ __global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, f
result_data[i] = data1[index1] + data2[index2];
}
__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
int* strides1 = (int*)malloc(max_ndim * sizeof(int));
@ -87,8 +87,8 @@ __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, floa
int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
stride1 *= broadcasted_shape[i];
stride2 *= broadcasted_shape[i];
stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
}
int* d_broadcasted_shape;
@ -104,8 +104,8 @@ __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, floa
cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int));
cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
int number_of_blocks = (tensor1->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_broadcasted_shape, d_strides1, d_strides2, max_ndim, tensor1->size);
int number_of_blocks = (broadcasted_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_broadcasted_shape, d_strides1, d_strides2, max_ndim, broadcasted_size);
cudaError_t error = cudaGetLastError();
if (error != cudaSuccess) {
@ -187,6 +187,67 @@ __host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
cudaDeviceSynchronize();
}
__global__ void sub_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= size) return;
int index1 = 0, index2 = 0;
int linear_index = i;
for (int j = max_ndim - 1; j >= 0; j--) {
int pos = linear_index % broadcasted_shape[j];
linear_index /= broadcasted_shape[j];
if (strides1[j] != 0) index1 += pos * strides1[j];
if (strides2[j] != 0) index2 += pos * strides2[j];
}
result_data[i] = data1[index1] - data2[index2];
}
__host__ void sub_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
int* strides1 = (int*)malloc(max_ndim * sizeof(int));
int* strides2 = (int*)malloc(max_ndim * sizeof(int));
if (strides1 == NULL || strides2 == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
int stride1 = 1, stride2 = 1;
for (int i = max_ndim - 1; i >= 0; i--) {
int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - max_ndim + i] : 1;
int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
}
int* d_broadcasted_shape;
int* d_strides1;
int* d_strides2;
cudaMalloc((void**)&d_broadcasted_shape, max_ndim * sizeof(int));
cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
cudaMalloc((void**)&d_strides1, max_ndim * sizeof(int));
cudaMemcpy(d_strides1, strides1, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int));
cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
int number_of_blocks = (broadcasted_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
sub_broadcasted_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, broadcasted_size);
cudaError_t error = cudaGetLastError();
if (error != cudaSuccess) {
printf("CUDA error: %s\n", cudaGetErrorString(error));
exit(-1);
}
cudaDeviceSynchronize();
cudaFree(d_broadcasted_shape);
}
__global__ void elementwise_mul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
int i = blockIdx.x * blockDim.x + threadIdx.x;

5
norch/csrc/cuda.h
View file

@ -8,7 +8,10 @@
__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* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size);
__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);
__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
__global__ void sub_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size);
__host__ void sub_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
__global__ void sum_tensor_cuda_kernel(float* data, float* result_data);
__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data);

27
norch/csrc/tensor.cpp
View file

@ -157,7 +157,7 @@ extern "C" {
if (strcmp(tensor1->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void **)&result_data, broadcasted_size * sizeof(float));
add_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape);
add_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
else {
@ -295,16 +295,27 @@ 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);
int broadcasted_size = 1;
for (int i = 0; i < max_ndim; i++) {
broadcasted_size *= broadcasted_shape[i];
}
sub_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
if (strcmp(tensor1->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void **)&result_data, broadcasted_size * sizeof(float));
sub_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
else {
float* result_data = (float*)malloc(broadcasted_size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
sub_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
}
Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2) {

BIN
norch/libtensor.so
View file

Binary file not shown.

21
tests/test_autograd.py
View file

@ -58,7 +58,7 @@ class TestTensorAutograd(unittest.TestCase):
self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
def test_broadcasting_addition_autograd(self):
def test_broadcasted_addition_autograd(self):
"""
Test autograd for broadcasting addition: tensor1 + tensor2
"""
@ -115,7 +115,7 @@ class TestTensorAutograd(unittest.TestCase):
self.assertTrue(utils.compare_torch(norch_tensor1_grad_sub, torch_tensor1_grad_sub))
self.assertTrue(utils.compare_torch(norch_tensor2_grad_sub, torch_tensor2_grad_sub))
def test_broadcasting_subtraction_autograd(self):
def test_broadcasted_subtraction_autograd(self):
"""
Test autograd for broadcasting subtraction: tensor1 - tensor2
"""
@ -135,6 +135,23 @@ class TestTensorAutograd(unittest.TestCase):
self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
# reversed order broadcasting
norch_result = (norch_tensor2 - norch_tensor1).sum()
norch_result.backward()
norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device) # Shape (1, 2, 3)
torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device) # Shape (3)
torch_result = (torch_tensor2 - torch_tensor1).sum()
torch_result.backward()
torch_tensor1_grad = torch_tensor1.grad
torch_tensor2_grad = torch_tensor2.grad
self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
def test_division(self):

10
tests/test_operations.py
View file

@ -52,6 +52,7 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
# reversed order broadcasting
norch_tensor1 = norch.Tensor([[0, 2]]).to(self.device)
norch_tensor2 = norch.Tensor([[3, 4], [5, -1]]).to(self.device)
norch_result = norch_tensor1 + norch_tensor2
@ -61,7 +62,6 @@ class TestTensorOperations(unittest.TestCase):
torch_tensor2 = torch.tensor([[3, 4], [5, -1]]).to(self.device)
torch_expected = torch_tensor1 + torch_tensor2
print(torch_result, torch_expected)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
norch_result = norch_tensor2 + norch_tensor1
@ -101,6 +101,14 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
# reversed order broadcasting
norch_result = norch_tensor2 - norch_tensor1
torch_result = utils.to_torch(norch_result).to(self.device)
torch_expected = torch_tensor2 - torch_tensor1
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_division_by_scalar(self):
"""