onednn-src 0.1.13

/******************************************************************************
* Copyright 2025 ZTE Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
******************************************************************************/

#include <math.h>
#include <riscv_vector.h>

#include "common/float16.hpp"
#include "common/nstl.hpp"

#include "cpu/rv64/rvv_softmax.hpp"

namespace dnnl {
namespace impl {
namespace cpu {
namespace rv64 {

rvv_softmax_fwd_t::rvv_softmax_fwd_t(const pd_t *apd) : primitive_t(apd) {
    if (pd()->use_jit_) {
        affine_kernel_.reset(new jit_rvv_softmax_affine_kernel_t());
    }
}

namespace {

// f32 compute kernel
void compute_softmax_f32_rvv(const float *src, float *dst, dim_t len,
        bool is_logsoftmax, bool is_softmax_inf_as_zero,
        const jit_rvv_softmax_affine_kernel_t *affine_kernel) {
    float max_val = -INFINITY;
    for (dim_t i = 0; i < len; ++i)
        max_val = src[i] > max_val ? src[i] : max_val;

    if (is_logsoftmax) {
        float sum_exp = 0.f;
        for (dim_t i = 0; i < len; ++i) {
            sum_exp += expf(src[i] - max_val);
        }
        const float log_sum = logf(sum_exp);

        if (affine_kernel) {
            jit_rvv_softmax_affine_kernel_t::call_params_t p;
            p.src = src;
            p.dst = dst;
            p.len = len;
            p.sub = max_val + log_sum;
            p.mul = 1.f;
            (*affine_kernel)(&p);
        } else {
            for (dim_t i = 0; i < len; ++i)
                dst[i] = src[i] - max_val - log_sum;
        }
    } else {
        float sum_exp = 0.f;
        const bool all_minus_inf
                = is_softmax_inf_as_zero && (max_val == -INFINITY);
        for (dim_t i = 0; i < len; ++i) {
            float e = all_minus_inf ? 0.f : expf(src[i] - max_val);
            dst[i] = e;
            sum_exp += e;
        }

        const float inv_sum = sum_exp ? (1.0f / sum_exp) : 1.0f;
        if (affine_kernel) {
            jit_rvv_softmax_affine_kernel_t::call_params_t p;
            p.src = dst;
            p.dst = dst;
            p.len = len;
            p.sub = 0.f;
            p.mul = inv_sum;
            (*affine_kernel)(&p);
        } else {
            for (dim_t i = 0; i < len; ++i)
                dst[i] *= inv_sum;
        }
    }
}

#if defined(DNNL_RISCV_USE_ZVFH_INTRINSICS)
// f16 compute kernel
void compute_softmax_f16_rvv(const dnnl::impl::float16_t *src,
        dnnl::impl::float16_t *dst, dim_t len, bool is_logsoftmax,
        bool is_softmax_inf_as_zero) {

    float max_val
            = (float)nstl::numeric_limits<dnnl::impl::float16_t>::lowest();
    for (dim_t i = 0; i < len; ++i) {
        float val = (float)src[i];
        if (val > max_val) max_val = val;
    }

    if (is_logsoftmax) {
        float sum_exp = 0.f;
        for (dim_t i = 0; i < len; ++i) {
            sum_exp += expf((float)src[i] - max_val);
        }
        const float log_sum = logf(sum_exp);

        for (dim_t i = 0; i < len;) {
            size_t vl = __riscv_vsetvl_e16m1((size_t)(len - i));
            vfloat16m1_t v_src
                    = __riscv_vle16_v_f16m1((const _Float16 *)(src + i), vl);
            vfloat32m2_t v_f32 = __riscv_vfwcvt_f_f_v_f32m2(v_src, vl);
            vfloat32m2_t v_res = __riscv_vfsub_vf_f32m2(v_f32, max_val, vl);
            v_res = __riscv_vfsub_vf_f32m2(v_res, log_sum, vl);
            vfloat16m1_t v_out = __riscv_vfncvt_f_f_w_f16m1(v_res, vl);
            __riscv_vse16_v_f16m1((_Float16 *)(dst + i), v_out, vl);
            i += (dim_t)vl;
        }
    } else {
        float *tmp_dst = new float[len];
        float sum_exp = 0.f;
        const bool all_minus_inf
                = is_softmax_inf_as_zero && (max_val == -INFINITY);
        for (dim_t i = 0; i < len; ++i) {
            float e = all_minus_inf ? 0.f : expf((float)src[i] - max_val);
            tmp_dst[i] = e;
            sum_exp += e;
        }
        const float inv_sum = sum_exp ? (1.0f / sum_exp) : 1.0f;

        for (dim_t i = 0; i < len;) {
            size_t vl = __riscv_vsetvl_e16m1((size_t)(len - i));

            vfloat32m2_t v_f32 = __riscv_vle32_v_f32m2(tmp_dst + i, vl);
            vfloat32m2_t v_res = __riscv_vfmul_vf_f32m2(v_f32, inv_sum, vl);
            vfloat16m1_t v_out = __riscv_vfncvt_f_f_w_f16m1(v_res, vl);
            __riscv_vse16_v_f16m1((_Float16 *)(dst + i), v_out, vl);

            i += (dim_t)vl;
        }
        delete[] tmp_dst;
    }
}
#endif

} // namespace

status_t rvv_softmax_fwd_t::execute_forward(const exec_ctx_t &ctx) const {
    auto src = CTX_IN_MEM(const void *, DNNL_ARG_SRC);
    auto dst = CTX_OUT_MEM(void *, DNNL_ARG_DST);

    const auto &rsp = pd()->rsp_;
    const bool is_softmax_inf_as_zero
            = pd()->alg_kind() == alg_kind::softmax_accurate_inf_as_zero;

    switch (rsp.data_type) {
        case data_type::f32: {
            const float *src_f32 = static_cast<const float *>(src);
            float *dst_f32 = static_cast<float *>(dst);

            const dim_t outer_stride = pd()->axis_size(true) * rsp.inner_size;
            const int nthr = pd()->nthr_;

            if (rsp.inner_size == 1) {
                parallel_nd(rsp.outer_size, [&](dim_t outer) {
                    const dim_t base = outer * outer_stride;
                    compute_softmax_f32_rvv(src_f32 + base, dst_f32 + base,
                            rsp.axis_size, rsp.is_logsoftmax,
                            is_softmax_inf_as_zero, affine_kernel_.get());
                });
            } else {
                auto scratch = ctx.get_scratchpad_grantor().template get<char>(
                        memory_tracking::names::key_softmax_interim_store);

                parallel(nthr, [&](int ithr, int nthr) {
                    float *tmp = reinterpret_cast<float *>(scratch)
                            + static_cast<size_t>(ithr)
                                    * static_cast<size_t>(rsp.axis_size);

                    const dim_t work_amount = rsp.outer_size * rsp.inner_size;
                    dim_t start {0}, end {0};
                    balance211(work_amount, nthr, ithr, start, end);
                    for (dim_t idx = start; idx < end; ++idx) {
                        const dim_t outer = idx / rsp.inner_size;
                        const dim_t i = idx % rsp.inner_size;
                        const dim_t base = outer * outer_stride + i;

                        // gather -> tmp
                        for (dim_t a = 0; a < rsp.axis_size; ++a)
                            tmp[a] = src_f32[base + a * rsp.inner_size];

                        // contiguous kernel (in-place)
                        compute_softmax_f32_rvv(tmp, tmp, rsp.axis_size,
                                rsp.is_logsoftmax, is_softmax_inf_as_zero,
                                affine_kernel_.get());

                        // write back
                        for (dim_t a = 0; a < rsp.axis_size; ++a)
                            dst_f32[base + a * rsp.inner_size] = tmp[a];
                    }
                });
            }
        } break;
#if defined(DNNL_RISCV_USE_ZVFH_INTRINSICS)
        case data_type::f16: {
            const auto *src_f16
                    = static_cast<const dnnl::impl::float16_t *>(src);
            auto *dst_f16 = static_cast<dnnl::impl::float16_t *>(dst);

            const dim_t outer_stride = pd()->axis_size(true) * rsp.inner_size;
            const int nthr = pd()->nthr_;

            if (rsp.inner_size == 1) {
                parallel_nd(rsp.outer_size, [&](dim_t outer) {
                    const dim_t base = outer * outer_stride;
                    compute_softmax_f16_rvv(src_f16 + base, dst_f16 + base,
                            rsp.axis_size, rsp.is_logsoftmax,
                            is_softmax_inf_as_zero);
                });
            } else {
                auto scratch = ctx.get_scratchpad_grantor().template get<char>(
                        memory_tracking::names::key_softmax_interim_store);

                parallel(nthr, [&](int ithr, int nthr) {
                    auto *tmp
                            = reinterpret_cast<dnnl::impl::float16_t *>(scratch)
                            + static_cast<size_t>(ithr)
                                    * static_cast<size_t>(rsp.axis_size);

                    const dim_t work_amount = rsp.outer_size * rsp.inner_size;
                    dim_t start {0}, end {0};
                    balance211(work_amount, nthr, ithr, start, end);

                    size_t stride_bytes = rsp.inner_size * sizeof(_Float16);

                    for (dim_t idx = start; idx < end; ++idx) {
                        const dim_t outer = idx / rsp.inner_size;
                        const dim_t i = idx % rsp.inner_size;
                        const dim_t base = outer * outer_stride + i;

                        for (dim_t a = 0; a < rsp.axis_size;) {
                            size_t vl = __riscv_vsetvl_e16m1(rsp.axis_size - a);
                            vfloat16m1_t v = __riscv_vlse16_v_f16m1(
                                    (const _Float16 *)(src_f16 + base
                                            + a * rsp.inner_size),
                                    stride_bytes, vl);
                            __riscv_vse16_v_f16m1((_Float16 *)(tmp + a), v, vl);
                            a += (dim_t)vl;
                        }

                        compute_softmax_f16_rvv(tmp, tmp, rsp.axis_size,
                                rsp.is_logsoftmax, is_softmax_inf_as_zero);

                        for (dim_t a = 0; a < rsp.axis_size;) {
                            size_t vl = __riscv_vsetvl_e16m1(rsp.axis_size - a);
                            vfloat16m1_t v = __riscv_vle16_v_f16m1(
                                    (const _Float16 *)(tmp + a), vl);
                            __riscv_vsse16_v_f16m1(
                                    (_Float16 *)(dst_f16 + base
                                            + a * rsp.inner_size),
                                    stride_bytes, v, vl);
                            a += (dim_t)vl;
                        }
                    }
                });
            }
        } break;
#endif
        default: return status::unimplemented;
    }

    return status::success;
}

} // namespace rv64
} // namespace cpu
} // namespace impl
} // namespace dnnl