Issue/259 add softmax operator

include/infiniop.h

@@ -14,6 +14,7 @@
 #include "infiniop/ops/relu.h"
 #include "infiniop/ops/rms_norm.h"
 #include "infiniop/ops/rope.h"
+#include "infiniop/ops/softmax.h"
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/tensor_descriptor.h"

include/infiniop/ops/softmax.h

@@ -0,0 +1,20 @@
+#ifndef __INFINIOP_SOFTMAX_API_H__
+#define __INFINIOP_SOFTMAX_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopSoftmaxDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateSoftmaxDescriptor(infiniopHandle_t handle,
+                                                            infiniopSoftmaxDescriptor_t *desc_ptr,
+                                                            infiniopTensorDescriptor_t y_desc,
+                                                            infiniopTensorDescriptor_t x_desc,
+                                                            int axis);
+
+__C infiniStatus_t infiniopGetSoftmaxWorkspaceSize(infiniopSoftmaxDescriptor_t desc, size_t *size);
+
+__C infiniStatus_t infiniopSoftmax(infiniopSoftmaxDescriptor_t desc, void *workspace, size_t workspace_size, void *y, const void *x, void *stream);
+
+__C infiniStatus_t infiniopDestroySoftmaxDescriptor(infiniopSoftmaxDescriptor_t desc);
+
+#endif

src/infiniop/ops/softmax/cpu/softmax_cpu.cc

@@ -0,0 +1,94 @@
+#include "softmax_cpu.h"
+#include "../../../devices/cpu/common_cpu.h"
+
+namespace op::softmax::cpu {
+
+Descriptor::~Descriptor() = default;
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t y,
+    infiniopTensorDescriptor_t x,
+    int axis) {
+
+    auto handle = reinterpret_cast<device::cpu::Handle *>(handle_);
+    auto dtype = y->dtype();
+
+    const auto &x_shape = x->shape();
+    const auto &y_shape = y->shape();
+
+    CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16);
+
+    CHECK_SAME_SHAPE(y_shape, x_shape);
+
+    auto result = SoftmaxInfo::create(y, x, axis);
+    CHECK_RESULT(result);
+
+    *desc_ptr = new Descriptor(
+        dtype,
+        result.take(),
+        0,
+        nullptr,
+        handle->device, handle->device_id);
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+template <typename T>
+void softmax_cpu(const SoftmaxInfo &info,
+                 const void *x, void *y, int axis) {
+    int dim_size = info.dim_size;
+    int stride = info.stride;
+    int other_size = info.other_size;
+    auto input = reinterpret_cast<const T *>(x);
+    auto output = reinterpret_cast<T *>(y);
+
+    auto compute_softmax = [&](int i) {
+        int tid = i % stride + (i - i % stride) * dim_size;
+
+        float max_data = -INFINITY;
+        for (int j = 0; j < dim_size; j++) {
+            int index = tid + j * stride;
+            max_data = fmax(max_data, utils::cast<float>(input[index]));
+        }
+
+        float sum_data = 0.0f;
+        for (int j = 0; j < dim_size; j++) {
+            int index = tid + j * stride;
+            sum_data += std::exp(utils::cast<float>(input[index]) - max_data);
+        }
+
+        for (int j = 0; j < dim_size; j++) {
+            int index = tid + j * stride;
+            float result = std::exp(utils::cast<float>(input[index]) - max_data) / sum_data;
+            output[index] = utils::cast<T>(result);
+        }
+    };
+#pragma omp parallel for
+    for (int i = 0; i < other_size; i++) {
+        compute_softmax(i);
+    }
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *workspace,
+    size_t workspace_size,
+    void *y,
+    const void *x,
+    void *stream_) const {
+    switch (_dtype) {
+    case INFINI_DTYPE_F16:
+        softmax_cpu<fp16_t>(_info, x, y, _info.axis);
+        return INFINI_STATUS_SUCCESS;
+    case INFINI_DTYPE_F32:
+        softmax_cpu<float>(_info, x, y, _info.axis);
+        return INFINI_STATUS_SUCCESS;
+    case INFINI_DTYPE_BF16:
+        softmax_cpu<bf16_t>(_info, x, y, _info.axis);
+        return INFINI_STATUS_SUCCESS;
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+}
+} // namespace op::softmax::cpu

src/infiniop/ops/softmax/cpu/softmax_cpu.h

@@ -0,0 +1,8 @@
+#ifndef __SOFTMAX_CPU_H__
+#define __SOFTMAX_CPU_H__
+
+#include "../softmax.h"
+
+DESCRIPTOR(cpu)
+
+#endif // __SOFTMAX_CPU_H__

src/infiniop/ops/softmax/cuda/softmax_cuda.cu

@@ -0,0 +1,54 @@
+#include "../../../devices/cuda/cuda_common.cuh"
+#include "softmax_cuda.cuh"
+#include "softmax_kernel.cuh"
+
+namespace op::softmax::cuda {
+
+struct Descriptor::Opaque {
+    std::shared_ptr<device::cuda::Handle::Internal> internal;
+};
+
+Descriptor::~Descriptor() {
+    delete _opaque;
+}
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t y,
+    infiniopTensorDescriptor_t x,
+    int axis) {
+    auto dtype = y->dtype();
+    auto handle = reinterpret_cast<device::cuda::Handle *>(handle_);
+    auto result = SoftmaxInfo::create(y, x, axis);
+    CHECK_RESULT(result);
+    CHECK_SAME_SHAPE(y->shape(), x->shape());
+    CHECK_DTYPE(y->dtype(), x->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16);
+    *desc_ptr = new Descriptor(
+        dtype,
+        result.take(),
+        0,
+        new Opaque{handle->internal()},
+        handle->device, handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *workspace,
+    size_t workspace_size,
+    void *y,
+    const void *x,
+    void *stream_) const {
+
+    switch (_dtype) {
+    case INFINI_DTYPE_F16:
+        return softmax_dispatch<half>(_info, y, x, stream_);
+    case INFINI_DTYPE_F32:
+        return softmax_dispatch<float>(_info, y, x, stream_);
+    case INFINI_DTYPE_BF16:
+        return softmax_dispatch<__nv_bfloat16>(_info, y, x, stream_);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+}
+} // namespace op::softmax::cuda

src/infiniop/ops/softmax/cuda/softmax_cuda.cuh

@@ -0,0 +1,8 @@
+#ifndef __SOFTMAX_CUDA_CUH__
+#define __SOFTMAX_CUDA_CUH__
+
+#include "../softmax.h"
+
+DESCRIPTOR(cuda)
+
+#endif // __SOFTMAX_CUDA_CUH__