From a159133c8f282d6272853e619dea01c8fd0b70e4 Mon Sep 17 00:00:00 2001
From: hezihang <hezihang@lixiang.com>
Date: Thu, 22 Feb 2024 10:27:33 +0800
Subject: [PATCH 1/2] add naive softmax kernel implement

---
 test/CMakeLists.txt | 14 ++++++++++++
 test/softmax.cu     | 55 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)
 create mode 100644 test/softmax.cu

diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 0a94437..33ec39b 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -94,3 +94,17 @@ target_link_libraries(
     -lcublas
     -lone_tensor
     )
+
+add_executable(
+    softmax
+    "softmax.cu"
+    )
+# set link lib for executed file
+target_link_libraries(
+    softmax
+    # vec_add_shared
+    -lcuda
+    -lcudart
+    -lcublas
+    -lone_tensor
+    )
\ No newline at end of file
diff --git a/test/softmax.cu b/test/softmax.cu
new file mode 100644
index 0000000..0d071df
--- /dev/null
+++ b/test/softmax.cu
@@ -0,0 +1,55 @@
+#include <one_tensor.h>
+#include <cassert>
+
+__global__ void softmax_kernel(float *input, float *output, size_t n)
+{
+    size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid < n)
+    {
+        float sum = 0.0;
+        for (int j = 0; j < n; j++)
+        {
+            sum += expf(input[j]);
+        }
+        output[tid] = expf(input[tid]) / sum;
+    }
+}
+
+int main()
+{
+    Shape problem_size(4, 4);
+
+    OneTensor<float> cuda_input(problem_size);
+    OneTensor<float> cuda_output(problem_size);
+
+    std::vector<float> input(problem_size.nums());
+
+    std::fill_n(input.begin(), problem_size.nums(), 0.1);
+
+    for (int i = 0; i < cuda_input.shape.d0; i++)
+    {
+        for (int j = 0; j < cuda_input.shape.d1; j++)
+        {
+            int index = i * cuda_input.shape.d1 + j;
+            cuda_input.SetHostData(input[i], index);
+        }
+    }
+
+    cuda_output.FillHostData(0.0f);
+
+    cuda_input.sync_device();
+    cuda_output.sync_device();
+
+    cudaStream_t stream;
+    cudaStreamCreate(&stream);
+
+    // softmax core
+    int blockSize = 128;
+    int gridSize = (int)ceil(float(problem_size.nums()) / blockSize);
+
+    softmax_kernel<<<gridSize, blockSize, 0, stream>>>(cuda_input.deviceData<float>(), cuda_output.deviceData<float>(), problem_size.nums());
+
+    cuda_output.sync_device(false);
+    cuda_output.HostDataView();
+    return 0;
+}
\ No newline at end of file

From d160c86dcf10d79c67e776e7fbb53a35def60e4e Mon Sep 17 00:00:00 2001
From: hezihang <hezihang@lixiang.com>
Date: Fri, 23 Feb 2024 14:50:20 +0800
Subject: [PATCH 2/2] add naive safe softmax

---
 test/CMakeLists.txt  | 14 ++++++++++
 test/safe_softmax.cu | 63 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 test/safe_softmax.cu

diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 33ec39b..1aca859 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -107,4 +107,18 @@ target_link_libraries(
     -lcudart
     -lcublas
     -lone_tensor
+    )
+
+add_executable(
+    safe_softmax
+    "safe_softmax.cu"
+    )
+# set link lib for executed file
+target_link_libraries(
+    safe_softmax
+    # vec_add_shared
+    -lcuda
+    -lcudart
+    -lcublas
+    -lone_tensor
     )
\ No newline at end of file
diff --git a/test/safe_softmax.cu b/test/safe_softmax.cu
new file mode 100644
index 0000000..4a931fd
--- /dev/null
+++ b/test/safe_softmax.cu
@@ -0,0 +1,63 @@
+#include <one_tensor.h>
+#include <cassert>
+
+__global__ void safe_softmax_kernel(float *input, float *output, size_t n)
+{
+    size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid < n)
+    {
+        float sum = 0.0;
+        float max_val = -MAXFLOAT;
+        for (int i = 0; i < n; i++)
+        {
+            if (input[i] > max_val)
+            {
+                max_val = input[i];
+            }
+        }
+        for (int i = 0; i < n; i++)
+        {
+            sum += expf(input[i] - max_val);
+        }
+        output[tid] = expf(input[tid] - max_val) / sum;
+    }
+}
+
+int main()
+{
+    Shape problem_size(4, 4);
+
+    OneTensor<float> cuda_input(problem_size);
+    OneTensor<float> cuda_output(problem_size);
+
+    std::vector<float> input(problem_size.nums());
+
+    std::fill_n(input.begin(), problem_size.nums(), 0.1);
+
+    for (int i = 0; i < cuda_input.shape.d0; i++)
+    {
+        for (int j = 0; j < cuda_input.shape.d1; j++)
+        {
+            int index = i * cuda_input.shape.d1 + j;
+            cuda_input.SetHostData(input[i], index);
+        }
+    }
+
+    cuda_output.FillHostData(0.0f);
+
+    cuda_input.sync_device();
+    cuda_output.sync_device();
+
+    cudaStream_t stream;
+    cudaStreamCreate(&stream);
+
+    // softmax core
+    int blockSize = 128;
+    int gridSize = (int)ceil(float(problem_size.nums()) / blockSize);
+
+    safe_softmax_kernel<<<gridSize, blockSize, 0, stream>>>(cuda_input.deviceData<float>(), cuda_output.deviceData<float>(), problem_size.nums());
+
+    cuda_output.sync_device(false);
+    cuda_output.HostDataView();
+    return 0;
+}
\ No newline at end of file