// 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" {
// inputs: [0]=hidden_states bf16 [T,n_embd], [1]=router_weight bf16
// [n_expert,n_embd], [2]=gate_proj bf16 [n_expert,n_ff,n_embd], [3]=up_proj
// bf16 [n_expert,n_ff,n_embd], [4]=down_proj bf16 [n_expert,n_embd,n_ff];
// 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],
        inputs[1],
        inputs[2], inputs[3], inputs[4],
        reinterpret_cast<float*>(output),
        var_axes[0], kNEmbd, kNFf, kNExpert, kNExpertUsed, /*is_q8=*/false);
}
} // extern "C"