/* Copyright 2020 Canaan Inc.
 *
 * 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.
 */
#pragma once
#ifndef ENABLE_K510_RUNTIME
#include "nncase/ir/ops/k510/gnne_meshnet.h"
#include <nncase/ir/ops/k510/op_utils.h>
#endif
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <iterator>
#include <nncase/kernels/kernel_context.h>
#include <nncase/kernels/kernel_utils.h>
#include <nncase/kernels/tensor_compute.h>
#include <nncase/runtime/bfloat16.h>
#include <nncase/runtime/k510/runtime_types.h>
#include <nncase/runtime/nnil.h>
#include <nncase/runtime/runtime_op_utility.h>
#include <stdexcept>
#include <type_traits>

#ifdef NNCASE_HALIDE
#include <hkg/export/HalideBuffer.h>
#include <hkg/export/halide_gnne_conv2d.h>
#include <hkg/export/halide_gnne_conv2d_depthwise.h>

template <>
HALIDE_ALWAYS_INLINE halide_type_t halide_type_of<nncase::bfloat16>()
{
    return halide_type_t(halide_type_bfloat, 16);
}

#endif

#ifdef NNCASE_OPENMP
#include <omp.h>
#endif

using namespace nncase::runtime::k510;
#ifndef ENABLE_K510_RUNTIME
using namespace nncase::ir::k510;
#endif
namespace nncase::kernels::k510
{
#define UNFOLD_OUTPUT_ASSIGN(_output_ptr) \
    auto *o0 = _output_ptr + 0;           \
    auto *o1 = _output_ptr + 1;           \
    auto *o2 = _output_ptr + 2;           \
    auto *o3 = _output_ptr + 3;

#define UNFOLD_ASSIGN_4(_input_ptr, _output_ptr) \
    auto *i0 = _input_ptr + 0;                   \
    auto *i1 = _input_ptr + 1;                   \
    auto *i2 = _input_ptr + 2;                   \
    auto *i3 = _input_ptr + 3;                   \
    UNFOLD_OUTPUT_ASSIGN(_output_ptr)

#define UNFOLD_CAST_4()         \
    *o0 = static_cast<TO>(*i0); \
    *o1 = static_cast<TO>(*i1); \
    *o2 = static_cast<TO>(*i2); \
    *o3 = static_cast<TO>(*i3);

#define UNFOLD_ROUND()                      \
    *o0 = bfloat16::round_to_bfloat16(*i0); \
    *o1 = bfloat16::round_to_bfloat16(*i1); \
    *o2 = bfloat16::round_to_bfloat16(*i2); \
    *o3 = bfloat16::round_to_bfloat16(*i3);

#define DIV_4(arg)                      \
    auto res = div((arg), 4);           \
    size_t res_quot = size_t(res.quot); \
    size_t res_rem = size_t(res.rem);

inline bfloat16 apply_gnne_activation(float value, bfloat16 x0, bfloat16 kl, bfloat16 bl, bfloat16 kr, bfloat16 br)
{
    if (value < x0)
    {
        return bfloat16::round_to_bfloat16(value * kl + bl);
    }
    else
    {
        return bfloat16::round_to_bfloat16(value * kr + br);
    }
}

#ifndef ENABLE_K510_RUNTIME
inline void gnne_transpose(const bfloat16 *input, bfloat16 *output, const runtime_shape_t &in_shape, const mfu_trans_permute &gnne_perm)
{
    auto perm = to_axis(gnne_perm);
    runtime_shape_t out_shape(4);
    for (size_t i = 0; i < 4; i++)
    {
        out_shape[i] = in_shape[perm[i]];
    }
    DIV_4(out_shape[3])
    size_t move_distance = std::accumulate(in_shape.begin() + perm[3] + 1, in_shape.end(), 1, std::multiplies<size_t>());

    auto out_img_size = out_shape[2] * out_shape[3];
    for (size_t b = 0; b < out_shape[0]; b++)
    {
#ifdef NNCASE_OPENMP
#pragma omp parallel for num_threads(kernels::default_kernel_context().num_threads)
#endif
        for (size_t c = 0; c < out_shape[1]; c++)
        {
            runtime_shape_t index(4);
            index[perm[0]] = b;
            index[perm[1]] = c;
            auto *output_ptr = output + b * out_shape[1] * out_img_size + c * out_img_size;
            for (size_t h = 0; h < out_shape[2]; h++)
            {
                index[perm[2]] = h;
                auto *input_ptr = input + linear_index(in_shape, index);
                for (size_t w = 0; w < res_quot; w++)
                {
                    auto *i0 = input_ptr + 0 * move_distance;
                    auto *i1 = input_ptr + 1 * move_distance;
                    auto *i2 = input_ptr + 2 * move_distance;
                    auto *i3 = input_ptr + 3 * move_distance;
                    UNFOLD_OUTPUT_ASSIGN(output_ptr);
                    *o0 = *i0;
                    *o1 = *i1;
                    *o2 = *i2;
                    *o3 = *i3;
                    input_ptr += 4 * move_distance;
                    output_ptr += 4;
                }
                for (size_t w = 0; w < res_rem; w++)
                {
                    *output_ptr = *input_ptr;
                    input_ptr += move_distance;
                    ++output_ptr;
                }
            }
        }
    }
}
#endif

inline void gnne_conv2d_ref(const bfloat16 *input, bfloat16 *output, const bfloat16 *weights, const float *psum, const bfloat16 *act, const runtime_shape_t &in_shape,
    int32_t groups, int32_t out_channels, int32_t filter_h, int32_t filter_w, int32_t stride_h, int32_t stride_w, int32_t dilation_h, int32_t dilation_w,
    const padding &padding_h, const padding &padding_w, value_range<bfloat16> fused_clamp, bool psum_is_uninitialized)
{
    const auto out_h = detail::get_windowed_output_size(in_shape[2], filter_h, stride_h, dilation_h, padding_h);
    const auto out_w = detail::get_windowed_output_size(in_shape[3], filter_w, stride_w, dilation_w, padding_w);
    const auto g_ic = in_shape[1] / groups;
    const auto g_oc = out_channels / groups;

    for (size_t batch = 0; batch < in_shape[0]; batch++)
    {
        const bfloat16 *in_batch_p = input + batch * in_shape[1] * in_shape[2] * in_shape[3];
        for (int32_t og = 0; og < groups; og++)
        {
            const bfloat16 *in_group_p = in_batch_p + og * g_ic * in_shape[2] * in_shape[3];
            const bfloat16 *w_group_p = weights + og * g_oc * g_ic * filter_h * filter_w;
            for (int32_t oc = 0; oc < g_oc; oc++)
            {
                const bfloat16 *w_oc_p = w_group_p + oc * g_ic * filter_h * filter_w;
                for (size_t oy = 0; oy < out_h; oy++)
                {
                    for (size_t ox = 0; ox < out_w; ox++)
                    {
                        const int32_t in_y_origin = (oy * stride_h) - padding_h.before;
                        const int32_t in_x_origin = (ox * stride_w) - padding_w.before;
                        const int32_t filter_y_start = std::max(0, (-in_y_origin + dilation_h - 1) / dilation_h);
                        const int32_t filter_y_end = std::min(filter_h, ((int32_t)in_shape[2] - in_y_origin + dilation_h - 1) / dilation_h);
                        const int32_t filter_x_start = std::max(0, (-in_x_origin + dilation_w - 1) / dilation_w);
                        const int32_t filter_x_end = std::min(filter_w, ((int32_t)in_shape[3] - in_x_origin + dilation_w - 1) / dilation_w);
                        float value = 0.f;

                        for (int32_t ic = 0; ic < (int32_t)g_ic; ic++)
                        {
                            const bfloat16 *in_c_p = in_group_p + ic * in_shape[2] * in_shape[3];
                            const bfloat16 *w_ic_p = w_oc_p + ic * filter_h * filter_w;
                            for (int32_t ky = filter_y_start; ky < filter_y_end; ky++)
                            {
                                for (int32_t kx = filter_x_start; kx < filter_x_end; kx++)
                                {
                                    const int32_t in_y = in_y_origin + dilation_h * ky;
                                    const int32_t in_x = in_x_origin + dilation_w * kx;
                                    const bfloat16 in_v = in_c_p[in_y * in_shape[3] + in_x];
                                    const bfloat16 w = w_ic_p[ky * filter_w + kx];
                                    value = value + (float)in_v * w;
                                }
                            }
                        }

                        int32_t base = og * g_oc + oc;
                        if (!psum_is_uninitialized)
                            value += *psum++;
                        auto result = apply_gnne_activation(value, act[base * 5], act[base * 5 + 1], act[base * 5 + 2], act[base * 5 + 3], act[base * 5 + 4]);
                        *output = detail::apply_activation(result, fused_clamp);
                        output++;
                    }
                }
            }
        }
    }
}
#ifdef NNCASE_HALIDE

#define HALIDE_GNNE_CONV2D_NXM_S1_S2(KH, KW)                                                                                                \
    if (filter_h == (KH) && filter_w == (KW))                                                                                               \
    {                                                                                                                                       \
        const auto out_h = detail::get_windowed_output_size(in_shape[2], filter_h, stride_h, dilation_h, padding_h);                        \
        const auto out_w = detail::get_windowed_output_size(in_shape[3], filter_w, stride_w, dilation_w, padding_w);                        \
        Halide::Runtime::Buffer<bfloat16> _input_buffer(const_cast<bfloat16 *>(input), in_shape[3], in_shape[2], in_shape[1], in_shape[0]); \
        Halide::Runtime::Buffer<bfloat16> _weights_buffer(const_cast<bfloat16 *>(weights), filter_w, filter_h, in_shape[1], out_channels);  \
        Halide::Runtime::Buffer<float> _psum_buffer(const_cast<float *>(psum), out_w, out_h, out_channels, in_shape[0]);                    \
        Halide::Runtime::Buffer<bfloat16> _act_buffer(const_cast<bfloat16 *>(act), 5, out_channels);                                        \
        bfloat16 v_range[2] = { fused_clamp.min, fused_clamp.max };                                                                         \
        Halide::Runtime::Buffer<bfloat16> _value_range_buffer(v_range, 2);                                                                  \
        Halide::Runtime::Buffer<bfloat16> _Psumed_buffer(output, out_w, out_h, out_channels, in_shape[0]);                                  \
        if (stride_h == stride_w && (stride_h == 1 || stride_h == 2))                                                                       \
        {                                                                                                                                   \
            halide_gnne_conv2d_##KH##x##KW(_input_buffer, _weights_buffer, _psum_buffer, _act_buffer,                                       \
                _value_range_buffer, psum_is_uninitialized,                                                                                 \
                padding_h.before, padding_h.after,                                                                                          \
                padding_w.before, padding_w.after,                                                                                          \
                stride_h, stride_w, _Psumed_buffer);                                                                                        \
            return;                                                                                                                         \
        }                                                                                                                                   \
    }

#define HALIDE_GNNE_CONV2D_DEPTHWISE_NXM_S1_S2(KH, KW)                                                                                      \
    if (filter_h == (KH) && filter_w == (KW))                                                                                               \
    {                                                                                                                                       \
        const auto out_h = detail::get_windowed_output_size(in_shape[2], filter_h, stride_h, dilation_h, padding_h);                        \
        const auto out_w = detail::get_windowed_output_size(in_shape[3], filter_w, stride_w, dilation_w, padding_w);                        \
        Halide::Runtime::Buffer<bfloat16> _input_buffer(const_cast<bfloat16 *>(input), in_shape[3], in_shape[2], in_shape[1], in_shape[0]); \
        Halide::Runtime::Buffer<bfloat16> _weights_buffer(const_cast<bfloat16 *>(weights), filter_w, filter_h, 1, out_channels);            \
        Halide::Runtime::Buffer<float> _psum_buffer(const_cast<float *>(psum), out_w, out_h, out_channels, in_shape[0]);                    \
        Halide::Runtime::Buffer<bfloat16> _act_buffer(const_cast<bfloat16 *>(act), 5, out_channels);                                        \
        bfloat16 v_range[2] = { fused_clamp.min, fused_clamp.max };                                                                         \
        Halide::Runtime::Buffer<bfloat16> _value_range_buffer(v_range, 2);                                                                  \
        Halide::Runtime::Buffer<bfloat16> _Psumed_buffer(output, out_w, out_h, out_channels, in_shape[0]);                                  \
        if (stride_h == stride_w && (stride_h == 1 || stride_h == 2))                                                                       \
        {                                                                                                                                   \
            halide_gnne_conv2d_depthwise_##KH##x##KW(_input_buffer, _weights_buffer, _psum_buffer, _act_buffer,                             \
                _value_range_buffer, psum_is_uninitialized,                                                                                 \
                padding_h.before, padding_h.after,                                                                                          \
                padding_w.before, padding_w.after,                                                                                          \
                stride_h, stride_w, _Psumed_buffer);                                                                                        \
            return;                                                                                                                         \
        }                                                                                                                                   \
    }

