diff --git a/build/cpu.o b/build/cpu.o index 3299b62..6429509 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 cb40e39..60b85e1 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 53ef710..bcc09bc 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 4d65981..1705e20 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 6ad4633..694d6fd 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 e5b3c20..daa0668 100644 --- a/norch/autograd/functions.py +++ b/norch/autograd/functions.py @@ -216,6 +216,7 @@ class MaxBackward: grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0] else: + if not self.keepdim: # Remove dimensions of size 1 from the gradient tensor. input_shape = [s for i, s in enumerate(input_shape) if i != self.axis] @@ -225,8 +226,12 @@ class MaxBackward: grad_output_shape.insert(self.axis, 1) grad_output = gradient.reshape(grad_output_shape) grad_output = grad_output + self.input[0].zeros_like() - + + print(self.input[0]) max_value = self.input[0].max() + print(max_value) + max_values = self.input[0].max(axis=self.axis, keepdim=True) + print('\n\n', max_values, '@@@') mask = self.input[0].equal(max_value) grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0] diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index 6d65273..9584c23 100644 --- a/norch/csrc/cpu.cpp +++ b/norch/csrc/cpu.cpp @@ -306,6 +306,45 @@ void equal_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { } } +void equal_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 + 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; + } + + // Perform element-wise equal with broadcasting + 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--) { + 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]) ? 1.0f : 0.0f; + } + + // Free strides + free(strides1); + free(strides2); +} + void ones_like_tensor_cpu(Tensor* tensor, float* result_data) { diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h index b006ca7..83bd83f 100644 --- a/norch/csrc/cpu.h +++ b/norch/csrc/cpu.h @@ -22,6 +22,7 @@ void tensor_pow_scalar_cpu(Tensor* tensor, float exponent, float* result_data); void log_tensor_cpu(Tensor* tensor, float* result_data); void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data); void equal_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); +void equal_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size); void ones_like_tensor_cpu(Tensor* tensor, float* result_data); void zeros_like_tensor_cpu(Tensor* tensor, float* result_data); void transpose_1D_tensor_cpu(Tensor* tensor, float* result_data); diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index 5726c69..7d056ab 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -564,6 +564,67 @@ __host__ void equal_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_ cudaDeviceSynchronize(); } +__global__ void equal_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]) ? 1.0f : 0.0f; +} + +__host__ void equal_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; + equal_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 ones_like_tensor_cuda_kernel(float* data, float* result_data, int size) { diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h index 712e600..73eb831 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -52,6 +52,9 @@ __global__ void equal_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size); __host__ void equal_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + __global__ void equal_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 equal_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size); + __global__ void ones_like_tensor_cuda_kernel(float* data, float* result_data, int size); __host__ void ones_like_tensor_cuda(Tensor* tensor, float* result_data); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 8564b2a..570aefa 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -1084,6 +1084,54 @@ extern "C" { } } + Tensor* equal_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; + } + + int broadcasted_size = 1; + for (int i = 0; i < max_ndim; i++) { + broadcasted_size *= broadcasted_shape[i]; + } + + if (strcmp(tensor1->device, "cuda") == 0) { + float* result_data; + cudaMalloc((void **)&result_data, broadcasted_size * sizeof(float)); + equal_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); + } + + equal_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size); + return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device); + } + } + Tensor* ones_like_tensor(Tensor* tensor) { char* device = (char*)malloc(strlen(tensor->device) + 1); diff --git a/norch/csrc/tensor.h b/norch/csrc/tensor.h index bdcf25d..5669d4f 100644 --- a/norch/csrc/tensor.h +++ b/norch/csrc/tensor.h @@ -31,6 +31,7 @@ extern "C" { Tensor* scalar_pow_tensor(float base, Tensor* tensor); Tensor* log_tensor(Tensor* tensor); Tensor* equal_tensor(Tensor* tensor1, Tensor* tensor2); + Tensor* equal_broadcasted_tensor(Tensor* tensor1, Tensor* tensor2); void to_device(Tensor* tensor, char* device); Tensor* ones_like_tensor(Tensor* tensor); Tensor* zeros_like_tensor(Tensor* tensor); diff --git a/norch/libtensor.so b/norch/libtensor.so index 90f2346..d1b1cd0 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/tensor.py b/norch/tensor.py index e8963e8..22cc111 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -656,25 +656,59 @@ class Tensor: return self.equal(other) else: return False - - if self.shape != other.shape: - return False - Tensor._C.equal_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] - Tensor._C.equal_tensor.restype = ctypes.POINTER(CTensor) + broadcasted_shape_add = [] - result_tensor_ptr = Tensor._C.equal_tensor(self.tensor, other.tensor) + # 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 - 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 + + 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_add.append(max(dim1, dim2)) + return broadcasted_shape_add, True + + broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape) + if needs_broadcasting: + # Call equal_broadcasted_tensor if broadcasting is needed + Tensor._C.equal_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.equal_broadcasted_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.equal_broadcasted_tensor(self.tensor, other.tensor) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = broadcasted_shape_add.copy() + result_data.ndim = len(broadcasted_shape_add) + + result_data.device = self.device + result_data.numel = 1 + for s in result_data.shape: + result_data.numel *= s + + else: + Tensor._C.equal_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.equal_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.equal_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 + return result_data - def log(self): Tensor._C.log_tensor.argtypes = [ctypes.POINTER(CTensor)] Tensor._C.log_tensor.restype = ctypes.POINTER(CTensor) diff --git a/tests/test_operations.py b/tests/test_operations.py index 008136d..6a12367 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -546,6 +546,21 @@ class TestTensorOperations(unittest.TestCase): self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_broadcasted_equal(self): + """ + Test broadcasted equal two tensors: tensor1.equal(tensor2) + """ + norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[10, 10]], [[5, 6]]]).to(self.device) + norch_result = norch_tensor1.equal(norch_tensor2) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[10, 10]], [[5, 6]]]).to(self.device) + torch_expected = (torch_tensor1 == torch_tensor2).float() + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_zeros_like(self):