Function do_softmax

dnn/src/fallback/softmax/opr_impl.cpp:10–142

namespace fallback {

static void do_softmax(
        const float* sptr, float* dptr, size_t A, size_t B, size_t C,
        _megdnn_workspace workspace) {
    constexpr auto float_min = std::numeric_limits<float>::min();
    constexpr auto step = GI_SIMD_LEN_BYTE / sizeof(float);
    // TODO: When C=2,3,4..., src_ptr span is relatively large, the performance may
    // be poor

    if (C != 1) {
        WorkspaceBundle workspace_bundle{
                workspace.raw_ptr, {A * C * sizeof(float), A * C * sizeof(float)}};
        float* max = workspace_bundle.get_workspace(0).raw_ptr->as<float>();
        GI_FLOAT32_t v_max = GiBroadcastFloat32(float_min);
        size_t i = 0;
        for (; i + step <= A * C; i += step)
            GiStoreFloat32(max + i, v_max);
        for (; i < A * C; i++)
            max[i] = float_min;

        for (size_t a = 0; a < A; a++) {
            for (size_t b = 0; b < B; b++) {
                auto max_ptr = max + a * C;
                auto limit = max_ptr + C;
                auto src_ptr = sptr + a * B * C + b * C;

                for (; max_ptr + step <= limit; max_ptr += step, src_ptr += step) {
                    GI_FLOAT32_t v_p = GiLoadFloat32(src_ptr);
                    GI_FLOAT32_t v_max = GiLoadFloat32(max_ptr);
                    v_max = GiMaximumFloat32(v_max, v_p);
                    GiStoreFloat32(max_ptr, v_max);
                }
                for (; max_ptr < limit; ++max_ptr, ++src_ptr) {
                    *max_ptr = std::max(*src_ptr, *max_ptr);
                }
            }
        }

        float* sum = workspace_bundle.get_workspace(1).raw_ptr->as<float>();
        memset(sum, 0, A * C * sizeof(float));
        for (size_t a = 0; a < A; a++) {
            for (size_t b = 0; b < B; b++) {
                auto max_ptr = max + a * C;
                auto limit = max_ptr + C;
                auto sum_ptr = sum + a * C;
                auto src_ptr = sptr + a * B * C + C * b;
                auto dst_ptr = dptr + a * B * C + C * b;
                for (; max_ptr + step <= limit; max_ptr += step, sum_ptr += step,
                                                src_ptr += step, dst_ptr += step) {
                    GI_FLOAT32_t v_p = GiLoadFloat32(src_ptr);
                    GI_FLOAT32_t v_max = GiLoadFloat32(max_ptr);
                    GI_FLOAT32_t v_sum = GiLoadFloat32(sum_ptr);
                    v_p = GiExpPsFloat32(GiSubtractFloat32(v_p, v_max));
                    v_sum = GiAddFloat32(v_p, v_sum);
                    GiStoreFloat32(dst_ptr, v_p);
                    GiStoreFloat32(sum_ptr, v_sum);
                }
                for (; max_ptr < limit; ++max_ptr, ++sum_ptr, ++src_ptr, ++dst_ptr) {
                    *dst_ptr = exp(*src_ptr - *max_ptr);