#endif

inline void gnne_conv2d(const bfloat16 *input, bfloat16 *output, const bfloat16 *weights, const float *psum, const bfloat16 *act, const runtime_shape_t &in_shape,
    int32_t groups, int32_t out_channels, int32_t filter_h, int32_t filter_w, int32_t stride_h, int32_t stride_w, int32_t dilation_h, int32_t dilation_w,
    const padding &padding_h, const padding &padding_w, value_range<bfloat16> fused_clamp, bool psum_is_uninitialized)
{
#ifdef NNCASE_HALIDE
    if (dilation_h == 1 && dilation_w == 1)
    {
        if (groups == 1)
        {
            // clang-format off
        HALIDE_GNNE_CONV2D_NXM_S1_S2(1, 1)
        else HALIDE_GNNE_CONV2D_NXM_S1_S2(3, 3)
        else HALIDE_GNNE_CONV2D_NXM_S1_S2(5, 5)
        else HALIDE_GNNE_CONV2D_NXM_S1_S2(7, 7)
            // clang-format on
        }

        if ((size_t)groups == in_shape[1] && groups == out_channels)
        {
            // clang-format off
        HALIDE_GNNE_CONV2D_DEPTHWISE_NXM_S1_S2(1, 1)
        else HALIDE_GNNE_CONV2D_DEPTHWISE_NXM_S1_S2(3, 3)
        else HALIDE_GNNE_CONV2D_DEPTHWISE_NXM_S1_S2(5, 5)
        else HALIDE_GNNE_CONV2D_DEPTHWISE_NXM_S1_S2(7, 7)
            // clang-format on
        }
    }

#endif
    gnne_conv2d_ref(input, output, weights, psum, act, in_shape,
        groups, out_channels,
        filter_h, filter_w,
        stride_h, stride_w,
        dilation_h, dilation_w,
        padding_h, padding_w,
        fused_clamp, psum_is_uninitialized);
}

inline void gnne_matmul_ref(const bfloat16 *input_a, const bfloat16 *input_b, bfloat16 *output, const bfloat16 *act, int32_t batch, int32_t a_rows, int32_t a_cols, int32_t b_cols, const value_range<bfloat16> &fused_activation)
{
    for (int32_t ob = 0; ob < batch; ob++)
    {
        for (int32_t oy = 0; oy < a_rows; oy++)
        {
            for (int32_t ox = 0; ox < b_cols; ox++)
            {
                float value = 0.f;

                for (int32_t i = 0; i < a_cols; i++)
                {
                    const auto a = input_a[ob * a_rows * a_cols + oy * a_cols + i];
                    const auto b = input_b[ob * a_cols * b_cols + i * b_cols + ox];
                    value = value + (float)a * b;
                }

                auto result = apply_gnne_activation(value, act[ob * 5 + 0], act[ob * 5 + 1], act[ob * 5 + 2], act[ob * 5 + 3], act[ob * 5 + 4]);
                output[ob * a_rows * b_cols + oy * b_cols + ox] = detail::apply_activation(result, fused_activation);
            }
        }
    }
}

inline void gnne_matmul(const bfloat16 *input_a, const bfloat16 *input_b, bfloat16 *output, const bfloat16 *act, int32_t batch, int32_t a_rows, int32_t a_cols, int32_t b_cols, const value_range<bfloat16> &fused_activation)
{
#ifdef NNCASE_HALIDE

#endif
    gnne_matmul_ref(input_a, input_b, output, act, batch, a_rows, a_cols, b_cols, fused_activation);
}

template <class TI, class TO>
void gnne_load_impl(const TI *input, TO *output, const runtime_shape_t &in_shape, const int32_t &channel_axis, const gsl::byte *deq_params)
{
    auto size = detail::compute_size(in_shape);
    DIV_4(size)
    NNCASE_UNUSED auto *output_ptr = output;
    NNCASE_UNUSED const auto *input_ptr = input;
    if constexpr (std::is_same_v<TI, TO> || (std::is_same_v<TI, bfloat16> && std::is_same_v<TO, float>))
    {
        for (size_t i = 0; i < res_quot; i++)
        {
            UNFOLD_ASSIGN_4(input_ptr, output_ptr);
            UNFOLD_CAST_4();
            input_ptr += 4;
            output_ptr += 4;
        }
        for (size_t i = 0; i < res_rem; i++)
        {
            *output_ptr = static_cast<TO>(*input_ptr);
            output_ptr++;
            input_ptr++;
        }
    }
    else if constexpr (std::is_same_v<TI, float> && std::is_same_v<TO, bfloat16>)
    {
        for (size_t i = 0; i < res_quot; i++)
        {
            UNFOLD_ASSIGN_4(input_ptr, output_ptr);
            UNFOLD_ROUND();
            input_ptr += 4;
            output_ptr += 4;
        }
        for (size_t i = 0; i < res_rem; i++)
        {
            *output_ptr = bfloat16::round_to_bfloat16(*input_ptr);
            output_ptr++;
            input_ptr++;
        }
    }
    else if constexpr ((std::is_same_v<TI, int8_t> || std::is_same_v<TI, uint8_t>)&&(std::is_same_v<TO, bfloat16> || std::is_same_v<TO, float>))
    {
        size_t size_per_channel = 1;
        for (int32_t i = channel_axis + 1; i < 4; i++)
        {
            size_per_channel *= in_shape[i];
        }
        for (size_t i = 0; i < detail::compute_size(in_shape); i++)
        {
            auto *scale = reinterpret_cast<const int16_t *>(&deq_params[(i / size_per_channel) * 4]);
            auto *shift = reinterpret_cast<const int8_t *>(&deq_params[(i / size_per_channel) * 4 + 2]);
            auto *bias = reinterpret_cast<const TI *>(&deq_params[(i / size_per_channel) * 4 + 3]);
            auto value = (input[i] - *bias) * (*scale) / pow(2, *shift);
            output[i] = static_cast<TO>(value);
        }
    }
    else
    {
        std::runtime_error("unsupported convert type in gnne_load!");
    }
}

#define GNNE_LOAD_IMPL_LV2(input_t, output_t) \
    if (out_type == to_datatype<output_t>())  \
    gnne_load_impl(reinterpret_cast<const input_t *>(input), reinterpret_cast<output_t *>(output), in_shape, channel_axis, deq_params)

#define GNNE_LOAD_IMPL_LV1(input_t)                 \
    if (in_type == to_datatype<input_t>())          \
    {                                               \
        GNNE_LOAD_IMPL_LV2(input_t, uint8_t);       \
        else GNNE_LOAD_IMPL_LV2(input_t, int8_t);   \
        else GNNE_LOAD_IMPL_LV2(input_t, bfloat16); \
        else GNNE_LOAD_IMPL_LV2(input_t, float);    \
    }

inline void gnne_load(datatype_t in_type, datatype_t out_type, const gsl::byte *input, gsl::byte *output, const runtime_shape_t &in_shape, const int32_t &channel_axis, const gsl::byte *deq_params)
{
    // clang-format off
    GNNE_LOAD_IMPL_LV1(uint8_t)
    else GNNE_LOAD_IMPL_LV1(int8_t)
    else GNNE_LOAD_IMPL_LV1(bfloat16)
    else GNNE_LOAD_IMPL_LV1(float)
    else std::runtime_error("unsupported data type in gnne_load!");
    // clang-format on
}

