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

sum axis cpu and broadcast forward

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-16 16:35:55 -03:00
parent 64953db11d
commit 3c8da6a245
14 changed files with 266 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

66
norch/csrc/tensor.cpp
View file

@ -124,7 +124,44 @@ extern "C" {
}
}
Tensor* sum_tensor(Tensor* tensor) {
Tensor* add_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;
}
// Allocate memory for result tensor
float* result_data = (float*)malloc(tensor1->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
Tensor* sum_tensor(Tensor* tensor, int axis) {
char* device = (char*)malloc(strlen(tensor->device) + 1);
if (device != NULL) {
@ -133,14 +170,25 @@ extern "C" {
fprintf(stderr, "Memory allocation failed\n");
exit(-1);
}
int ndim = 1;
int* shape = (int*)malloc(ndim * sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
int ndim;
int* shape;
if (axis == -1) {
shape = (int*) malloc(sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
shape[0] = 1;
ndim = 1;
} else {
shape = (int*) malloc((tensor->ndim - 1) * sizeof(int));
for (int i = 0, j = 0; i < tensor->ndim; ++i) {
if (i != axis) {
shape[j++] = tensor->shape[i];
}
}
ndim = tensor->ndim - 1;
}
shape[0] = 1;
if (strcmp(tensor->device, "cuda") == 0) {
@ -155,7 +203,7 @@ extern "C" {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
sum_tensor_cpu(tensor, result_data);
sum_tensor_cpu(tensor, result_data, axis);
return create_tensor(result_data, shape, ndim, device);
}
}