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

min cuda and min cuda axis

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-23 01:39:37 -03:00
parent 9708054ecd
commit fb6429aad6
2 changed files with 42 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

4
norch/csrc/cuda.h
View file

@ -21,6 +21,10 @@
__global__ void max_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size);
__host__ void max_tensor_cuda(Tensor* tensor, float* result_data, int axis);
__global__ void min_tensor_cuda_kernel(float* data, float* result_data);
__global__ void min_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size);
__host__ void min_tensor_cuda(Tensor* tensor, float* result_data, int axis);
__global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
__host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);

55
norch/csrc/tensor.cpp
View file

@ -365,7 +365,6 @@ extern "C" {
int ndim;
int* shape;
if (axis == -1) {
shape = (int*) malloc(sizeof(int));
shape[0] = 1;
ndim = 1;
@ -379,26 +378,20 @@ extern "C" {
ndim = tensor->ndim - 1;
}
int size = 1;
int axis_size = 1;
for (int i = 0; i < ndim; i++) {
size *= shape[i];
axis_size *= shape[i];
}
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void**)&result_data, size * sizeof(float));
//cudaMemset(result_data, INFINITY, size * sizeof(float));
//min_tensor_cuda(tensor, 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);
if (axis == -1) {
cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
} else {
cudaMalloc((void**)&result_data, axis_size * sizeof(float));
}
min_tensor_cpu(tensor, result_data, size, shape, axis);
min_tensor_cuda(tensor, result_data, axis);
if (keepdim) {
if (axis == -1){
@ -414,11 +407,39 @@ extern "C" {
}
shape[axis] = 1;
ndim = tensor->ndim;
}
}
}
return create_tensor(result_data, shape, ndim, device);
}
}
else {
float* result_data = (float*)malloc(axis_size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
min_tensor_cpu(tensor, result_data, axis_size, shape, axis);
if (keepdim) {
if (axis == -1){
ndim = tensor->ndim;
shape = (int*) malloc((tensor->ndim) * sizeof(int));
for (int i = 0; i < tensor->ndim; i++) {
shape[i] = 1;
}
} else {
shape = (int*) malloc((tensor->ndim) * sizeof(int));
for (int i = 0; i < tensor->ndim; i++) {
shape[i] = tensor->shape[i];
}
shape[axis] = 1;
ndim = tensor->ndim;
}
}
return create_tensor(result_data, shape, ndim, device);
}
}
Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2) {