template <class TI, class TO>
void gnne_store_impl(const TI *input, TO *output, const runtime_shape_t &in_shape, const int32_t &channel_axis, const bfloat16 *q_params, const runtime_shape_t &strides)
{
    auto size = detail::compute_size(in_shape);
    DIV_4(size)
    auto *output_ptr = output;
    const auto *input_ptr = input;
    if (std::is_same_v<TI, TO> || (std::is_same_v<TI, float> && std::is_same_v<TO, bfloat16>)
        || (std::is_same_v<TI, bfloat16> && std::is_same_v<TO, float>))
    {
        if (detail::compute_size(strides) == 1)
        {
            for (size_t i = 0; i < res_quot; i++)
            {
                UNFOLD_ASSIGN_4(input_ptr, output_ptr);
                UNFOLD_CAST_4();
                output_ptr += 4;
                input_ptr += 4;
            }
            for (size_t i = 0; i < res_rem; i++)
            {
                *output_ptr = static_cast<TO>(*input_ptr);
                output_ptr++;
                input_ptr++;
            }
        }
        else
        {
            for (size_t n = 0; n < in_shape[0]; n++)
            {
                if (n % strides[0] == 0)
                {
                    for (size_t c = 0; c < in_shape[1]; c++)
                    {
                        if (c % strides[1] == 0)
                        {
                            for (size_t h = 0; h < in_shape[2]; h++)
                            {
                                if (h % strides[2] == 0)
                                {
                                    for (size_t w = 0; w < in_shape[3]; w++)
                                    {
                                        if (w % strides[3] == 0)
                                        {
                                            *output_ptr++ = static_cast<TO>(*(input_ptr + (n * in_shape[1] * in_shape[2] * in_shape[3] + c * in_shape[2] * in_shape[3] + h * in_shape[3] + w)));
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    else if ((std::is_same_v<TO, int8_t> || std::is_same_v<TO, uint8_t>)&&std::is_same_v<TI, bfloat16>)
    {
        size_t size_per_channel = 1;
        for (int32_t i = channel_axis + 1; i < 4; i++)
        {
            size_per_channel *= in_shape[i];
        }
        for (size_t i = 0; i < detail::compute_size(in_shape); i++)
        {
            auto value = (float)input[i] * q_params[(i / size_per_channel) * 2] + q_params[(i / size_per_channel) * 2 + 1];
            if (std::is_same_v<TO, int8_t>)
                value = std::clamp(value, -128.f, 127.f);
            else
                value = std::clamp(value, 0.f, 255.f);
            output[i] = static_cast<TO>(value);
        }
    }
    else
    {
        std::runtime_error("unsupported convert type in gnne_store!");
    }
}

#define GNNE_STORE_IMPL_LV2(input_t, output_t) \
    if (out_type == to_datatype<output_t>())   \
    gnne_store_impl(reinterpret_cast<const input_t *>(input), reinterpret_cast<output_t *>(output), in_shape, channel_axis, q_params, strides)

#define GNNE_STORE_IMPL_LV1(input_t)                 \
    if (in_type == to_datatype<input_t>())           \
    {                                                \
        GNNE_STORE_IMPL_LV2(input_t, uint8_t);       \
        else GNNE_STORE_IMPL_LV2(input_t, int8_t);   \
        else GNNE_STORE_IMPL_LV2(input_t, bfloat16); \
        else GNNE_STORE_IMPL_LV2(input_t, float);    \
    }

inline void gnne_store(datatype_t in_type, datatype_t out_type, const gsl::byte *input, gsl::byte *output, const runtime_shape_t &in_shape, const int32_t &channel_axis, const bfloat16 *q_params, const runtime_shape_t &strides)
{
    // clang-format off
    GNNE_STORE_IMPL_LV1(uint8_t)
    else GNNE_STORE_IMPL_LV1(int8_t)
    else GNNE_STORE_IMPL_LV1(bfloat16)
    else GNNE_STORE_IMPL_LV1(float)
    else std::runtime_error("unsupported data type in gnne_store!");
    // clang-format on
}

template <class TBinaryOp, class TOutputOp>
void gnne_pdp_reduce(const bfloat16 *input, bfloat16 *output, const runtime_shape_t &in_shape, int32_t filter_h, int32_t filter_w,
    int32_t stride_h, int32_t stride_w, const padding &padding_h, const padding &padding_w, TBinaryOp &&binary_op, TOutputOp &&window_op)
{
    const auto out_h = kernels::detail::get_windowed_output_size(in_shape[2], filter_h, stride_h, 1, padding_h);
    const auto out_w = kernels::detail::get_windowed_output_size(in_shape[3], filter_w, stride_w, 1, padding_w);
    runtime_shape_t out_shape { in_shape[0], in_shape[1], (size_t)out_h, (size_t)out_w };

    for (size_t batch = 0; batch < in_shape[0]; batch++)
    {
        for (size_t oc = 0; oc < in_shape[1]; oc++)
        {
            for (size_t oy = 0; oy < out_h; oy++)
            {
                for (size_t ox = 0; ox < out_w; ox++)
                {
                    const int32_t in_y_origin = (oy * stride_h) - padding_h.before;
                    const int32_t in_x_origin = (ox * stride_w) - padding_w.before;
                    const int32_t filter_y_start = std::max(0, -in_y_origin);
                    const int32_t filter_y_end = std::min(filter_h, (int32_t)in_shape[2] - in_y_origin);
                    const int32_t filter_x_start = std::max(0, -in_x_origin);
                    const int32_t filter_x_end = std::min(filter_w, (int32_t)in_shape[3] - in_x_origin);
                    float value = input[linear_index(in_shape, { batch, oc, (size_t)(in_y_origin + filter_y_start), (size_t)((in_x_origin + filter_x_start)) })];
                    int32_t kernel_count = 0;

                    for (int32_t ky = filter_y_start; ky < filter_y_end; ky++)
                    {
                        for (int32_t kx = filter_x_start; kx < filter_x_end; kx++)
                        {
                            const int32_t in_y = in_y_origin + ky;
                            const int32_t in_x = in_x_origin + kx;

                            const float in_v = input[linear_index(in_shape, { batch, oc, (size_t)in_y, (size_t)in_x })];

                            if (!((ky == filter_y_start) && (kx == filter_x_start)))
                                value = binary_op(value, in_v);
                            kernel_count++;
                        }
                    }

                    output[linear_index(out_shape, { batch, oc, (size_t)oy, (size_t)ox })] = bfloat16::round_to_bfloat16(window_op(value, kernel_count));
                }
            }
        }
    }
}

template <class TReducer>
void gnne_reduce(const bfloat16 *input, bfloat16 *output, bfloat16 init_value, const runtime_shape_t &in_shape, const runtime_shape_t &reduced_shape, TReducer &&reducer)
{
    auto output_size = kernels::detail::compute_size(reduced_shape);
    auto out_tmp = std::make_unique<float[]>(output_size);
    // op is sub, maybe a error with op sub cauesd by order of operation
    // 1 - (2 - 1) != 1 - 2 - 1
    if (reducer(1, 2) == -1)
    {
        std::fill(out_tmp.get(), out_tmp.get() + output_size, init_value);

        for (size_t d0 = 0; d0 < in_shape[0]; d0++)
        {
            for (size_t d1 = 0; d1 < in_shape[1]; d1++)
            {
                for (size_t d2 = 0; d2 < in_shape[2]; d2++)
                {
                    for (size_t d3 = 0; d3 < in_shape[3]; d3++)
                    {
                        runtime_shape_t in_off = { d0, d1, d2, d3 };
                        auto out_off = kernels::detail::get_reduced_offset(in_off, reduced_shape);
                        const auto a = input[linear_index(in_shape, in_off)];
                        auto &b = out_tmp[linear_index(reduced_shape, out_off)];
                        b = reducer(b, (float)a);
                    }
                }
            }
        }
    }
    else
    {
        NNCASE_UNUSED auto *output_tmp_ptr = out_tmp.get();
        auto input_size = kernels::detail::compute_size(in_shape);
        auto block_size = input_size / output_size;
#ifdef NNCASE_OPENMP
#pragma omp parallel for num_threads(kernels::default_kernel_context().num_threads)
#endif
        for (size_t i = 0; i < output_size; i++)
        {
            auto input_ptr = input + i * block_size;
            auto output_ptr = out_tmp.get() + i;
            *output_ptr = init_value;
            // sub memory reduce
            for (size_t j = 0; j < block_size; ++j)
            {
                *output_ptr = reducer(*output_ptr, static_cast<float>(*input_ptr));
                ++input_ptr;
            }
            ++output_ptr;
        }
    }

    // output round
    for (auto i = 0U; i < output_size; i++)
    {
        output[i] = bfloat16::round_to_bfloat16(out_tmp[i]);
    }
}
#ifndef ENABLE_K510_RUNTIME
inline void gnne_mn_map(const bfloat16 *input, bfloat16 *output, bfloat16 *seg_fitting_param, const runtime_shape_t &in_shape, const meshnet::graph &graph)
{
    std::vector<const meshnet::node *> nodes;
    graph.dfs_visit(graph.vector_in0, [&](const meshnet::node &rnode) { nodes.push_back(&rnode); });

    auto output_size = kernels::detail::compute_size(in_shape);
    auto output_tmp = std::make_unique<bfloat16[]>(output_size);
    std::copy(input, input + output_size, output_tmp.get());

    auto unary = [&](const meshnet::node &node) {
        std::function<bfloat16(bfloat16)> f;
        switch (node.id())
        {
        case meshnet::node_id::exp0:
            f = [](bfloat16 a) { return std::exp(a); };
            break;
        case meshnet::node_id::square0:
            f = [](bfloat16 a) { return a * a; };
            break;
        case meshnet::node_id::log0:
        {
            auto log = static_cast<const meshnet::log *>(&node);
            if (log->mode() == mn_log_mode::positive)
                f = [](bfloat16 a) { return std::log(a); };
            else if (log->mode() == mn_log_mode::nan)
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::log(a);
                    else
                        return bfloat16::nan();
                };
            else
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::log(a);
                    else
                        return -std::log(-a);
                };
            break;
        }
        case meshnet::node_id::round0:
        {
            auto rnd = static_cast<const meshnet::round *>(&node);
            if (rnd->mode() == mn_round_mode::round)
                f = [](bfloat16 a) { return std::round(a); };
            else if (rnd->mode() == mn_round_mode::ceil)
                f = [](bfloat16 a) { return std::ceil(a); };
            else
                f = [](bfloat16 a) { return std::floor(a); };
            break;
        }
        case meshnet::node_id::sqrt0:
        {
            auto sqrt = static_cast<const meshnet::sqrt *>(&node);
            if (sqrt->mode() == mn_sqrt_mode::positive)
                f = [](bfloat16 a) { return std::sqrt(a); };
            else if (sqrt->mode() == mn_sqrt_mode::nan)
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::sqrt(a);
                    else
                        return bfloat16::nan();
                };
            else
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::sqrt(a);
                    else
                        return -std::sqrt(-a);
                };
            break;
        }
        case meshnet::node_id::triangle0:
        {
            auto triangle = static_cast<const meshnet::triangle *>(&node);
            if (triangle->mode() == mn_triangle_mode::cos)
                f = [](bfloat16 a) { return std::cos(a); };
            else
                f = [](bfloat16 a) { return std::sin(a); };
            break;
        }
        case meshnet::node_id::unary_logic0:
        {
            auto logic = static_cast<const meshnet::unary_logic *>(&node);
            if (logic->mode() == mn_unary_logic_mode::abs)
                f = [](bfloat16 a) { return std::abs(a); };
            else if (logic->mode() == mn_unary_logic_mode::neg)
                f = [](bfloat16 a) { return -a; };
            else
                f = [](bfloat16 a) { return a > 0 ? bfloat16(1) : (a < 0 ? bfloat16(-1) : bfloat16(0)); };
            break;
        }
        case meshnet::node_id::vector_out0:
        case meshnet::node_id::vector_in0:
            f = [](bfloat16 a) { return a; };
            break;
        case meshnet::node_id::reg0:
            f = [&](bfloat16 a) {
                auto index = 15;
                for (auto i = 0; i < 15; i++)
                {
                    if (a < seg_fitting_param[32 + i])
                    {
                        index = i;
                        break;
                    }
                }
                return a * seg_fitting_param[2 * index] + seg_fitting_param[2 * index + 1];
            };
            break;
        case meshnet::node_id::binary_logic0:
        {
            auto logic = static_cast<const meshnet::binary_logic *>(&node);
            if (logic->islogic() == mn_binary_logic_islogic::logic)
            {
                if (logic->mode() == mn_binary_logic_mode::knot)
                {
                    f = [](bfloat16 a) { return bfloat16(!a); };
                }
                else
                {
                    throw std::runtime_error("Not supported logical mode for unary");
                }
            }
            break;
        }
        default:
            throw std::runtime_error("Not supported meshnet id: " + to_string(node.id()));
            break;
        }

        return f;
    };

    for (auto &node : nodes)
    {
        auto op = unary(*node);
        for (auto i = 0U; i < output_size; i++)
            output_tmp[i] = op(output_tmp[i]);
    }

    std::copy(output_tmp.get(), output_tmp.get() + output_size, output);
}
#endif
inline int get_pad_index(const runtime_shape_t &in_shape, int n, int c, int h, int w, pad_mode_t mode, const runtime_paddings_t &paddings)
{
    int ret_index = 0;
    switch (mode)
    {
    case pad_reflect:
        if (h < (int)paddings[2].before && w < (int)paddings[3].before)
        {
            ret_index = n * in_shape[0] + c * in_shape[1] + (paddings[2].before - h) * in_shape[2] + paddings[3].before - w;
        }
        else if (h < (int)paddings[2].before && w >= (int)paddings[3].before)
        {
            if (w > (int)(in_shape[3] + paddings[3].before - 1))
            {
                auto dis_w = w - (in_shape[3] + paddings[3].before - 1);
                ret_index = n * in_shape[0] + c * in_shape[1] + (paddings[2].before - h) * in_shape[2] + in_shape[3] - dis_w - 1;
            }
            else
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (paddings[2].before - h) * in_shape[2] + w - paddings[3].before;
            }
        }
        else if (h >= (int)paddings[2].before && w < (int)paddings[3].before)
        {
            if (h > (int)(in_shape[2] + paddings[2].before - 1))
            {
                auto dis_h = h - (in_shape[2] + paddings[2].before - 1);
                ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - dis_h - 1) * in_shape[2] + paddings[3].before - w;
            }
            else
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + h * in_shape[2] + paddings[3].before - w;
            }
        }
        else if (h >= (int)paddings[2].before && w >= (int)paddings[3].before)
        {
            auto dis_h = h - (in_shape[2] + paddings[2].before - 1);
            auto dis_w = w - (in_shape[3] + paddings[3].before - 1);
            if (h <= (int)(in_shape[2] + paddings[2].before - 1) && w > (int)(in_shape[3] + paddings[3].before - 1))
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (h - paddings[2].before - 1) * in_shape[2] + in_shape[3] - dis_w - 1;
            }
            else if (h > (int)(in_shape[2] + paddings[2].before - 1) && w <= (int)(in_shape[3] + paddings[3].before - 1))
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - dis_h - 1) * in_shape[2] + w - paddings[3].before - 1;
            }
            else if (h > (int)(in_shape[2] + paddings[2].before - 1) && w > (int)(in_shape[3] + paddings[3].before - 1))
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - dis_h - 1) * in_shape[2] + in_shape[3] - dis_w - 1;
            }
        }
        break;
    case pad_symmetric:
        if (h < (int)paddings[2].before && w < (int)paddings[3].before)
        {
            ret_index = n * in_shape[0] + c * in_shape[1] + (paddings[2].before - h - 1) * in_shape[2] + paddings[3].before - w - 1;
        }
        else if (h < (int)paddings[2].before && w >= (int)paddings[3].before)
        {
            if (w > (int)(in_shape[3] + paddings[3].before - 1))
            {
                auto dis_w = w - (in_shape[3] + paddings[3].before - 1);
                ret_index = n * in_shape[0] + c * in_shape[1] + (paddings[2].before - h - 1) * in_shape[2] + in_shape[3] - dis_w;
            }
            else
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (paddings[2].before - h - 1) * in_shape[2] + w - paddings[3].before;
            }
        }
        else if (h >= (int)paddings[2].before && w < (int)paddings[3].before)
        {
            if (h > (int)(in_shape[2] + paddings[2].before - 1))
            {
                auto dis_h = h - (in_shape[2] + paddings[2].before - 1);
                ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - dis_h) * in_shape[2] + paddings[3].before - w - 1;
            }
            else
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + h * in_shape[2] + paddings[3].before - w;
            }
        }
        else if (h >= (int)paddings[2].before && w >= (int)paddings[3].before)
        {
            auto dis_h = h - (in_shape[2] + paddings[2].before - 1);
            auto dis_w = w - (in_shape[3] + paddings[3].before - 1);
            if (h <= (int)(in_shape[2] + paddings[2].before - 1) && w > (int)(in_shape[3] + paddings[3].before - 1))
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (h - paddings[2].before) * in_shape[2] + in_shape[3] - dis_w;
            }
            else if (h > (int)(in_shape[2] + paddings[2].before - 1) && w <= (int)(in_shape[3] + paddings[3].before - 1))
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - dis_h) * in_shape[2] + w - paddings[3].before;
            }
            else if (h > (int)(in_shape[2] + paddings[2].before - 1) && w > (int)(in_shape[3] + paddings[3].before - 1))
            {
                ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - dis_h) * in_shape[2] + in_shape[3] - dis_w;
            }
        }
        break;
    // case pad_edge:
    //     if (h < (int)paddings[2].before && w < (int)paddings[3].before)
    //     {
    //         ret_index = n * in_shape[0] + c * in_shape[1] + 0 * in_shape[2] + 0;
    //     }
    //     else if (h < (int)paddings[2].before && w >= (int)paddings[3].before)
    //     {
    //         if (w > (int)(in_shape[3] + paddings[3].before - 1))
    //         {
    //             ret_index = n * in_shape[0] + c * in_shape[1] + 0 * in_shape[2] + in_shape[3] - 1;
    //         }
    //         else
    //         {
    //             ret_index = n * in_shape[0] + c * in_shape[1] + 0 * in_shape[2] + w - paddings[3].before;
    //         }
    //     }
    //     else if (h >= (int)paddings[2].before && w < (int)paddings[3].before)
    //     {
    //         if (h > (int)(in_shape[2] + paddings[2].before - 1))
    //         {
    //             ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - 1) * in_shape[2] + 0;
    //         }
    //         else
    //         {
    //             ret_index = n * in_shape[0] + c * in_shape[1] + (h - paddings[2].before) * in_shape[2];
    //         }
    //     }
    //     else if (h >= (int)paddings[2].before && w >= (int)paddings[3].before)
    //     {
    //         auto dis_h = h - (in_shape[2] + paddings[2].before - 1);
    //         auto dis_w = w - (in_shape[3] + paddings[3].before - 1);
    //         if (h <= (int)(in_shape[2] + paddings[2].before - 1) && w > (int)(in_shape[3] + paddings[3].before - 1))
    //         {
    //             ret_index = n * in_shape[0] + c * in_shape[1] + (h - paddings[2].before) * in_shape[2] + in_shape[3] - 1;
    //         }
    //         else if (h > (int)(in_shape[2] + paddings[2].before - 1) && w <= (int)(in_shape[3] + paddings[3].before - 1))
    //         {
    //             ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - 1) * in_shape[2] + w - paddings[3].before;
    //         }
    //         else if (h > (int)(in_shape[2] + paddings[2].before - 1) && w > (int)(in_shape[3] + paddings[3].before - 1))
    //         {
    //             ret_index = n * in_shape[0] + c * in_shape[1] + (in_shape[2] - 1) * in_shape[2] + in_shape[3] - 1;
    //         }
    //     }
    //     break;
    default:
        std::runtime_error("unsupport pad mode");
    };
    return ret_index;
}

