diff --git a/build/cpu.o b/build/cpu.o index e346900..d681b76 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 004bd90..069d281 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/build/libtensor.so b/build/libtensor.so index 9c8edfd..4118958 100755 Binary files a/build/libtensor.so and b/build/libtensor.so differ diff --git a/build/tensor.o b/build/tensor.o index 9a67429..3dbd4ca 100644 Binary files a/build/tensor.o and b/build/tensor.o differ diff --git a/norch/__pycache__/__init__.cpython-39.pyc b/norch/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..52180c0 Binary files /dev/null and b/norch/__pycache__/__init__.cpython-39.pyc differ diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 683c329..f6d3f53 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/__pycache__/tensor.cpython-39.pyc b/norch/__pycache__/tensor.cpython-39.pyc index 16c0798..02bd848 100644 Binary files a/norch/__pycache__/tensor.cpython-39.pyc and b/norch/__pycache__/tensor.cpython-39.pyc differ diff --git a/norch/autograd/__pycache__/functions.cpython-39.pyc b/norch/autograd/__pycache__/functions.cpython-39.pyc new file mode 100644 index 0000000..7e8e38c Binary files /dev/null and b/norch/autograd/__pycache__/functions.cpython-39.pyc differ diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index 3bea802..0365de5 100644 --- a/norch/csrc/cpu.cpp +++ b/norch/csrc/cpu.cpp @@ -45,6 +45,26 @@ void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { } } + +void 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]; + + for (int batch = 0; batch < tensor2->shape[0]; batch++) { + + for (int i = 0; i < tensor1->shape[0]; i++) { + 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)]; + } + result_data[(batch * result_data_offset) + (i * tensor2->shape[2] + j)] = sum; + } + } + } +} + void pow_tensor_cpu(Tensor* tensor, float power, float* result_data) { for (int i = 0; i < tensor->size; i++) { diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h index f50b0d9..b6275b3 100644 --- a/norch/csrc/cpu.h +++ b/norch/csrc/cpu.h @@ -8,6 +8,7 @@ void sum_tensor_cpu(Tensor* tensor1, float* result_data); void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); +void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); void pow_tensor_cpu(Tensor* tensor, float power, float* result_data); void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data); void ones_like_tensor_cpu(Tensor* tensor, float* result_data); diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index 071828f..afa9820 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -5,7 +5,6 @@ #define THREADS_PER_BLOCK 128 #define TILE_SIZE 32 -#define SHMEM_SIZE THREADS_PER_BLOCK * sizeof(float) __host__ void cpu_to_cuda(Tensor* tensor) { @@ -61,7 +60,7 @@ __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size) { - __shared__ float partial_sum[SHMEM_SIZE]; + __shared__ float partial_sum[THREADS_PER_BLOCK * sizeof(float)]; int tid = threadIdx.x; int i = blockIdx.x * blockDim.x + threadIdx.x; diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index c361161..ed4eeb7 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -327,9 +327,10 @@ extern "C" { } Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2) { + //MxN @ NxP = MxP // Check if tensors have compatible shapes for matrix multiplication if (tensor1->shape[1] != tensor2->shape[0]) { - fprintf(stderr, "Incompatible shapes for matrix multiplication\n"); + fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]); exit(1); } @@ -387,6 +388,69 @@ extern "C" { } } + Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) { + //MxN @ BATCHxNxP = BATCHxMxP + // Check if tensors have compatible shapes for matrix multiplication + if (tensor1->shape[1] != tensor2->shape[1]) { + fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + int ndim = 3; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + shape[0] = tensor2->shape[0];; + shape[1] = tensor1->shape[0]; + shape[2] = tensor2->shape[2]; + + int size = 1; + for (int i = 0; i < ndim; i++) { + size *= shape[i]; + } + + float* result_data = (float*)malloc(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, size * sizeof(float)); + matmul_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + batched_matmul_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + + } + + Tensor* pow_tensor(Tensor* tensor, float power) { char* device = (char*)malloc(strlen(tensor->device) + 1); if (device != NULL) { diff --git a/norch/tensor.py b/norch/tensor.py index 94568e1..32a5bd4 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -291,22 +291,38 @@ class Tensor: return self def __matmul__(self, other): - if self.ndim != 2 or other.ndim != 2: - raise ValueError("Matrix multiplication requires 2D tensors") + if other.ndim == 3: + #batched 3D matmul - if self.shape[1] != other.shape[0]: - raise ValueError("Incompatible shapes for matrix multiplication") + Tensor._C.batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.batched_matmul_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.batched_matmul_tensor(self.tensor, other.tensor) - Tensor._C.matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] - Tensor._C.matmul_tensor.restype = ctypes.POINTER(CTensor) + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = [other.shape[0], self.shape[0], other.shape[2]] + result_data.ndim = 3 + result_data.device = self.device + + else: + #2D matmul + if self.ndim != 2 or other.ndim != 2: + raise ValueError("Matrix multiplication requires 2D tensors") - result_tensor_ptr = Tensor._C.matmul_tensor(self.tensor, other.tensor) + if self.shape[1] != other.shape[0]: + raise ValueError("Incompatible shapes for matrix multiplication") - result_data = Tensor() - result_data.tensor = result_tensor_ptr - result_data.shape = [self.shape[0], other.shape[1]] - result_data.ndim = 2 - result_data.device = self.device + Tensor._C.matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.matmul_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.matmul_tensor(self.tensor, other.tensor) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = [self.shape[0], other.shape[1]] + result_data.ndim = 2 + result_data.device = self.device result_data.requires_grad = self.requires_grad or other.requires_grad if result_data.requires_grad: diff --git a/test.py b/test.py index 7efccf6..a65593d 100644 --- a/test.py +++ b/test.py @@ -20,8 +20,25 @@ if __name__ == "__main__": import random import numpy as np + a = norch.Tensor([ + [[1.234, 2.123], [3.635, 4.456], [5.678, 6.789]], + [[7.890, 8.901], [9.012, 1.234], [2.345, 3.456]], + [[4.567, 5.678], [6.789, 7.890], [8.901, 9.012]], + [[1.234, 2.345], [3.456, 4.567], [5.678, 6.789]], + [[7.890, 8.901], [9.012, 1.234], [2.345, 3.456]] + ]) + + b = norch.Tensor([ + [1.234, 2.123, 1.5], + [5.678, 6.789, 1.293], + [3.635, 4.456, 1.0202], + [7.890, 8.901, 1.91], + ]) + + result = b @ a + print(result) - #a = Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda") + #a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda") #b = Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda") #c = Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda") @@ -37,7 +54,6 @@ if __name__ == "__main__": print(a.grad)""" """#print(a) - """ N = 10 a = norch.Tensor([[1 for _ in range(N)] for _ in range(N)]) #b = norch.Tensor([[random.uniform(0, 1) for _ in range(N)] for _ in range(N)]) @@ -56,7 +72,7 @@ if __name__ == "__main__": print("\n\n") """ - + """ a = [[random.uniform(0, 1) for _ in range(N)] for _ in range(N)] b = [[random.uniform(0, 1) for _ in range(N)] for _ in range(N)] ini = time.time()