// File: kernel.cpp
#include "gemm_w8a8ch_contract.h"
#include "innerproduct_arm.h"
#include "mat.h"
#include "option.h"
#include <cstring>

namespace {
constexpr int N = 1000;
constexpr int K = 1280;
// input_scale is constant across every workload for this definition (checked
// at generation time), so it's baked here rather than plumbed through the ABI.
constexpr float input_scale = 0.01;
} // namespace

int ncnn::gemm_w8a8ch_kernel(
    const Mat& A, Mat& top,
    const Mat& B, const Mat& weight_scales,
    const Option& opt)
{
    // ncnn's int8 scale fields are *quantization* multipliers (int8 =
    // round(float*scale)); this baseline's weight_scales/input_scale are
    // *dequantization* multipliers (real = int8*scale) per the reference —
    // invert them, or the output is silently wrong (not a crash).
    Mat weight_int8_scales;
    weight_int8_scales.create(N, (size_t)4u);
    if (weight_int8_scales.empty()) return -1;
    const float* ws = (const float*)weight_scales.data;
    for (int p = 0; p < N; ++p)
        weight_int8_scales[p] = 1.0f / ws[p];

    Mat bottom_int8_scales;
    bottom_int8_scales.create(1, (size_t)4u);
    if (bottom_int8_scales.empty()) return -1;
    bottom_int8_scales[0] = 1.0f / input_scale;

    InnerProduct_arm ip;
    ip.num_output       = N;
    ip.bias_term        = 0;
    ip.weight_data_size = N * K;
    // Plain truthy gate here (no >100 distinction like Convolution) — InnerProduct_arm
    // never produces int8 output itself regardless of this value, see below.
    ip.int8_scale_term  = 1;
    ip.activation_type  = 0;
    ip.activation_params = Mat();
    // B is flat (N*K), row-major with K contiguous — exactly the layout
    // InnerProduct_arm expects (reshape(K, N) done internally in create_pipeline).
    ip.weight_data             = const_cast<Mat&>(B);
    ip.weight_data_int8_scales = weight_int8_scales;
    ip.bottom_blob_int8_scales = bottom_int8_scales;

    if (ip.create_pipeline(opt) != 0) return -1;

    // A is a genuine 2D Mat (w=K, h=M) — InnerProduct_arm handles the full
    // (M,K)->(M,N) gemm in one call when dims==2, no manual row loop needed.
    // InnerProduct_arm's int8 path only quantizes the *input*; it always
    // dequantizes the int32 accumulator back to float32 (never calls
    // Requantize/float2int8), which is exactly the plain dequantized result
    // this task wants — just copy it into the pre-allocated output Mat.
    Mat local_top_fp32;
    if (ip.forward(A, local_top_fp32, opt) != 0) return -1;

    const int M = top.h;
    for (int m = 0; m < M; ++m)
        std::memcpy((float*)top.row(m), local_top_fp32.row(m), N * sizeof(float));
    return 0;
}