template <class T>
inline void dsp_pad(const T *input, T *output, const runtime_shape_t &in_shape, const runtime_paddings_t &paddings, T pad_value, pad_mode_t mode)
{
    runtime_shape_t out_shape = { in_shape[0] + paddings[0].sum(),
        in_shape[1] + paddings[1].sum(),
        in_shape[2] + paddings[2].sum(),
        in_shape[3] + paddings[3].sum() };

    for (int d0 = 0; d0 < (int)out_shape[0]; d0++)
    {
        auto d0_origin = -paddings[0].before;
        auto in0 = input + ((int)d0_origin + d0) * in_shape[1] * in_shape[2] * in_shape[3];

        for (int d1 = 0; d1 < (int)out_shape[1]; d1++)
        {
            auto d1_origin = -paddings[1].before;
            auto in1 = in0 + ((int)d1_origin + d1) * in_shape[2] * in_shape[3];

            for (int d2 = 0; d2 < (int)out_shape[2]; d2++)
            {
                auto d2_origin = -paddings[2].before;
                auto in2 = in1 + ((int)d2_origin + d2) * in_shape[3];

                for (int d3 = 0; d3 < (int)out_shape[3]; d3++)
                {
                    auto d3_origin = -paddings[3].before;

                    if (d0 < paddings[0].before || d0 >= (int)out_shape[0] - paddings[0].after
                        || d1 < paddings[1].before || d1 >= (int)out_shape[1] - paddings[1].after
                        || d2 < paddings[2].before || d2 >= (int)out_shape[2] - paddings[2].after
                        || d3 < paddings[3].before || d3 >= (int)out_shape[3] - paddings[3].after)
                    {
                        if (mode == pad_constant)
                        {
                            *output++ = pad_value;
                        }
                        else
                        {
                            *output++ = in2[get_pad_index(in_shape, d0, d1, d2, d3, mode, paddings)];
                        }
                    }
                    else
                        *output++ = in2[d3_origin + d3];
                }
            }
        }
    }
}

