From fb6429aad6eb11ecebda59dd5a4a2cdb4c2635a0 Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Thu, 23 May 2024 01:39:37 -0300
Subject: [PATCH] min cuda and min cuda axis

---
 norch/csrc/cuda.h     |  4 ++++
 norch/csrc/tensor.cpp | 55 ++++++++++++++++++++++++++++++-------------
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h
index b0ccbf2..766ce67 100644
--- a/norch/csrc/cuda.h
+++ b/norch/csrc/cuda.h
@@ -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);
     
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index d94a3da..15a66bb 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -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) {