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

batched and broadcasted matmul

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-02 20:21:19 -03:00
parent 31be051028
commit b7cf03828e
11 changed files with 136 additions and 12 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/libtensor.so
View file

Binary file not shown.

BIN
build/tensor.o
View file

Binary file not shown.

BIN
norch/__pycache__/tensor.cpython-38.pyc
View file

Binary file not shown.

BIN
norch/autograd/__pycache__/functions.cpython-38.pyc
View file

Binary file not shown.

2
norch/autograd/functions.py
View file

@ -35,8 +35,6 @@ class MatmulBackward:
x, y = self.input
return [gradient @ y.T, x.T @ gradient]
class PowBackward:
def __init__(self, x, power):
self.input = [x]

22
norch/csrc/cpu.cpp
View file

@ -46,7 +46,7 @@ void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
}
void batched_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];
@ -65,6 +65,26 @@ void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_d
}
}
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];
for (int batch = 0; batch < tensor2->shape[0]; batch++) {
for (int i = 0; i < tensor1->shape[1]; i++) {
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)];
}
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++) {

1
norch/csrc/cpu.h
View file

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

74
norch/csrc/tensor.cpp
View file

@ -388,11 +388,11 @@ extern "C" {
}
}
Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
Tensor* broadcasted_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]);
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]);
exit(1);
}
@ -435,7 +435,75 @@ extern "C" {
float* result_data;
cudaMalloc((void **)&result_data, size * sizeof(float));
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 {
float* result_data = (float*)malloc(size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
broadcasted_batched_matmul_tensor_cpu(tensor1, tensor2, result_data);
return create_tensor(result_data, shape, ndim, device);
}
}
Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
//BATCHxMxN @ BATCHxNxP = BATCHxMxP
// Check if tensors have compatible shapes for matrix multiplication
if (tensor1->shape[0] != tensor2->shape[0]) {
fprintf(stderr, "Tensors must have same batch dimension for batch matmul %d and %d\n", tensor1->shape[0], tensor2->shape[0]);
exit(1);
}
if (tensor1->shape[2] != 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[1];
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));
//batched_matmul_tensor_cuda(tensor1, tensor2, result_data);
return create_tensor(result_data, shape, ndim, device);
}
else {

20
norch/tensor.py
View file

@ -291,8 +291,22 @@ class Tensor:
return self
def __matmul__(self, other):
if other.ndim == 3:
#batched 3D matmul
if self.ndim < 3 and other.ndim == 3:
#broadcasted 2D x 3D matmul
Tensor._C.broadcasted_batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.broadcasted_batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.broadcasted_batched_matmul_tensor(self.tensor, other.tensor)
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
elif self.ndim == 3 and other.ndim == 3:
#broadcasted 3D x 3D matmul
Tensor._C.batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
@ -301,7 +315,7 @@ class Tensor:
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = [other.shape[0], self.shape[0], other.shape[2]]
result_data.shape = [other.shape[0], self.shape[1], other.shape[2]]
result_data.ndim = 3
result_data.device = self.device

29
test.py
View file

@ -26,17 +26,40 @@ if __name__ == "__main__":
[[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]]
])
], requires_grad=True)
b = norch.Tensor([
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],
])
],[
[1.234, 2.123, 1.5],
[5.678, 6.789, 1.293],
[3.635, 4.456, 1.0202],
[7.890, 8.901, 1.91],
],[
[1.234, 2.123, 1.5],
[5.678, 6.789, 1.293],
[3.635, 4.456, 1.0202],
[7.890, 8.901, 1.91],
],[
[1.234, 2.123, 1.5],
[5.678, 6.789, 1.293],
[3.635, 4.456, 1.0202],
[7.890, 8.901, 1.91],
],[
[1.234, 2.123, 1.5],
[5.678, 6.789, 1.293],
[3.635, 4.456, 1.0202],
[7.890, 8.901, 5.91],
]])
result = b @ a
print(result)
#c = result.sum()
#c.backward()
#print(a.grad)
#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")