template <class T>
inline void gnne_pad(const T *input, T *output, const runtime_shape_t &in_shape, const runtime_paddings_t &paddings, T pad_value, [[maybe_unused]] pad_mode_t mode)
{
    assert(mode == pad_constant);
    runtime_shape_t out_shape = { in_shape[0] + paddings[0].sum(),
        in_shape[1] + paddings[1].sum(),
        in_shape[2] + paddings[2].sum(),
        in_shape[3] + paddings[3].sum() };

    for (int d0 = 0; d0 < (int)out_shape[0]; d0++)
    {
        auto d0_origin = -paddings[0].before;
        auto in0 = input + ((int)d0_origin + d0) * in_shape[1] * in_shape[2] * in_shape[3];

        for (int d1 = 0; d1 < (int)out_shape[1]; d1++)
        {
            auto d1_origin = -paddings[1].before;
            auto in1 = in0 + ((int)d1_origin + d1) * in_shape[2] * in_shape[3];

            for (int d2 = 0; d2 < (int)out_shape[2]; d2++)
            {
                auto d2_origin = -paddings[2].before;
                auto in2 = in1 + ((int)d2_origin + d2) * in_shape[3];

                for (int d3 = 0; d3 < (int)out_shape[3]; d3++)
                {
                    auto d3_origin = -paddings[3].before;

                    if (d0 < paddings[0].before || d0 >= (int)out_shape[0] - paddings[0].after
                        || d1 < paddings[1].before || d1 >= (int)out_shape[1] - paddings[1].after
                        || d2 < paddings[2].before || d2 >= (int)out_shape[2] - paddings[2].after
                        || d3 < paddings[3].before || d3 >= (int)out_shape[3] - paddings[3].after)
                    {
                        *output++ = pad_value;
                    }
                    else
                        *output++ = in2[d3_origin + d3];
                }
            }
        }
    }
}
#ifndef ENABLE_K510_RUNTIME
inline void gnne_mn_binary(bfloat16 *input_a, bfloat16 *input_b, bfloat16 *output, const std::array<bfloat16, 4> &constants, const runtime_shape_t &in_a_shape, const runtime_shape_t &in_b_shape, const runtime_shape_t &out_shape, const meshnet::graph &graph)
{
    std::vector<const meshnet::node *> nodes;
    graph.dfs_visit(graph.vector_in0, [&](const meshnet::node &rnode) { nodes.push_back(&rnode); });

    auto unary = [](const meshnet::node &node) {
        std::function<bfloat16(bfloat16)> f;
        switch (node.id())
        {
        case meshnet::node_id::exp0:
            f = [](bfloat16 a) { return std::exp(a); };
            break;
        case meshnet::node_id::square0:
            f = [](bfloat16 a) { return a * a; };
            break;
        case meshnet::node_id::log0:
        {
            auto log = static_cast<const meshnet::log *>(&node);
            if (log->mode() == mn_log_mode::positive)
                f = [](bfloat16 a) { return std::log(a); };
            else if (log->mode() == mn_log_mode::nan)
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::log(a);
                    else
                        return bfloat16::nan();
                };
            else
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::log(a);
                    else
                        return -std::log(-a);
                };
            break;
        }
        case meshnet::node_id::round0:
        {
            auto rnd = static_cast<const meshnet::round *>(&node);
            if (rnd->mode() == mn_round_mode::round)
                f = [](bfloat16 a) { return std::round(a); };
            else if (rnd->mode() == mn_round_mode::ceil)
                f = [](bfloat16 a) { return std::ceil(a); };
            else
                f = [](bfloat16 a) { return std::floor(a); };
            break;
        }
        case meshnet::node_id::sqrt0:
        {
            auto sqrt = static_cast<const meshnet::sqrt *>(&node);
            if (sqrt->mode() == mn_sqrt_mode::positive)
                f = [](bfloat16 a) { return std::sqrt(a); };
            else if (sqrt->mode() == mn_sqrt_mode::nan)
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::sqrt(a);
                    else
                        return bfloat16::nan();
                };
            else
                f = [](bfloat16 a) {
                    if (a > 0)
                        return std::sqrt(a);
                    else
                        return -std::sqrt(-a);
                };
            break;
        }
        case meshnet::node_id::triangle0:
        {
            auto triangle = static_cast<const meshnet::triangle *>(&node);
            if (triangle->mode() == mn_triangle_mode::cos)
                f = [](bfloat16 a) { return std::cos(a); };
            else
                f = [](bfloat16 a) { return std::sin(a); };
            break;
        }
        case meshnet::node_id::unary_logic0:
        {
            auto logic = static_cast<const meshnet::unary_logic *>(&node);
            if (logic->mode() == mn_unary_logic_mode::abs)
                f = [](bfloat16 a) { return std::abs(a); };
            else if (logic->mode() == mn_unary_logic_mode::neg)
                f = [](bfloat16 a) { return -a; };
            else
                f = [](bfloat16 a) { return a > 0 ? bfloat16(1) : (a < 0 ? bfloat16(-1) : bfloat16(0)); };
            break;
        }
        default:
            break;
        }

        return f;
    };

    auto binary = [](const meshnet::node &node) {
        std::function<bfloat16(bfloat16, bfloat16)> f;
        switch (node.id())
        {
        case meshnet::node_id::add_sub0:
        case meshnet::node_id::add_sub1:
        case meshnet::node_id::add_sub2:
        case meshnet::node_id::add_sub3:
        {
            auto add_sub = static_cast<const meshnet::add_sub *>(&node);
            if (add_sub->neg_input_a() && add_sub->neg_input_b())
                f = [](bfloat16 a, bfloat16 b) { return -a - b; };
            else if (add_sub->neg_input_a() && !add_sub->neg_input_b())
                f = [](bfloat16 a, bfloat16 b) { return -a + b; };
            else if (!add_sub->neg_input_a() && add_sub->neg_input_b())
                f = [](bfloat16 a, bfloat16 b) { return a - b; };
            else
                f = [](bfloat16 a, bfloat16 b) { return a + b; };
            break;
        }
        case meshnet::node_id::mul0:
        case meshnet::node_id::mul1:
        case meshnet::node_id::mul2:
        case meshnet::node_id::mul3:
        {
            auto mul = static_cast<const meshnet::mul *>(&node);
            if (mul->neg_output())
                f = [](bfloat16 a, bfloat16 b) { return -a * b; };
            else
                f = [](bfloat16 a, bfloat16 b) { return a * b; };
            break;
        }
        case meshnet::node_id::div0:
        {
            auto div = static_cast<const meshnet::div *>(&node);
            if (div->neg_output())
                f = [](bfloat16 a, bfloat16 b) { return -a / b; };
            else

                f = [](bfloat16 a, bfloat16 b) { return a / b; };
            break;
        }
        case meshnet::node_id::cmp0:
        case meshnet::node_id::cmp1:
        {
            auto cmp = static_cast<const meshnet::cmp *>(&node);
            if (cmp->mode() == mn_cmp_mode::min)
                f = [](bfloat16 a, bfloat16 b) { return std::min(a, b); };
            else if (cmp->mode() == mn_cmp_mode::max)
                f = [](bfloat16 a, bfloat16 b) { return std::max(a, b); };
            else if (cmp->mode() == mn_cmp_mode::eq)
                f = [](bfloat16 a, bfloat16 b) { return a == b ? (bfloat16)1 : (bfloat16)0; };
            else if (cmp->mode() == mn_cmp_mode::gt)
                f = [](bfloat16 a, bfloat16 b) { return a > b ? (bfloat16)1 : (bfloat16)0; };
            else if (cmp->mode() == mn_cmp_mode::ge)
                f = [](bfloat16 a, bfloat16 b) { return a >= b ? (bfloat16)1 : (bfloat16)0; };
            else if (cmp->mode() == mn_cmp_mode::lt)
                f = [](bfloat16 a, bfloat16 b) { return a < b ? (bfloat16)1 : (bfloat16)0; };
            else if (cmp->mode() == mn_cmp_mode::le)
                f = [](bfloat16 a, bfloat16 b) { return a <= b ? (bfloat16)1 : (bfloat16)0; };
            else
                throw std::runtime_error("cmp mode not implemented!");
            break;
        }
        case meshnet::node_id::binary_logic0:
        {
            auto logic = static_cast<const meshnet::binary_logic *>(&node);
            if (logic->islogic() == mn_binary_logic_islogic::bitwise)
            {
                if (logic->mode() == mn_binary_logic_mode::kand)
                    f = [](bfloat16 a, bfloat16 b) {
                        uint16_t tmp = a.raw() & b.raw();
                        return *reinterpret_cast<bfloat16 *>(&tmp);
                    };
                else if (logic->mode() == mn_binary_logic_mode::kor)
                    f = [](bfloat16 a, bfloat16 b) {
                        uint16_t tmp = a.raw() | b.raw();
                        return *reinterpret_cast<bfloat16 *>(&tmp);
                    };
                else if (logic->mode() == mn_binary_logic_mode::kxor)
                    f = [](bfloat16 a, bfloat16 b) {
                        uint16_t tmp = a.raw() ^ b.raw();
                        return *reinterpret_cast<bfloat16 *>(&tmp);
                    };
            }
            else
            {
                if (logic->mode() == mn_binary_logic_mode::kand)
                    f = [](bfloat16 a, bfloat16 b) {
                        uint16_t tmp = a.raw() && b.raw();
                        return *reinterpret_cast<bfloat16 *>(&tmp);
                    };
                else if (logic->mode() == mn_binary_logic_mode::kor)
                    f = [](bfloat16 a, bfloat16 b) {
                        uint16_t tmp = a.raw() || b.raw();
                        return *reinterpret_cast<bfloat16 *>(&tmp);
                    };
                else if (logic->mode() == mn_binary_logic_mode::kxor)
                    f = [](bfloat16 a, bfloat16 b) {
                        uint16_t tmp = (a.raw() != 0) ^ (b.raw() != 0);
                        return *reinterpret_cast<bfloat16 *>(&tmp);
                    };
            }
            break;
        }
        default:
            break;
        }

        return f;
    };

    auto ternary = [](const meshnet::node &node) {
        std::function<bfloat16(bfloat16, bfloat16, bfloat16)> f;
        switch (node.id())
        {
        case meshnet::node_id::select0:
        case meshnet::node_id::select1:
        {
            f = [](bfloat16 a, bfloat16 b, bfloat16 c) { return c ? a : b; };
            break;
        }
        default:
            break;
        }

        return f;
    };

    std::unordered_map<const meshnet::output_connector *, bfloat16 *> output_data;
    std::unordered_map<const meshnet::output_connector *, runtime_shape_t> output_shape;
    std::vector<std::unique_ptr<bfloat16[]>> output_tmp;

    for (auto &node : nodes)
    {
        if (node->id() == meshnet::node_id::vector_out0)
        {
            output_shape.emplace(&node->output_at(0), in_a_shape);
            output_data.emplace(&node->output_at(0), input_a);
        }
        else if (node->id() == meshnet::node_id::vector_out1)
        {
            output_shape.emplace(&node->output_at(0), in_b_shape);
            output_data.emplace(&node->output_at(0), input_b);
        }
        else if (node->id() == meshnet::node_id::dup0 || node->id() == meshnet::node_id::dup1 || node->id() == meshnet::node_id::dup2
            || node->id() == meshnet::node_id::dup3)
        {
            auto out_connector = node->input_at(0).connection();
            output_shape.emplace(&node->output_at(0), output_shape[out_connector]);
            output_data.emplace(&node->output_at(0), output_data[out_connector]);

            output_shape.emplace(&node->output_at(1), output_shape[out_connector]);
            output_data.emplace(&node->output_at(1), output_data[out_connector]);
        }
        else if (node->id() == meshnet::node_id::constant0)
        {
            runtime_shape_t out_shape { 1, 1, 1, 1 };
            output_shape.emplace(&node->output_at(0), out_shape);
            output_data.emplace(&node->output_at(0), const_cast<bfloat16 *>(&constants[0]));
        }
        else if (node->id() == meshnet::node_id::constant1)
        {
            runtime_shape_t out_shape { 1, 1, 1, 1 };
            output_shape.emplace(&node->output_at(0), out_shape);
            output_data.emplace(&node->output_at(0), const_cast<bfloat16 *>(&constants[1]));
        }
        else if (node->id() == meshnet::node_id::constant2)
        {
            runtime_shape_t out_shape { 1, 1, 1, 1 };
            output_shape.emplace(&node->output_at(0), out_shape);
            output_data.emplace(&node->output_at(0), const_cast<bfloat16 *>(&constants[2]));
        }
        else if (node->id() == meshnet::node_id::constant3)
        {
            runtime_shape_t out_shape { 1, 1, 1, 1 };
            output_shape.emplace(&node->output_at(0), out_shape);
            output_data.emplace(&node->output_at(0), const_cast<bfloat16 *>(&constants[3]));
        }
        else if (node->id() == meshnet::node_id::vector_in0)
        {
            // do broadcast at vector_in0 if needed.
            if (output_shape.at(node->input_at(0).connection()) != out_shape)
            {
                size_t output_offset = 0;
                auto out_size = kernels::detail::compute_size(out_shape);
                output_tmp.push_back(std::make_unique<bfloat16[]>(out_size));

                // lowering transform guarantees shape of gnne_mn_broadcast is {n, c, 1, 1}.
                for (size_t batch = 0; batch < out_shape[0]; batch++)
                {
                    auto in_batch = input_a + batch / (out_shape[0] / in_a_shape[0]) * in_a_shape[1] * in_a_shape[2] * in_a_shape[3];
                    for (size_t oc = 0; oc < out_shape[1]; oc++)
                    {
                        auto in_c = in_batch + oc / (out_shape[1] / in_a_shape[1]) * in_a_shape[2] * in_a_shape[3];
                        for (size_t oh = 0; oh < out_shape[2]; oh++)
                        {
                            auto in_h = in_c + oh / (out_shape[2] / in_a_shape[2]) * in_a_shape[3];
                            for (size_t ow = 0; ow < out_shape[3]; ow++)
                            {
                                output_tmp.back()[output_offset++] = *(in_h + ow / (out_shape[3] / in_a_shape[3]));
                            }
                        }
                    }
                }
            }
        }
        else if (node->id() == meshnet::node_id::select0)
        {
            auto x = const_cast<bfloat16 *>(&constants[0]);
            auto in_shape_0 = output_shape.at(node->input_at(1).connection());
            auto in_shape_1 = output_shape.at(node->input_at(2).connection());
            auto out_shape = [&]() {
                runtime_shape_t ret(4);
                for (auto i = 0U; i < 4; i++)
                    ret[i] = std::max(in_shape_0[i], in_shape_1[i]);
                return ret;
            }();
            auto out_size = kernels::detail::compute_size(out_shape);
            output_tmp.push_back(std::make_unique<bfloat16[]>(out_size));

            auto op = ternary(*node);
            for (auto d0 = 0U; d0 < out_shape[0]; d0++)
            {
                auto index_0_a = d0 * in_shape_0[0] / out_shape[0];
                auto index_0_b = d0 * in_shape_1[0] / out_shape[0];

                for (auto d1 = 0U; d1 < out_shape[1]; d1++)
                {
                    auto index_1_a = d1 * in_shape_0[1] / out_shape[1];
                    auto index_1_b = d1 * in_shape_1[1] / out_shape[1];

                    for (auto d2 = 0U; d2 < out_shape[2]; d2++)
                    {
                        auto index_2_a = d2 * in_shape_0[2] / out_shape[2];
                        auto index_2_b = d2 * in_shape_1[2] / out_shape[2];

                        for (auto d3 = 0U; d3 < out_shape[3]; d3++)
                        {
                            auto index_3_a = d3 * in_shape_0[3] / out_shape[3];
                            auto index_3_b = d3 * in_shape_1[3] / out_shape[3];

                            auto in_a_offset = linear_index(in_shape_0, { index_0_a, index_1_a, index_2_a, index_3_a });
                            auto in_b_offset = linear_index(in_shape_1, { index_0_b, index_1_b, index_2_b, index_3_b });
                            auto out_offset = linear_index(out_shape, { d0, d1, d2, d3 });

                            output_tmp.back()[out_offset] = op(
                                x[0],
                                output_data.at(node->input_at(1).connection())[in_a_offset],
                                output_data.at(node->input_at(2).connection())[in_b_offset]);
                        }
                    }
                }
            }

            output_shape.emplace(&node->output_at(0), out_shape);
            output_data.emplace(&node->output_at(0), output_tmp.back().get());
        }
        else
        {
            if (node->inputs().size() == 1)
            {
                runtime_shape_t out_shape = output_shape.at(node->input_at(0).connection());
                auto out_size = kernels::detail::compute_size(out_shape);
                output_tmp.push_back(std::make_unique<bfloat16[]>(out_size));

                auto op = unary(*node);
                for (auto i = 0U; i < out_size; i++)
                    output_tmp.back()[i] = op(output_data.at(node->input_at(0).connection())[i]);

                output_shape.emplace(&node->output_at(0), out_shape);
                output_data.emplace(&node->output_at(0), output_tmp.back().get());
            }
            else
            {
                auto in_shape_0 = output_shape.at(node->input_at(0).connection());
                auto in_shape_1 = output_shape.at(node->input_at(1).connection());
                auto out_shape = [&]() {
                    runtime_shape_t ret(4);
                    for (auto i = 0U; i < 4; i++)
                        ret[i] = std::max(in_shape_0[i], in_shape_1[i]);
                    return ret;
                }();
                auto out_size = kernels::detail::compute_size(out_shape);
                output_tmp.push_back(std::make_unique<bfloat16[]>(out_size));

                auto op = binary(*node);
                for (auto d0 = 0U; d0 < out_shape[0]; d0++)
                {
                    auto index_0_a = d0 * in_shape_0[0] / out_shape[0];
                    auto index_0_b = d0 * in_shape_1[0] / out_shape[0];

                    for (auto d1 = 0U; d1 < out_shape[1]; d1++)
                    {
                        auto index_1_a = d1 * in_shape_0[1] / out_shape[1];
                        auto index_1_b = d1 * in_shape_1[1] / out_shape[1];

                        for (auto d2 = 0U; d2 < out_shape[2]; d2++)
                        {
                            auto index_2_a = d2 * in_shape_0[2] / out_shape[2];
                            auto index_2_b = d2 * in_shape_1[2] / out_shape[2];

                            for (auto d3 = 0U; d3 < out_shape[3]; d3++)
                            {
                                auto index_3_a = d3 * in_shape_0[3] / out_shape[3];
                                auto index_3_b = d3 * in_shape_1[3] / out_shape[3];

                                auto in_a_offset = linear_index(in_shape_0, { index_0_a, index_1_a, index_2_a, index_3_a });
                                auto in_b_offset = linear_index(in_shape_1, { index_0_b, index_1_b, index_2_b, index_3_b });
                                auto out_offset = linear_index(out_shape, { d0, d1, d2, d3 });
                                output_tmp.back()[out_offset] = op(output_data.at(node->input_at(0).connection())[in_a_offset],
                                    output_data.at(node->input_at(1).connection())[in_b_offset]);
                            }
                        }
                    }
                }

                output_shape.emplace(&node->output_at(0), out_shape);
                output_data.emplace(&node->output_at(0), output_tmp.back().get());
            }
        }
    }

    std::copy(output_tmp.back().get(), output_tmp.back().get() + runtime::compute_size(out_shape), output);
}
#endif
inline void gnne_conv2d_transpose(const bfloat16 *input, bfloat16 *output, const bfloat16 *weights, const bfloat16 *act, const runtime_shape_t &in_shape,
    const int32_t groups, const runtime_shape_t &out_shape, int32_t filter_h, int32_t filter_w, int32_t stride_h, int32_t stride_w, int32_t dilation_h, int32_t dilation_w,
    const padding &padding_h, const padding &padding_w, const value_range<bfloat16> &fused_activation, int32_t output_channels)
{
    auto output_size = kernels::detail::compute_size(out_shape);
    auto out_tmp = std::make_unique<float[]>(output_size);
    std::fill(out_tmp.get(), out_tmp.get() + output_size, 0.f);

    const auto g_ic = in_shape[1] / groups;
    const auto g_oc = out_shape[1] / groups;

    for (size_t batch = 0; batch < in_shape[0]; batch++)
    {
        float *out_batch_p = out_tmp.get() + batch * out_shape[1] * out_shape[2] * out_shape[3];

        for (int32_t g = 0; g < groups; g++)
        {
            float *out_group_p = out_batch_p + g * g_oc * out_shape[2] * out_shape[3];
            const bfloat16 *w_group_p = weights + g * g_oc * g_ic * filter_h * filter_w;

            for (size_t ic = 0; ic < g_ic; ic++)
            {
                for (size_t iy = 0; iy < in_shape[2]; iy++)
                {
                    for (size_t ix = 0; ix < in_shape[3]; ix++)
                    {
                        const int32_t out_y_origin = (iy * stride_h) - padding_h.before;
                        const int32_t out_x_origin = (ix * stride_w) - padding_w.before;
                        const int32_t filter_y_start = std::max(0, (-out_y_origin + dilation_h - 1) / dilation_h);
                        const int32_t filter_y_end = std::min(filter_h, ((int32_t)out_shape[2] - out_y_origin + dilation_h - 1) / dilation_h);
                        const int32_t filter_x_start = std::max(0, (-out_x_origin + dilation_w - 1) / dilation_w);
                        const int32_t filter_x_end = std::min(filter_w, ((int32_t)out_shape[3] - out_x_origin + dilation_w - 1) / dilation_w);
                        const bfloat16 in_v = *input++;

                        for (size_t oc = 0; oc < g_oc; oc++)
                        {
                            float *out_c_p = out_group_p + oc * out_shape[2] * out_shape[3];
                            const bfloat16 *w_oc_p = w_group_p + oc * g_ic * filter_h * filter_w;
                            const bfloat16 *w_ic_p = w_oc_p + ic * filter_h * filter_w;

                            for (int32_t ky = filter_y_start; ky < filter_y_end; ky++)
                            {
                                for (int32_t kx = filter_x_start; kx < filter_x_end; kx++)
                                {
                                    const int32_t out_y = out_y_origin + dilation_h * ky;
                                    const int32_t out_x = out_x_origin + dilation_w * kx;

                                    const bfloat16 w = w_ic_p[ky * filter_w + kx];

                                    out_c_p[out_y * out_shape[3] + out_x] += (float)in_v * w;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    auto channel_size = output_size / out_shape[0] / output_channels;
    for (size_t i = 0; i < output_size; i++)
    {
        auto oc = (i / channel_size) % output_channels;
        output[i] = apply_gnne_activation(out_tmp[i], act[oc * 5], act[oc * 5 + 1], act[oc * 5 + 2], act[oc * 5 + 3], act[oc * 5 + 4]);
        output[i] = detail::apply_activation(output[i], fused_activation);
    }
}

inline void gnne_resize_nearest_neighbor(const bfloat16 *input, bfloat16 *output, const runtime_shape_t &in_shape, int32_t out_h, int32_t out_w)
{
    auto height_scale = (float)in_shape[2] / out_h;
    auto width_scale = (float)in_shape[3] / out_w;

    const auto in_image_size = in_shape[2] * in_shape[3];
    const auto out_image_size = out_h * out_w;
    for (size_t batch = 0; batch < in_shape[0]; batch++)
    {
        auto *begin_input_ptr = input + batch * in_shape[1] * in_image_size;
        auto *begin_output_ptr = output + batch * in_shape[1] * out_image_size;
#ifdef NNCASE_OPENMP
#pragma omp parallel for num_threads(kernels::default_kernel_context().num_threads)
#endif
        for (size_t oc = 0; oc < in_shape[1]; oc++)
        {
            auto *input_ptr = begin_input_ptr + oc * in_image_size;
            auto *output_ptr = begin_output_ptr + oc * out_image_size;

            for (int oy = 0; oy < out_h; oy++)
            {
                auto in_y = std::min((int32_t)floorf(oy * height_scale), (int32_t)in_shape[2] - 1);
                auto *in_row = input_ptr + in_y * in_shape[3];

                for (int ox = 0; ox < out_w; ox++)
                {
                    auto in_x = std::min((int32_t)floorf(ox * width_scale), (int32_t)in_shape[3] - 1);
                    *output_ptr++ = in_row[in_x];
                }
            }
        }
    }
}

inline void gnne_resize_bilinear(const bfloat16 *input, bfloat16 *output, const runtime_shape_t &in_shape, int32_t out_h, int32_t out_w, bool align_corners, bool half_pixel_centers)
{
    auto height_scale = (float)in_shape[2] / out_h;
    auto width_scale = (float)in_shape[3] / out_w;
    if (align_corners && out_h > 1)
        height_scale = (float)(in_shape[2] - 1) / (out_h - 1);
    if (align_corners && out_w > 1)
        width_scale = (float)(in_shape[3] - 1) / (out_w - 1);

    const auto in_img_size = in_shape[2] * in_shape[3];
    const auto out_img_size = out_w * out_h;

    for (size_t batch = 0; batch < in_shape[0]; batch++)
    {
        auto in_batch = input + (size_t)batch * in_shape[1] * in_img_size;
        auto *begin_output_ptr = output + batch * in_shape[1] * out_w * out_h;
#ifdef NNCASE_OPENMP
#pragma omp parallel for num_threads(kernels::default_kernel_context().num_threads)
#endif
        for (size_t oc = 0; oc < in_shape[1]; oc++)
        {
            auto in_c = in_batch + (size_t)oc * in_img_size;
            auto *output_ptr = begin_output_ptr + oc * out_img_size;
            for (int oy = 0; oy < out_h; oy++)
            {
                auto in_y = half_pixel_centers == true ? (oy + 0.5f) * height_scale - 0.5f : oy * height_scale;
                auto in_y0 = std::max((int)in_y, 0);
                auto in_y1 = std::min((int32_t)(std::ceil(in_y)), (int32_t)in_shape[2] - 1);

                for (int ox = 0; ox < out_w; ox++)
                {
                    auto in_x = half_pixel_centers == true ? (ox + 0.5f) * width_scale - 0.5f : ox * width_scale;
                    auto in_x0 = std::max((int)in_x, 0);
                    auto in_x1 = std::min((int32_t)(std::ceil(in_x)), (int32_t)in_shape[3] - 1);

                    auto v0 = in_c[in_y0 * in_shape[3] + in_x0];
                    auto v1 = in_c[in_y1 * in_shape[3] + in_x0];
                    auto v2 = in_c[in_y0 * in_shape[3] + in_x1];
                    auto v3 = in_c[in_y1 * in_shape[3] + in_x1];

                    auto a0 = (1 - (in_y - in_y0)) * (1 - (in_x - in_x0));
                    auto a1 = (in_y - in_y0) * (1 - (in_x - in_x0));
                    auto a2 = (1 - (in_y - in_y0)) * (in_x - in_x0);
                    auto a3 = (in_y - in_y0) * (in_x - in_x0);

                    *output_ptr = bfloat16::round_to_bfloat16(v0 * a0 + v1 * a1 + v2 * a2 + v3 * a3);
                    ++output_ptr;
                }
            }
        }
    }
}

template <class TI, class TO>
void dsp_convert_impl(const TI *input, const runtime_shape_t &in_shape, TO *output)
{
    if ((std::is_same_v<TI, float> || std::is_same_v<TI, bfloat16> || std::is_same_v<TI, int8_t> || std::is_same_v<TI, uint8_t>)&&(std::is_same_v<TO, float> || std::is_same_v<TO, bfloat16> || std::is_same_v<TO, int8_t> || std::is_same_v<TO, uint8_t>))
    {
        for (size_t i = 0; i < detail::compute_size(in_shape); i++)
            output[i] = static_cast<TO>(input[i]);
    }
    else
    {
        std::runtime_error("unsupported convert type in dsp_convert!");
    }
}

#define DSP_CONVERT_IMPL_LV2(input_t, output_t) \
    if (out_type == to_datatype<output_t>())    \
    dsp_convert_impl(reinterpret_cast<const input_t *>(input), in_shape, reinterpret_cast<output_t *>(output))

#define DSP_CONVERT_IMPL_LV1(input_t)            \
    if (in_type == to_datatype<input_t>())       \
    {                                            \
        DSP_CONVERT_IMPL_LV2(input_t, uint8_t);  \
        DSP_CONVERT_IMPL_LV2(input_t, int8_t);   \
        DSP_CONVERT_IMPL_LV2(input_t, bfloat16); \
        DSP_CONVERT_IMPL_LV2(input_t, float);    \
    }

inline void dsp_convert(datatype_t in_type, const runtime_shape_t &in_shape, datatype_t out_type, const gsl::byte *input, gsl::byte *output)
{
    DSP_CONVERT_IMPL_LV1(uint8_t);
    DSP_CONVERT_IMPL_LV1(int8_t);
    DSP_CONVERT_IMPL_LV1(bfloat16);
    DSP_CONVERT_IMPL_LV1(float);
    err(std::errc::not_supported);
}

/**
 * @brief dsp_slice
 * NOTE begin-end must same step with stride, because the each step inc add one strides.
 * @tparam T
 * @param in_shape
 * @param begin
 * @param end
 * @param strides if input begin end with tensor shape size, the strides should be {1,1,1,1}
 * @param input
 * @param output
 */
#ifndef ENABLE_K510_RUNTIME
template <class T>
inline void dsp_slice(const runtime_shape_t &in_shape, const nncase::ir::axis_t &begin, const nncase::ir::axis_t &end, const nncase::ir::axis_t &strides, const T *input, T *output)
{
    auto loop_cond = [](int32_t i, int32_t stop, int32_t stride) {
        return stride > 0 ? i < stop : i > stop;
    };

    for (int32_t d0 = begin[0]; loop_cond(d0, end[0], strides[0]); d0 += strides[0])
    {
        auto d0_origin = input + (size_t)d0 * in_shape[1] * in_shape[2] * in_shape[3];
        for (int d1 = begin[1]; loop_cond(d1, end[1], strides[1]); d1 += strides[1])
        {
            auto d1_origin = d0_origin + (size_t)d1 * in_shape[2] * in_shape[3];
            for (int32_t d2 = begin[2]; loop_cond(d2, end[2], strides[2]); d2 += strides[2])
            {
                auto d2_origin = d1_origin + (size_t)d2 * in_shape[3];
                for (int32_t d3 = begin[3]; loop_cond(d3, end[3], strides[3]); d3 += strides[3])
                {
                    *output++ = d2_origin[d3];
                }
            }
        }
    }
}
#endif

template <class T>
inline void dsp_clamp(const runtime_shape_t &in_shape, const T *input, const T *input_low, const T *input_high, T *output)
{
    auto min = reinterpret_cast<const T *>(input_low);
    auto max = reinterpret_cast<const T *>(input_high);

    size_t size = in_shape[0] * in_shape[1] * in_shape[2] * in_shape[3];
    for (size_t i = 0; i < size; i++)
    {
        *(output + i) = *(input + i) >= *max ? *max : *(input + i);
        *(output + i) = *(input + i) <= *min ? *min : *(input + i);
    }
}

template <class TI>
inline void dsp_dequantize(const runtime_shape_t &param_shape, const TI *input, float *output, size_t size, const gsl::byte *param)
{
    size_t channel = param_shape[2];
    size_t per_channel_count = size / channel;

    for (size_t i = 0; i < size; i++)
    {
        auto scale = reinterpret_cast<int16_t *>(param[(i / per_channel_count) * 4]);
        auto shift = reinterpret_cast<int8_t *>(param[(i / per_channel_count) * 4 + 2]);
        auto bias = reinterpret_cast<TI *>(param[(i / per_channel_count) * 4 + 3]);
        auto value = (input[i] - *bias) * (*scale) / pow(2, *shift);
        output[i] = static_cast<float>(value);
    }
}

template <class TO>
inline void dsp_quantize(const runtime_shape_t &param_shape, const float *input, TO *output, size_t size, const bfloat16 *param)
{
    size_t channel = param_shape[2];
    size_t per_channel_count = size / channel;

    for (size_t i = 0; i < size; i++)
    {
        auto value = input[i] * param[(i / per_channel_count) * 2] + param[(i / per_channel_count) * 2 + 1];
        if (std::is_same_v<TO, int8_t>)
            value = std::clamp(value, -128.f, 127.f);
        else
            value = std::clamp(value, 0.f, 255.f);
        output[i] = static_cast<TO>(value);
    }
}

template <typename TI>
inline void gnne_activation(const TI *input, bfloat16 *output, const bfloat16 *act,
    const runtime_shape_t &in_shape, int32_t out_channels, value_range<bfloat16> fused_clamp)
{
    const auto img_size = in_shape[2] * in_shape[3];
    const auto batch_size = in_shape[1] * img_size;
    for (size_t batch = 0; batch < in_shape[0]; batch++)
    {
        auto *begin_input_ptr = input + batch * batch_size;
        auto *begin_output_ptr = output + batch * out_channels * img_size;
#ifdef NNCASE_OPENMP
#pragma omp parallel for num_threads(kernels::default_kernel_context().num_threads)
#endif
        for (int oc = 0; oc < out_channels; oc++)
        {
            auto *input_ptr = begin_input_ptr + oc * img_size;
            auto *output_ptr = begin_output_ptr + oc * img_size;
            for (size_t i = 0; i < img_size; i++)
            {
                auto value = apply_gnne_activation(*input_ptr, act[oc * 5], act[oc * 5 + 1], act[oc * 5 + 2], act[oc * 5 + 3], act[oc * 5 + 4]);
                *output_ptr = detail::apply_activation(value, fused_clamp);
                ++input_ptr;
                ++output_ptr;
            }
        }
    }
}

/**
 * @brief gnne crop
 *
 * @param input nchw
 * @param output n*m, c, out_h, out_w
 * @param bbox n * [y0x0 y1x1]
 * @param in_shape
 * @param roi_amounts
 * @param out_h
 * @param out_w
 * @param resize_mode
 */
inline void gnne_crop(const bfloat16 *input, bfloat16 *output, bfloat16 *bbox, const runtime_shape_t &in_shape, size_t roi_amounts, int32_t out_h, int32_t out_w, image_resize_mode_t resize_mode, bool align_corners, bool half_pixel_centers)
{
    runtime_shape_t begin { 0, 0, 0, 0 };
    runtime_axis_t end { 0, 0, 0, 0 };
    runtime_shape_t in_strides = runtime::get_default_strides(in_shape);
    runtime_shape_t output_shape { in_shape[0], roi_amounts * in_shape[1], size_t(out_h), size_t(out_w) };

    /* padding zero strides */
    auto fill_empty = [](runtime_shape_t &shape) {
        for (int i = shape.size() - 2; i >= 0; i--)
        {
            shape[i] = (shape[i] == 0) ? (shape[i + 1]) : shape[i];
        }
    };
    fill_empty(in_strides);

    runtime_axis_t steps { 1, 1, 1, 1 };
    size_t out_offset = 0; // offset(out_strides, { batch, roi * ic + ic, 0, 0 });
    for (size_t batch = 0; batch < in_shape[0]; batch++)
    {
        for (size_t roi = 0; roi < roi_amounts; roi++)
        {
            bfloat16 y1 = bbox[roi * 4 + 0];
            bfloat16 x1 = bbox[roi * 4 + 1];
            bfloat16 y2 = bbox[roi * 4 + 2];
            bfloat16 x2 = bbox[roi * 4 + 3];
            begin[2] = (size_t)round(y1 * (in_shape[2] - 1));
            begin[3] = (size_t)round(x1 * (in_shape[3] - 1));
            end[2] = (size_t)round(y2 * (in_shape[2] - 1)) + 1;
            end[3] = (size_t)round(x2 * (in_shape[3] - 1)) + 1;
            for (size_t ic = 0; ic < in_shape[1]; ic++)
            {
                begin[0] = batch;
                end[0] = batch + 1;
                begin[1] = ic;
                end[1] = ic + 1;
                runtime_shape_t tmp_shape { 1, 1, size_t(end[2] - begin[2]), size_t(end[3] - begin[3]) };
                size_t tmp_size = runtime::compute_size(tmp_shape);
                bfloat16 *tmp = new bfloat16[tmp_size];
                runtime_shape_t out_strides = runtime::get_default_strides({ 1UL, 1UL, size_t(end[2] - begin[2]), size_t(end[3] - begin[3]) });
                fill_empty(out_strides);
                slice(dt_bfloat16, reinterpret_cast<const gsl::byte *>(input), reinterpret_cast<gsl::byte *>(tmp), in_shape,
                    in_strides, out_strides, begin, end, steps)
                    .unwrap_or_throw();

                if (resize_mode == image_resize_bilinear)
                {
                    gnne_resize_bilinear(tmp, output + out_offset, tmp_shape, out_h, out_w, align_corners, half_pixel_centers);
                }
                else if (resize_mode == image_resize_nearest_neighbor)
                {
                    gnne_resize_nearest_neighbor(tmp, output + out_offset, tmp_shape, out_h, out_w);
                }

                out_offset += out_h * out_w;
            }
        }
    }
}

inline void gnne_lstm(const bfloat16 *input, const bfloat16 *w_xc, const bfloat16 *act_xc, const bfloat16 *w_rc,
    const bfloat16 *act_rc, const bfloat16 *init_h, const bfloat16 *init_c, [[maybe_unused]] const bfloat16 *seg_fitting_param_ft,
    [[maybe_unused]] const bfloat16 *seg_fitting_param_gt, bfloat16 *output, bfloat16 *output_h, bfloat16 *output_c,
    const runtime_shape_t &in_shape, const runtime_shape_t &init_h_shape, const runtime_shape_t &init_c_shape, const runtime_shape_t &out_shape, const runtime_shape_t &w_xc_shape, const runtime_shape_t &w_rc_shape, lstm_direction direction)
{
    auto tanh = [&](float x) {
        for (size_t c = 0; c < 15; c++)
        {
            if (bfloat16(x) < seg_fitting_param_gt[32 + c])
            {
                return float(seg_fitting_param_gt[2 * c] * x + seg_fitting_param_gt[2 * c + 1]);
            }
        }
        return float(seg_fitting_param_gt[2 * 15] * x + seg_fitting_param_gt[2 * 15 + 1]);
    };
    auto sigmoid = [&](float x) {
        for (size_t c = 0; c < 15; c++)
        {
            if (bfloat16(x) < seg_fitting_param_ft[32 + c])
            {
                return float(seg_fitting_param_ft[2 * c] * x + seg_fitting_param_ft[2 * c + 1]);
            }
        }
        return float(seg_fitting_param_ft[2 * 15] * x + seg_fitting_param_ft[2 * 15 + 1]);
    };

    auto output_h_tmp = std::make_unique<bfloat16[]>(xt::compute_size(init_h_shape));
    auto output_c_tmp = std::make_unique<bfloat16[]>(xt::compute_size(init_c_shape));
    std::memcpy(output_h_tmp.get(), init_h, sizeof(bfloat16) * xt::compute_size(init_h_shape));
    std::memcpy(output_c_tmp.get(), init_c, sizeof(bfloat16) * xt::compute_size(init_c_shape));

    std::vector<uint32_t> seq_len_loop;
    for (uint32_t l = 0; l < in_shape[1]; l++)
        seq_len_loop.push_back(l);
    if (direction == lstm_direction::kReverse)
        std::reverse(seq_len_loop.begin(), seq_len_loop.end());
    for (uint32_t d = 0; d < out_shape[1]; d++)
    {
        if (d == 1)
            std::reverse(seq_len_loop.begin(), seq_len_loop.end());
        for (uint32_t b = 0; b < in_shape[2]; b++)
        {
            for (auto &l : seq_len_loop)
            {
                // g = w_xc_x + w_xc_x
                auto out_mul1 = std::vector<float>(out_shape[3] * 4);
                auto out_mul2 = std::vector<float>(out_shape[3] * 4);
                for (size_t o = 0; o < out_mul1.size(); o++)
                {
                    for (size_t i = 0; i < in_shape[3]; i++)
                    {
                        auto in_idx = i + b * in_shape[3] + l * in_shape[2] * in_shape[3];
                        auto w_idx = i + o * w_xc_shape[3] + d * w_xc_shape[2] * w_xc_shape[3];
                        out_mul1[o] += float(input[in_idx]) * float(w_xc[w_idx]);
                    }
                    auto base = act_xc + d * w_xc_shape[2] * 5;
                    out_mul1[o] = float(apply_gnne_activation(out_mul1[o], base[o * 5], base[o * 5 + 1], base[o * 5 + 2], base[o * 5 + 3], base[o * 5 + 4]));

                    for (size_t i = 0; i < out_shape[3]; i++)
                    {
                        auto in_idx = i + b * out_shape[3] + d * out_shape[2] * out_shape[3];
                        auto w_idx = i + o * w_rc_shape[3] + d * w_rc_shape[2] * w_rc_shape[3];
                        out_mul2[o] += float(output_h_tmp[in_idx]) * float(w_rc[w_idx]);
                    }
                    base = act_rc + d * w_rc_shape[2] * 5;
                    out_mul2[o] = float(apply_gnne_activation(out_mul2[o], base[o * 5], base[o * 5 + 1], base[o * 5 + 2], base[o * 5 + 3], base[o * 5 + 4]));

                    out_mul1[o] += out_mul2[o];
                }

                // ft = sigmoid(g[2])
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    out_mul1[o + out_shape[3] * 2] = sigmoid(out_mul1[o + out_shape[3] * 2]);
                }

                // ct = init_c * ft
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    out_mul1[o + out_shape[3] * 2] = out_mul1[o + out_shape[3] * 2] * output_c_tmp[o + b * out_shape[3] + d * out_shape[2] * out_shape[3]];
                }

                // it = sigmoid(g[0])
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    out_mul1[o + out_shape[3] * 0] = sigmoid(out_mul1[o + out_shape[3] * 0]);
                }

                // c_t = tanh(g[3])
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    out_mul1[o + out_shape[3] * 3] = tanh(out_mul1[o + out_shape[3] * 3]);
                }

                // c_t_it = it * c_t
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    out_mul1[o + out_shape[3] * 0] = out_mul1[o + out_shape[3] * 0] * out_mul1[o + out_shape[3] * 3];
                }

                // ct = ct + c_t_it
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    output_c_tmp[o + d * out_shape[2] * out_shape[3]] = bfloat16(out_mul1[o + out_shape[3] * 2] + out_mul1[o + out_shape[3] * 0]);
                }

                // ot = sigmoid(g[1])
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    out_mul1[o + out_shape[3] * 1] = sigmoid(out_mul1[o + out_shape[3] * 1]);
                }

                // tanh_ct = tanh(ct_o)
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    out_mul1[o + out_shape[3] * 3] = tanh(float(output_c_tmp[o + d * out_shape[2] * out_shape[3]]));
                }

                // ht = ot * tanh_ct
                for (size_t o = 0; o < out_shape[3]; o++)
                {
                    output_h_tmp[o + d * out_shape[2] * out_shape[3]] = bfloat16(out_mul1[o + out_shape[3] * 3] * out_mul1[o + out_shape[3] * 1]);
                }
                std::memcpy(output + b * out_shape[3] + d * out_shape[2] * out_shape[3] + l * out_shape[1] * out_shape[2] * out_shape[3], output_h_tmp.get() + d * out_shape[2] * out_shape[3], sizeof(bfloat16) * out_shape[3]);

                if (l == seq_len_loop.back())
                {
                    std::memcpy(output_h + b * out_shape[3] + d * out_shape[2] * out_shape[3], output_h_tmp.get() + d * out_shape[2] * out_shape[3], sizeof(bfloat16) * out_shape[3]);
                    std::memcpy(output_c + b * out_shape[3] + d * out_shape[2] * out_shape[3], output_c_tmp.get() + d * out_shape[2] * out_shape[3], sizeof(bfloat16) * out_shape[3]);
                }
            }
        }
    }
}

