// File: binding.cpp
#include "moe.h"

namespace {
constexpr int64_t kNEmbd       = 2048;
constexpr int64_t kNFf         = 1024;
constexpr int64_t kNExpert     = 64;
constexpr int64_t kNExpertUsed = 8;
} // namespace

extern "C" {
// [0]=hidden_states bf16, [1]=router_weight fp32,
// [2]=gate_proj block_q4_K, [3]=NULL (gate_scales), [4]=NULL (gate_mins),
// [5]=up_proj block_q4_K, [6]=NULL (up_scales), [7]=NULL (up_mins),
// [8]=down_proj block_q4_K, [9]=NULL (down_scales), [10]=NULL (down_mins);
// var_axes: [0]=n_tokens
int armbench_entry_moe(const void* const* inputs, void* output,
                       const int64_t* var_axes)
{
    return armbench_llamacpp_moe(
        inputs[0],
        reinterpret_cast<const float*>(inputs[1]),
        inputs[2], inputs[5], inputs[8],
        reinterpret_cast<float*>(output),
        var_axes[0], kNEmbd, kNFf, kNExpert, kNExpertUsed);
}
} // extern "C"