Merge pull request #53 from lucasdelimanogueira/tmp

Tmp

norch/autograd/functions.py

@@ -8,11 +8,50 @@ class AddBackward:
         return [gradient, gradient]
     
 class AddBroadcastedBackward:
-    pass
+    def __init__(self, x, y):
+        self.input = [x, y]
 
+    def backward(self, gradient):
+        x, y = self.input
+        grad_x = self._reshape_gradient(gradient, x.shape)
+        grad_y = self._reshape_gradient(gradient, y.shape)
+        return [grad_x, grad_y]
+    
+    def _reshape_gradient(self, gradient, shape):
+        # Reduce gradient dimensions to match the target shape dimensions
+        while len(gradient.shape) > len(shape):
+            gradient = gradient.sum(axis=0)
 
+        # Sum along axes where the target shape dimension is 1
+        for i in range(len(shape)):
+            if shape[i] == 1:
+                gradient = gradient.sum(axis=i)
+
+        return gradient
     
+
+class SubBroadcastedBackward:
+    def __init__(self, x, y):
+        self.input = [x, y]
+
+    def backward(self, gradient):
+        x, y = self.input
+        grad_x = self._reshape_gradient(gradient, x.shape)
+        grad_y = self._reshape_gradient(gradient, y.shape)
+        return [grad_x, -grad_y]
     
+    def _reshape_gradient(self, gradient, shape):
+        # Reduce gradient dimensions to match the target shape dimensions
+        while len(gradient.shape) > len(shape):
+            gradient = gradient.sum(axis=0)
+
+        # Sum along axes where the target shape dimension is 1
+        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]

norch/csrc/cpu.cpp

@@ -58,6 +58,45 @@ 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) {
+    int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
+
+    // Calculate strides for broadcasting
+    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 *= broadcasted_shape[i];
+        stride2 *= broadcasted_shape[i];
+    }
+
+    // Perform element-wise addition with broadcasting
+    for (int i = 0; i < tensor1->size; i++) {
+        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] = tensor1->data[index1] - tensor2->data[index2];
+    }
+
+    // Free strides
+    free(strides1);
+    free(strides2);
+}
+
 void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
     
     for (int i = 0; i < tensor1->size; i++) {

norch/csrc/cpu.h

@@ -7,6 +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);
 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 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);

norch/csrc/cuda.cu

@@ -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)];

norch/csrc/cuda.h

@@ -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);
 

norch/csrc/tensor.cpp

@@ -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) {
@@ -259,6 +266,43 @@ extern "C" {
         }
     }
 