inline void gnne_gru(const bfloat16 *input, const bfloat16 *w, const bfloat16 *r, const bfloat16 *b, bfloat16 *init_h,
    [[maybe_unused]] const bfloat16 *seg_fitting_param_ft, [[maybe_unused]] const bfloat16 *seg_fitting_param_gt, bfloat16 *output, bfloat16 *output_h,
    const runtime_shape_t &in_shape, const runtime_shape_t &w_shape, lstm_direction direction, bool linear_before_reset)
{
    const uint32_t seq_length = in_shape[1];
    const uint32_t batch_size = in_shape[2];
    const uint32_t input_size = in_shape[3];
    const uint32_t num_direction = w_shape[1];
    const uint32_t hidden_size = w_shape[2] / 3;

    auto tanh = [&](float x) {
        for (size_t c = 0; c < 15; c++)
        {
            if (bfloat16(x) < seg_fitting_param_gt[32 + c])
            {
                return float(seg_fitting_param_gt[2 * c] * x + seg_fitting_param_gt[2 * c + 1]);
            }
        }
        return float(seg_fitting_param_gt[2 * 15] * x + seg_fitting_param_gt[2 * 15 + 1]);
    };
    auto sigmoid = [&](float x) {
        for (size_t c = 0; c < 15; c++)
        {
            if (bfloat16(x) < seg_fitting_param_ft[32 + c])
            {
                return float(seg_fitting_param_ft[2 * c] * x + seg_fitting_param_ft[2 * c + 1]);
            }
        }
        return float(seg_fitting_param_ft[2 * 15] * x + seg_fitting_param_ft[2 * 15 + 1]);
    };

    auto x_gate_size = batch_size * input_size;
    auto w_gate_size = 3 * hidden_size * input_size;
    auto h_t_size = batch_size * hidden_size;
    auto r_gate_size = 3 * hidden_size * hidden_size;

    auto tmp_a = std::vector<float>(batch_size * hidden_size, 0.f);
    auto tmp_b = std::vector<float>(batch_size * hidden_size, 0.f);
    auto gate_z = std::vector<float>(batch_size * hidden_size, 0.f);
    auto gate_r = std::vector<float>(batch_size * hidden_size, 0.f);
    auto gate_h = std::vector<float>(batch_size * hidden_size, 0.f);

    auto x_i = input;
    auto h_t = init_h;
    auto w_i = w;
    auto r_i = r;
    auto b_i = b;

    std::vector<uint32_t> seq_len_loop;
    for (uint32_t l = 0; l < seq_length; l++)
        seq_len_loop.push_back(l);
    if (direction == lstm_direction::kReverse)
        std::reverse(seq_len_loop.begin(), seq_len_loop.end());

    for (uint32_t d = 0; d < num_direction; d++)
    {
        h_t = init_h + d * h_t_size;
        w_i = w + d * w_gate_size;
        r_i = r + d * r_gate_size;
        b_i = b + d * 6 * hidden_size * 5; //act 5 params
        if (d == 1)
            std::reverse(seq_len_loop.begin(), seq_len_loop.end());
        for (auto i : seq_len_loop)
        {
            x_i = input + i * x_gate_size;
            // clean gate_z gate_r gate_h
            std::fill(gate_z.begin(), gate_z.end(), 0.f);
            std::fill(gate_r.begin(), gate_r.end(), 0.f);
            std::fill(gate_h.begin(), gate_h.end(), 0.f);

            // clean tmp_a tmp_b
            std::fill(tmp_a.begin(), tmp_a.end(), 0.f);
            std::fill(tmp_b.begin(), tmp_b.end(), 0.f);
            // gate_z = x_i * w_i_z + b_w_z + h_t *r_i_z + b_r_z
            for (uint32_t bs = 0; bs < batch_size; bs++)
            {
                for (uint32_t hs = 0; hs < hidden_size; hs++)
                {
                    for (uint32_t is = 0; is < input_size; is++)
                    {
                        tmp_a[bs * hidden_size + hs] += x_i[bs * input_size + is] * w_i[hs * input_size + is];
                    }

                    tmp_a[bs * hidden_size + hs] = float(apply_gnne_activation(tmp_a[bs * hidden_size + hs], b_i[hs * 5],
                        b_i[hs * 5 + 1], b_i[hs * 5 + 2], b_i[hs * 5 + 3], b_i[hs * 5 + 4]));

                    for (uint32_t rs = 0; rs < hidden_size; rs++)
                    {
                        tmp_b[bs * hidden_size + hs] += h_t[bs * hidden_size + rs] * r_i[hs * hidden_size + rs];
                    }

                    tmp_b[bs * hidden_size + hs] = float(apply_gnne_activation(tmp_b[bs * hidden_size + hs], b_i[(3 * hidden_size + hs) * 5],
                        b_i[(3 * hidden_size + hs) * 5 + 1], b_i[(3 * hidden_size + hs) * 5 + 2], b_i[(3 * hidden_size + hs) * 5 + 3], b_i[(3 * hidden_size + hs) * 5 + 4]));
                    gate_z[bs * hidden_size + hs] = tmp_a[bs * hidden_size + hs] + tmp_b[bs * hidden_size + hs];
                }
            }
            // gate_z = sigmoid(gate_z);
            std::transform(gate_z.begin(), gate_z.end(), gate_z.begin(), sigmoid);

            // clear tmp_a tmp_b
            std::fill(tmp_a.begin(), tmp_a.end(), 0.f);
            std::fill(tmp_b.begin(), tmp_b.end(), 0.f);
            // gate_r = x_i * w_i_r + b_w_r + h_t *r_i_r + b_r_r
            for (uint32_t bs = 0; bs < batch_size; bs++)
            {
                for (uint32_t hs = 0; hs < hidden_size; hs++)
                {
                    for (uint32_t is = 0; is < input_size; is++)
                    {
                        tmp_a[bs * hidden_size + hs] += x_i[bs * input_size + is] * w_i[hidden_size * input_size + hs * input_size + is];
                    }

                    tmp_a[bs * hidden_size + hs] = float(apply_gnne_activation(tmp_a[bs * hidden_size + hs], b_i[(hidden_size + hs) * 5],
                        b_i[(hidden_size + hs) * 5 + 1], b_i[(hidden_size + hs) * 5 + 2], b_i[(hidden_size + hs) * 5 + 3], b_i[(hidden_size + hs) * 5 + 4]));

                    for (uint32_t rs = 0; rs < hidden_size; rs++)
                    {
                        tmp_b[bs * hidden_size + hs] += h_t[bs * hidden_size + rs] * r_i[hidden_size * hidden_size + hs * hidden_size + rs];
                    }

                    tmp_b[bs * hidden_size + hs] = float(apply_gnne_activation(tmp_b[bs * hidden_size + hs], b_i[(4 * hidden_size + hs) * 5],
                        b_i[(4 * hidden_size + hs) * 5 + 1], b_i[(4 * hidden_size + hs) * 5 + 2], b_i[(4 * hidden_size + hs) * 5 + 3], b_i[(4 * hidden_size + hs) * 5 + 4]));
                    gate_r[bs * hidden_size + hs] = tmp_a[bs * hidden_size + hs] + tmp_b[bs * hidden_size + hs];
                }
            }
            // gate_r = sigmoid(gate_r);
            std::transform(gate_r.begin(), gate_r.end(), gate_r.begin(), sigmoid);

            // clear tmp_a tmp_b
            std::fill(tmp_a.begin(), tmp_a.end(), 0.f);
            std::fill(tmp_b.begin(), tmp_b.end(), 0.f);
            // gate_h = x_i * w_i_h + b_w_h + gate_r·h_t *r_i_h + b_r_h
            for (uint32_t bs = 0; bs < batch_size; bs++)
            {
                for (uint32_t hs = 0; hs < hidden_size; hs++)
                {
                    for (uint32_t is = 0; is < input_size; is++)
                    {
                        tmp_a[bs * hidden_size + hs] += x_i[bs * input_size + is] * w_i[2 * hidden_size * input_size + hs * input_size + is];
                    }

                    tmp_a[bs * hidden_size + hs] = float(apply_gnne_activation(tmp_a[bs * hidden_size + hs], b_i[(2 * hidden_size + hs) * 5],
                        b_i[(2 * hidden_size + hs) * 5 + 1], b_i[(2 * hidden_size + hs) * 5 + 2], b_i[(2 * hidden_size + hs) * 5 + 3], b_i[(2 * hidden_size + hs) * 5 + 4]));

                    for (uint32_t rs = 0; rs < hidden_size; rs++)
                    {
                        if (!linear_before_reset)
                            tmp_b[bs * hidden_size + hs] += gate_r[bs * hidden_size + rs] * h_t[bs * hidden_size + rs] * r_i[2 * hidden_size * hidden_size + hs * hidden_size + rs];
                        else
                            tmp_b[bs * hidden_size + hs] += h_t[bs * hidden_size + rs] * r_i[2 * hidden_size * hidden_size + hs * hidden_size + rs];
                    }

                    tmp_b[bs * hidden_size + hs] = float(apply_gnne_activation(tmp_b[bs * hidden_size + hs], b_i[(5 * hidden_size + hs) * 5],
                        b_i[(5 * hidden_size + hs) * 5 + 1], b_i[(5 * hidden_size + hs) * 5 + 2], b_i[(5 * hidden_size + hs) * 5 + 3], b_i[(5 * hidden_size + hs) * 5 + 4]));

                    if (!linear_before_reset)
                        gate_h[bs * hidden_size + hs] = tmp_a[bs * hidden_size + hs] + tmp_b[bs * hidden_size + hs];
                    else
                        gate_h[bs * hidden_size + hs] = tmp_a[bs * hidden_size + hs] + gate_r[bs * hidden_size + hs] * tmp_b[bs * hidden_size + hs];
                }
            }
            // gate_h = tanh(gate_h);
            std::transform(gate_h.begin(), gate_h.end(), gate_h.begin(), tanh);

            for (uint32_t k = 0; k < batch_size * hidden_size; k++)
            {
                h_t[k] = bfloat16((1 - gate_z[k]) * gate_h[k] + gate_z[k] * float(h_t[k]));
                output[i * (num_direction * batch_size * hidden_size) + d * (batch_size * hidden_size) + k] = h_t[k];
            }
        }
        for (uint32_t k = 0; k < batch_size * hidden_size; k++)
        {
            output_h[d * (batch_size * hidden_size) + k] = h_t[k];
        }
    }
}
}