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

---
 norch/csrc/cuda.cu | 131 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 127 insertions(+), 4 deletions(-)

diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu
index fdaaba5..7c193fa 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -227,26 +227,26 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
 }
 
 __global__ void max_tensor_cuda_kernel(float* data, float* result_data, int size) {
-    __shared__ float partial_sum[THREADS_PER_BLOCK * sizeof(float)];
+    __shared__ float partial_max[THREADS_PER_BLOCK * sizeof(float)];
 
     int tid = threadIdx.x;
     int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-    partial_sum[tid] = (i < size) ? data[i] : 0;
+    partial_max[tid] = (i < size) ? data[i] : -FLT_MAX;
 
     __syncthreads();
 
     // Perform block-wise reduction
     for (int s = blockDim.x / 2; s > 0; s >>= 1) {
         if (tid < s) {
-            partial_sum[tid] = fmax(partial_sum[tid], partial_sum[tid + s]);
+            partial_max[tid] = fmax(partial_max[tid], partial_max[tid + s]);
         }
         __syncthreads();
     }
 
     // Write block sum to global memory
     if (tid == 0) {
-        result_data[blockIdx.x] = partial_sum[0];
+        result_data[blockIdx.x] = partial_max[0];
     }
 }
 
@@ -347,6 +347,129 @@ __host__ void max_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
     }
 }
 
+__global__ void min_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    __shared__ float partial_min[THREADS_PER_BLOCK * sizeof(float)];
+
+    int tid = threadIdx.x;
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+
+    partial_min[tid] = (i < size) ? data[i] : FLT_MAX;
+
+    __syncthreads();
+
+    // Perform block-wise reduction
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            partial_min[tid] = fmin(partial_min[tid], partial_min[tid + s]);
+        }
+        __syncthreads();
+    }
+
+    // Write block sum to global memory
+    if (tid == 0) {
+        result_data[blockIdx.x] = partial_min[0];
+    }
+}
+
+__device__ float atomicMinFloat(float* address, float val) {
+    int* address_as_int = (int*)address;
+    int old = *address_as_int, assumed;
+
+    do {
+        assumed = old;
+        old = atomicCAS(address_as_int, assumed,
+                        __float_as_int(fminf(val, __int_as_float(assumed))));
+    } while (assumed != old);
+
+    return __int_as_float(old);
+}
+
+__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) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < result_size) {
+        for (int i = 0; i < shape[axis]; i++) {
+            int index = 0;
+            int remainder = tid;
+            for (int k = ndim - 2; k >= 0; k--) {
+                index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1];
+                remainder /= shape[k < axis ? k : k + 1];
+            }
+            index += i * axis_stride;
+
+            atomicMinFloat(&result_data[tid], data[index]);
+        }
+    }
+}
+
+__host__ void min_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
+
+    if (axis == -1) {
+        cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
+        
+        int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+
+        // First-level reduction
+        min_tensor_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
+
+        // If necessary, perform multiple levels of reduction
+        while (num_blocks > 1) {
+            int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+            min_tensor_cuda_kernel<<<num_blocks_next, THREADS_PER_BLOCK>>>(result_data, result_data, num_blocks);
+            num_blocks = num_blocks_next;
+        }
+
+        cudaError_t error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("CUDA error: %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+        cudaDeviceSynchronize();
+        
+    } else {
+        int axis_stride = tensor->strides[axis];
+
+        // Calculate the size of the resulting tensor
+        int result_size = 1;
+        for (int i = 0; i < tensor->ndim; i++) {
+            if (i != axis) {
+                result_size *= tensor->shape[i];
+            }
+        }
+
+        // Allocate memory for strides and shape on the device
+        int* d_strides;
+        int* d_shape;
+        cudaMalloc(&d_strides, tensor->ndim * sizeof(int));
+        cudaMalloc(&d_shape, tensor->ndim * sizeof(int));
+        cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+        cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+        // Initialize result_data to 0
+        float inf = FLT_MAX;
+        cudaMemset(result_data, *reinterpret_cast<int*>(&inf), result_size * sizeof(float));
+
+        int num_threads = result_size;
+        int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+        min_tensor_cuda_kernel_axis<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size);
+
+        cudaError_t error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("CUDA error: %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+        cudaDeviceSynchronize();
+
+        // Free allocated memory
+        cudaFree(d_strides);
+        cudaFree(d_shape);
+    }
+}
+
+
+
 __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
    
     int i = blockIdx.x * blockDim.x + threadIdx.x;