diff --git a/build/cpu.o b/build/cpu.o index 9e6ddf7..2ee4a6f 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 8c288f1..b40f8c4 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/norch/__pycache__/__init__.cpython-38.pyc b/norch/__pycache__/__init__.cpython-38.pyc index a6e6d05..176d505 100644 Binary files a/norch/__pycache__/__init__.cpython-38.pyc and b/norch/__pycache__/__init__.cpython-38.pyc differ diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index 7c96f58..4e27010 100644 --- a/norch/csrc/cpu.cpp +++ b/norch/csrc/cpu.cpp @@ -148,8 +148,7 @@ void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { - int tensor2_offset = tensor2->shape[1] * tensor2->shape[2]; - int result_data_offset = tensor1->shape[0] * tensor2->shape[2]; + int result_data_stride = tensor1->shape[0] * tensor2->shape[2]; for (int batch = 0; batch < tensor2->shape[0]; batch++) { @@ -157,19 +156,16 @@ void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, flo for (int j = 0; j < tensor2->shape[2]; j++) { float sum = 0.0; for (int k = 0; k < tensor1->shape[1]; k++) { - sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[batch*tensor2_offset + (k * tensor2->shape[2] + j)]; + sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[batch*tensor2->strides[0] + (k * tensor2->shape[2] + j)]; } - result_data[(batch * result_data_offset) + (i * tensor2->shape[2] + j)] = sum; + result_data[(batch * result_data_stride) + (i * tensor2->shape[2] + j)] = sum; } } } } void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { - - int tensor1_offset = tensor1->shape[1] * tensor1->shape[2]; - int tensor2_offset = tensor2->shape[1] * tensor2->shape[2]; - int result_data_offset = tensor1->shape[1] * tensor2->shape[2]; + int result_data_stride = tensor1->shape[1] * tensor2->shape[2]; for (int batch = 0; batch < tensor2->shape[0]; batch++) { @@ -177,9 +173,9 @@ void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_d for (int j = 0; j < tensor2->shape[2]; j++) { float sum = 0.0; for (int k = 0; k < tensor1->shape[2]; k++) { - sum += tensor1->data[(batch * tensor1_offset) + i * tensor1->shape[2] + k] * tensor2->data[batch*tensor2_offset + (k * tensor2->shape[2] + j)]; + sum += tensor1->data[(batch * tensor1->strides[0]) + i * tensor1->shape[2] + k] * tensor2->data[batch*tensor2->strides[0] + (k * tensor2->shape[2] + j)]; } - result_data[(batch * result_data_offset) + (i * tensor2->shape[2] + j)] = sum; + result_data[(batch * result_data_stride) + (i * tensor2->shape[2] + j)] = sum; } } } diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index da241e4..a2bb8e0 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -456,7 +456,7 @@ __global__ void batched_matmul_tensor_cuda_kernel(float* data1, float* data2, fl } __host__ void batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { - + int batch_size = tensor2->shape[0]; int rows1 = tensor1->shape[1]; int cols1 = tensor1->shape[2]; @@ -476,6 +476,42 @@ __host__ void batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float cudaDeviceSynchronize(); } +__global__ void broadcasted_batched_matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int batch_size, int rows1, int cols1, int cols2) { + int batch = blockIdx.z; + int row = blockIdx.y * blockDim.y + threadIdx.y; + int col = blockIdx.x * blockDim.x + threadIdx.x; + + if (row < rows1 && col < cols2) { + float sum = 0.0f; + for (int k = 0; k < cols1; ++k) { + sum += data1[row * cols1 + k] * + data2[batch * cols1 * cols2 + k * cols2 + col]; + } + result_data[batch * rows1 * cols2 + row * cols2 + col] = sum; + } +} + +__host__ void broadcasted_batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int batch_size = tensor2->shape[0]; + int rows1 = tensor1->shape[0]; + int cols1 = tensor1->shape[1]; + int cols2 = tensor2->shape[2]; + + dim3 threadsPerBlock(16, 16); + dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y, batch_size); + broadcasted_batched_matmul_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, batch_size, rows1, cols1, cols2); + + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + __global__ void tensor_pow_scalar_cuda_kernel(float* data, float exponent, 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 7f0f49f..38a6977 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -40,6 +40,9 @@ __global__ void batched_matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int batch_size, int rows1, int cols1, int cols2); __host__ void batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + __global__ void broadcasted_batched_matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int batch_size, int rows1, int cols1, int cols2); + __host__ void broadcasted_batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + __global__ void tensor_pow_scalar_cuda_kernel(float* data, float exponent, float* result_data, int size); __host__ void tensor_pow_scalar_cuda(Tensor* tensor, float exponent, float* result_data); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 6ac29e2..4d4d2fb 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -751,7 +751,7 @@ extern "C" { } Tensor* broadcasted_batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) { - //MxN @ BATCHxNxP = BATCHxMxP + //BATCHxMxP = MxN @ BATCHxNxP // Check if tensors have compatible shapes for matrix multiplication if (tensor1->shape[1] != tensor2->shape[1]) { fprintf(stderr, "Incompatible shapes for broadcasted batched matrix multiplication %dx%d and %dx%dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1], tensor2->shape[2]); @@ -797,7 +797,7 @@ extern "C" { float* result_data; cudaMalloc((void **)&result_data, size * sizeof(float)); - ////broadcasted_batched_matmul_tensor_cuda(tensor1, tensor2, result_data); + broadcasted_batched_matmul_tensor_cuda(tensor1, tensor2, result_data); return create_tensor(result_data, shape, ndim, device); } else { @@ -813,7 +813,7 @@ extern "C" { } Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) { - //BATCHxMxN @ BATCHxNxP = BATCHxMxP + //BATCHxMxP = BATCHxMxN @ BATCHxNxP // Check if tensors have compatible shapes for matrix multiplication if (tensor1->shape[0] != tensor2->shape[0]) { diff --git a/norch/libtensor.so b/norch/libtensor.so index 921564d..d9dcf12 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/tests/test_operations.py b/tests/test_operations.py index ad11717..d443260 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -153,13 +153,13 @@ class TestTensorOperations(unittest.TestCase): """ Test matrix multiplication: tensor1 @ tensor2 """ - norch_tensor1 = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) - norch_tensor2 = norch.Tensor([[[1, 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) + norch_tensor1 = norch.Tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[1., 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) norch_result = norch_tensor1 @ norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) - torch_tensor1 = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) - torch_tensor2 = torch.tensor([[[1, 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) + torch_tensor1 = torch.tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[1., 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) torch_expected = torch_tensor1 @ torch_tensor2 self.assertTrue(utils.compare_torch(torch_result, torch_expected)) @@ -490,15 +490,15 @@ class TestTensorOperations(unittest.TestCase): Test matrix multiplication: MxP = NxM @ MxP """ # Creating batched tensors for Norch - norch_tensor1 = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) - norch_tensor2 = norch.Tensor([[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) + norch_tensor1 = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + norch_tensor2 = norch.Tensor([[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) norch_result = norch_tensor1 @ norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) # Converting to PyTorch tensors for comparison - torch_tensor1 = torch.tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) - torch_tensor2 = torch.tensor([[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) + torch_tensor1 = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_tensor2 = torch.tensor([[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) torch_expected = torch.matmul(torch_tensor1, torch_tensor2) @@ -509,14 +509,14 @@ class TestTensorOperations(unittest.TestCase): """ Test reshaping a tensor followed by matrix multiplication: (tensor.reshape(shape) @ other_tensor) """ - norch_tensor = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + norch_tensor = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) new_shape = [2, 4] norch_reshaped = norch_tensor.reshape(new_shape) norch_result = norch_reshaped @ norch_tensor torch_result = utils.to_torch(norch_result).to(self.device) - torch_tensor = torch.tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_tensor = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) torch_expected = torch_tensor.reshape(new_shape) @ torch_tensor self.assertTrue(utils.compare_torch(torch_result, torch_expected)) @@ -528,15 +528,15 @@ class TestTensorOperations(unittest.TestCase): B = 3 # Batch size # Creating batched tensors for Norch - norch_tensor1 = norch.Tensor([[[1, 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) - norch_tensor2 = norch.Tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + norch_tensor1 = norch.Tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) + norch_tensor2 = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) norch_result = norch_tensor1 @ norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) # Converting to PyTorch tensors for comparison - torch_tensor1 = torch.tensor([[[1, 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) - torch_tensor2 = torch.tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + torch_tensor1 = torch.tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) + torch_tensor2 = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) torch_expected = torch.matmul(torch_tensor1, torch_tensor2) @@ -551,15 +551,15 @@ class TestTensorOperations(unittest.TestCase): B = 3 # Batch size # Creating batched tensors for Norch - norch_tensor1 = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) - norch_tensor2 = norch.Tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + norch_tensor1 = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + norch_tensor2 = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) norch_result = norch_tensor1 @ norch_tensor2 torch_result = utils.to_torch(norch_result).to(self.device) # Converting to PyTorch tensors for comparison - torch_tensor1 = torch.tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device) - torch_tensor2 = torch.tensor([[[2, 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + torch_tensor1 = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_tensor2 = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) torch_expected = torch.matmul(torch_tensor1, torch_tensor2) @@ -572,12 +572,12 @@ class TestTensorOperations(unittest.TestCase): """ Test transposing a tensor followed by matrix multiplication: (tensor.transpose(dim1, dim2) @ other_tensor) """ - norch_tensor = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + norch_tensor = norch.Tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) dim1, dim2 = 0, 2 norch_result = norch_tensor.transpose(dim1, dim2) @ norch_tensor torch_result = utils.to_torch(norch_result).to(self.device) - torch_tensor = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor = torch.tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) torch_expected = torch_tensor.transpose(dim1, dim2) @ torch_tensor self.assertTrue(utils.compare_torch(torch_result, torch_expected)) @@ -587,7 +587,7 @@ class TestTensorOperations(unittest.TestCase): Test a combination of operations: (tensor.sum() + other_tensor) / scalar @ another_tensor followed by reshape """ norch_tensor1 = norch.Tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) - norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[1., 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) scalar = 2 new_shape = [2, 4] norch_result = ((norch_tensor1 + norch_tensor2) / scalar) @ norch_tensor1 @@ -595,7 +595,7 @@ class TestTensorOperations(unittest.TestCase): torch_result = utils.to_torch(norch_result).to(self.device) torch_tensor1 = torch.tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) - torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + torch_tensor2 = torch.tensor([[[1., 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) torch_expected = ((torch_tensor1 + torch_tensor2) / scalar) @ torch_tensor1 torch_expected = torch_expected.reshape(new_shape)