+    Tensor* sub_broadcasted_tensor(Tensor* tensor1, Tensor* tensor2) {
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
+
+        // Determine the broadcasted shape
+        int* broadcasted_shape = (int*)malloc(max_ndim * sizeof(int));
+        if (broadcasted_shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        for (int i = 0; i < max_ndim; i++) {
+            int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - 1 - i] : 1;
+            int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - 1 - i] : 1;
+            if (dim1 != dim2 && dim1 != 1 && dim2 != 1) {
+                fprintf(stderr, "Shapes are not compatible for broadcasting\n");
+                exit(1);
+            }
+            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);
+        }
+
+        sub_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
+
+        return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
+    }
+
     Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2) {
         if (tensor1->ndim != tensor2->ndim) {
             fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for element-wise multiplication\n", tensor1->ndim, tensor2->ndim);

norch/tensor.py

@@ -312,24 +312,61 @@ class Tensor:
         if isinstance(other, (int, float)):
             other = other * self.ones_like()
 
-        if self.shape != other.shape:
-            raise ValueError("Tensors must have the same shape for subtraction")
+        # Function to determine if broadcasting is needed and get the broadcasted shape
+        def broadcast_shape(shape1, shape2):
+            if shape1 == shape2:
+                return shape1, False
+            
+            max_len = max(len(shape1), len(shape2))
+            shape1 = [1] * (max_len - len(shape1)) + shape1
+            shape2 = [1] * (max_len - len(shape2)) + shape2
+
+            broadcasted_shape = []
+            for dim1, dim2 in zip(shape1, shape2):
+                if dim1 != dim2 and dim1 != 1 and dim2 != 1:
+                    raise ValueError("Shapes are not compatible for broadcasting")
+                broadcasted_shape.append(max(dim1, dim2))
+            return broadcasted_shape, True
+
+        broadcasted_shape, needs_broadcasting = broadcast_shape(self.shape, other.shape)
+
+        if needs_broadcasting:
+            # Call add_broadcasted_tensor if broadcasting is needed
+            Tensor._C.sub_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.sub_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
+
+            result_tensor_ptr = Tensor._C.sub_broadcasted_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = broadcasted_shape
+            result_data.ndim = len(broadcasted_shape)
+
+            result_data.device = self.device
+            result_data.numel = self.numel  # Update this to calculate the correct number of elements if broadcasting
+
+            result_data.requires_grad = self.requires_grad or other.requires_grad
+            if result_data.requires_grad:
+                result_data.grad_fn = SubBroadcastedBackward(self, other)
         
-        Tensor._C.sub_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
-        Tensor._C.sub_tensor.restype = ctypes.POINTER(CTensor)
+        else:
+            # Call add_tensor if shapes are identical
+            Tensor._C.sub_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.sub_tensor.restype = ctypes.POINTER(CTensor)
 
-        result_tensor_ptr = Tensor._C.sub_tensor(self.tensor, other.tensor)
+            result_tensor_ptr = Tensor._C.sub_tensor(self.tensor, other.tensor)
 
-        result_data = Tensor()
-        result_data.tensor = result_tensor_ptr
-        result_data.shape = self.shape.copy()
-        result_data.ndim = self.ndim
-        result_data.device = self.device
-        result_data.numel = self.numel
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = self.shape.copy()
+            result_data.ndim = self.ndim
 
-        result_data.requires_grad = self.requires_grad or other.requires_grad
-        if result_data.requires_grad:
-            result_data.grad_fn = SubBackward(self, other)
+            result_data.device = self.device
+            result_data.numel = self.numel  # Update this to calculate the correct number of elements if broadcasting
+
+            result_data.requires_grad = self.requires_grad or other.requires_grad
+            if result_data.requires_grad:
+                result_data.grad_fn = SubBackward(self, other)
 
         return result_data
     

tests/test_autograd.py

@@ -73,6 +73,28 @@ 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):
+        """
+        Test autograd for broadcasting subtraction: tensor1 - tensor2
+        """
+        norch_tensor1 = norch.Tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device)  # Shape (1, 2, 3)
+        norch_tensor2 = norch.Tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        norch_result = (norch_tensor1 - norch_tensor2).sum()
+        norch_result.backward()
+        norch_tensor1_grad = utils.to_torch(norch_tensor1.grad)
+        norch_tensor2_grad = utils.to_torch(norch_tensor2.grad)
+
+        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_tensor1 - torch_tensor2).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):
         """
         Test autograd from dividing two tensors: tensor1 / tensor2

tests/test_operations.py

@@ -46,7 +46,7 @@ class TestTensorOperations(unittest.TestCase):
 
         torch_tensor1 = torch.tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device)  # Shape (1, 2, 3)
         torch_tensor2 = torch.tensor([1, 1, 1]).to(self.device)  # Shape (3)
-        torch_expected = torch_tensor1 + torch_tensor2  # Broadcasting addition
+        torch_expected = torch_tensor1 + torch_tensor2
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
@@ -66,6 +66,22 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+    def test_broadcasting_subtraction(self):
+        """
+        Test subtraction of two tensors with broadcasting: tensor1 - tensor2
+        """
+        norch_tensor1 = norch.Tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device)  # Shape (1, 2, 3)
+        norch_tensor2 = norch.Tensor([1, 1, 1]).to(self.device)  # Shape (3)
+        norch_result = norch_tensor1 - norch_tensor2
+        torch_result = utils.to_torch(norch_result)
+
+        torch_tensor1 = torch.tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1, 1, 1]).to(self.device)  # Shape (3)
+        torch_expected = torch_tensor1 - torch_tensor2 
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
+
     def test_division_by_scalar(self):
         """
         Test division of a tensor by a scalar: tensor / scalar
@@ -192,7 +208,7 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
-    def test_sum_axis_no_keepdims(self):
+    def test_sum_axis(self):
         """
         Test summation of a tensor along a specific axis without keeping the dimensions
         """
@@ -203,8 +219,6 @@ class TestTensorOperations(unittest.TestCase):
         torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.sum(torch_tensor, dim=1)
 
-        print(torch_result)
-        print("\n\n", torch_expected)
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
     def test_transpose_T(self):