diff --git a/build/cpu.o b/build/cpu.o index 054ce2e..0b75e02 100644 Binary files a/build/cpu.o and b/build/cpu.o differ diff --git a/build/cuda.cu.o b/build/cuda.cu.o index 9809e05..7f95837 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/build/tensor.o b/build/tensor.o index 5df7819..b8560d8 100644 Binary files a/build/tensor.o and b/build/tensor.o differ diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index d3ad7c7..939d371 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc index b9738eb..53f8f0f 100644 Binary files a/norch/autograd/__pycache__/functions.cpython-38.pyc and b/norch/autograd/__pycache__/functions.cpython-38.pyc differ diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py index 0023139..d0aed64 100644 --- a/norch/autograd/functions.py +++ b/norch/autograd/functions.py @@ -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): diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index a2f107b..3b4e303 100644 --- a/norch/csrc/cpu.cpp +++ b/norch/csrc/cpu.cpp @@ -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--) { diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h index 2a97fd2..be875ed 100644 --- a/norch/csrc/cpu.h +++ b/norch/csrc/cpu.h @@ -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); diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index 9ebe6d3..467b8cc 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -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<<>>(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<<>>(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<<>>(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; diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h index 5d5e353..f7a883b 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -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); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 730d204..824737d 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -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) { diff --git a/norch/libtensor.so b/norch/libtensor.so index 9f15e45..a125f3c 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 7635e58..8cd105e 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -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): diff --git a/tests/test_operations.py b/tests/test_operations.py index e7dd833..d6552a6 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -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): """