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

broadcasted batched matmul cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-22 19:20:54 -03:00
parent b72c16f110
commit 65dc86aaa4
9 changed files with 71 additions and 36 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
norch/__pycache__/__init__.cpython-38.pyc
View file

Binary file not shown.

16
norch/csrc/cpu.cpp
View file

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

38
norch/csrc/cuda.cu
View file

@ -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<<<number_of_blocks, threadsPerBlock>>>(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;

3
norch/csrc/cuda.h
View file

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

6
norch/csrc/tensor.cpp
View file

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

BIN
norch/libtensor.so
View file

Binary file not shown.

44
tests/test_operations.py
View file

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