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

namespace {
constexpr int64_t kNHeads  = 16;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kKvHeads = 16;
} // namespace

extern "C" {
// inputs: [0]=Q bf16 [M,n_heads,head_dim], [1]=K bf16 [S,kv_heads,head_dim],
// [2]=V bf16 [S,kv_heads,head_dim]; var_axes: [0]=M, [1]=S
int armbench_entry_mha(const void* const* inputs, void* output,
                       const int64_t* var_axes)
{
    return armbench_llamacpp_mha(
        inputs[0],
        inputs[1],
        inputs[2],
        reinterpret_cast<float*>(output),
        var_axes[0], var_axes[1], kNHeads, kHeadDim, kKvHeads);
}
} // extern "C"