// Copyright 2026 Tencent
// SPDX-License-Identifier: BSD-3-Clause

#ifdef _MSC_VER
#define _CRT_SECURE_NO_DEPRECATE
#endif

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <float.h>
#include <string>
#include <vector>

// npy format header
#include "npy.hpp"

// ncnn public header
#include "datareader.h"
#include "layer.h"
#include "layer_type.h"
#include "net.h"

// ncnn private header
#include "../modelwriter.h"
#include "ncnnllm_quant.h"

class DataReaderFromEmpty : public ncnn::DataReader
{
public:
    virtual int scan(const char* /*format*/, void* /*p*/) const
    {
        return 0;
    }
    virtual size_t read(void* buf, size_t size) const
    {
        memset(buf, 0, size);
        return size;
    }
};

class QuantActStat
{
public:
    QuantActStat()
    {
        width = 0;
        count = 0;
        max_sample_count = 0;
    }

public:
    int width;
    int count;
    int max_sample_count;
    ncnn::Mat sum_abs;
    std::vector<float> samples;
};

class QuantNet : public ModelWriter
{
public:
    QuantNet();

public:
    std::vector<std::vector<std::string> > listspaths;
    std::vector<std::vector<int> > shapes;
    std::vector<int> input_blobs;
    std::vector<QuantActStat> gemm_act_stats;
    std::vector<QuantActStat> mha_act_stats;
    int quantize_num_threads;
    int file_type;
    int awq_steps;
    int awq_samples;
    int awq_inner_method;
    float awq_max_scale;
    int gptq_samples;
    float gptq_damp;
    bool use_calibration_dataset;

public:
    int init(int method);
    int collect_activation_stats();
    int save_table(const char* tablepath, int block_size, int weight_bits, int method) const;
};

QuantNet::QuantNet()
    : ModelWriter()
{
    quantize_num_threads = 1;
    file_type = 1;
    awq_steps = 20;
    awq_samples = 128;
    awq_inner_method = LLM_QUANT_METHOD_MINMAX;
    awq_max_scale = 16.f;
    gptq_samples = 128;
    gptq_damp = 0.01f;
    use_calibration_dataset = false;
}

static void show_usage(const char* argv0)
{
    fprintf(stderr, "Usage: %s [inparam] [inbin] [outtable] [(key=value)...]\n", argv0);
    fprintf(stderr, "       %s [inparam] [inbin] [caliblist] [outtable] [(key=value)...]\n", argv0);
    fprintf(stderr, "  method=minmax/mseclip/awq/gptq\n");
    fprintf(stderr, "  bits=4/6/8\n");
    fprintf(stderr, "  block=32/64/128\n");
    fprintf(stderr, "  thread=8\n");
    fprintf(stderr, "  type=1\n");
    fprintf(stderr, "  shape=[w,h,...]\n");
    fprintf(stderr, "Sample usage:\n");
    fprintf(stderr, "  %s model.param model.bin model.llm.table method=mseclip bits=4 block=64\n", argv0);
    fprintf(stderr, "  %s model.param model.bin calib.list model.llm.table method=awq bits=4 block=64 shape=[4096]\n", argv0);
    fprintf(stderr, "  %s model.param model.bin calib.list model.llm.table method=gptq bits=4 block=64 shape=[4096]\n", argv0);
}

static ncnn::Mat read_npy(const std::vector<int>& shape, const std::string& npypath)
{
    npy::npy_data<float> d;
    try
    {
        d = npy::read_npy<float>(npypath);
    }
    catch (const std::exception& e)
    {
        fprintf(stderr, "npy::read_npy %s exception: %s\n", npypath.c_str(), e.what());
        return ncnn::Mat();
    }

    const std::vector<unsigned long>& npy_shape = d.shape;
    const size_t dims = shape.size();

    if (dims != npy_shape.size())
    {
        fprintf(stderr, "%s expect %d dims, but got %d\n", npypath.c_str(), (int)dims, (int)npy_shape.size());
        return ncnn::Mat();
    }

    for (size_t i = 0; i < dims; i++)
    {
        if ((unsigned long)shape[i] != npy_shape[dims - 1 - i])
        {
            fprintf(stderr, "%s shape mismatch\n", npypath.c_str());
            return ncnn::Mat();
        }
    }

    if (dims == 1)
        return ncnn::Mat(shape[0], (void*)d.data.data()).reshape(shape[0]).clone();
    if (dims == 2)
        return ncnn::Mat(shape[0] * shape[1], (void*)d.data.data()).reshape(shape[0], shape[1]).clone();
    if (dims == 3)
        return ncnn::Mat(shape[0] * shape[1] * shape[2], (void*)d.data.data()).reshape(shape[0], shape[1], shape[2]).clone();
    if (dims == 4)
        return ncnn::Mat(shape[0] * shape[1] * shape[2] * shape[3], (void*)d.data.data()).reshape(shape[0], shape[1], shape[2], shape[3]).clone();

    fprintf(stderr, "%s dims %d is unsupported\n", npypath.c_str(), (int)dims);
    return ncnn::Mat();
}

static std::vector<std::vector<std::string> > parse_comma_path_list(char* s)
{
    std::vector<std::vector<std::string> > aps;

    char* pch = strtok(s, ",");
    while (pch != NULL)
    {
        FILE* fp = fopen(pch, "rb");
        if (!fp)
        {
            fprintf(stderr, "fopen %s failed\n", pch);
            break;
        }

        std::vector<std::string> paths;

        char line[1024];
        while (!feof(fp))
        {
            char* ss = fgets(line, 1024, fp);
            if (!ss)
                break;

            char filepath[256];
            int nscan = sscanf(line, "%255s", filepath);
            if (nscan != 1)
                continue;

            paths.push_back(std::string(filepath));
        }

        fclose(fp);

        aps.push_back(paths);

        pch = strtok(NULL, ",");
    }

    return aps;
}

static std::vector<std::vector<int> > parse_comma_int_array_list(char* s)
{
    std::vector<std::vector<int> > aai;

    char* pch = strtok(s, "[]");
    while (pch != NULL)
    {
        int v;
        int nconsumed = 0;
        int nscan = sscanf(pch, "%d%n", &v, &nconsumed);
        if (nscan == 1)
        {
            pch += nconsumed;

            std::vector<int> ai;
            ai.push_back(v);

            nscan = sscanf(pch, ",%d%n", &v, &nconsumed);
            while (nscan == 1)
            {
                pch += nconsumed;

                ai.push_back(v);

                nscan = sscanf(pch, ",%d%n", &v, &nconsumed);
            }

            aai.push_back(ai);
        }

        pch = strtok(NULL, "[]");
    }

    return aai;
}

static std::string make_qweight_filename(const char* tablepath, const char* key, std::string& qweight_path)
{
    const char* slash = strrchr(tablepath, '/');
#ifdef _WIN32
    const char* backslash = strrchr(tablepath, '\\');
    if (backslash && (!slash || backslash > slash))
        slash = backslash;
#endif

    const char* table_base = slash ? slash + 1 : tablepath;
    std::string qweight_name = std::string(table_base) + "." + key + ".qweight";
    for (size_t i = 0; i < qweight_name.size(); i++)
    {
        if (qweight_name[i] == '/' || qweight_name[i] == '\\' || qweight_name[i] == ':' || qweight_name[i] == ' ')
            qweight_name[i] = '_';
    }

    qweight_path = qweight_name;
    if (slash)
        qweight_path = std::string(tablepath, slash - tablepath + 1) + qweight_name;

    return qweight_name;
}

static int write_raw_mat_file(const char* path, const ncnn::Mat& m)
{
    FILE* fp = fopen(path, "wb");
    if (!fp)
    {
        fprintf(stderr, "fopen %s failed\n", path);
        return -1;
    }

    const size_t logical_count = (size_t)m.w * m.h * m.d * m.c;
    if (fwrite(m.data, m.elemsize, logical_count, fp) != logical_count)
    {
        fprintf(stderr, "fwrite %s failed\n", path);
        fclose(fp);
        return -1;
    }

    fclose(fp);

    return 0;
}

static int init_act_stat(QuantActStat& stat, int width, int max_sample_count)
{
    stat.width = width;
    stat.count = 0;
    stat.max_sample_count = max_sample_count;

    stat.sum_abs.create(width);
    if (stat.sum_abs.empty())
        return -100;

    stat.sum_abs.fill(0.f);
    stat.samples.clear();
    stat.samples.reserve((size_t)width * max_sample_count);

    return 0;
}

static void collect_act_row(QuantActStat& stat, const float* ptr)
{
    const int width = stat.width;
    float* sum_abs_ptr = stat.sum_abs;

    for (int k = 0; k < width; k++)
    {
        const float v = (float)fabs(ptr[k]);
        sum_abs_ptr[k] += v;
    }

    if ((int)(stat.samples.size() / width) < stat.max_sample_count)
    {
        for (int k = 0; k < width; k++)
            stat.samples.push_back(ptr[k]);
    }

    stat.count++;
}

static void collect_act_rows(QuantActStat& stat, const ncnn::Mat& m)
{
    if (m.dims == 1)
    {
        collect_act_row(stat, (const float*)m);
        return;
    }

    if (m.dims == 2)
    {
        for (int y = 0; y < m.h; y++)
            collect_act_row(stat, m.row(y));
        return;
    }

    if (m.dims == 3)
    {
        for (int q = 0; q < m.c; q++)
        {
            const ncnn::Mat m0 = m.channel(q);
            for (int y = 0; y < m.h; y++)
            {
                collect_act_row(stat, m0.row(y));
            }
        }
        return;
    }

    if (m.dims == 4)
    {
        for (int q = 0; q < m.c; q++)
        {
            const ncnn::Mat m0 = m.channel(q);
            for (int z = 0; z < m.d; z++)
            {
                const ncnn::Mat m1 = m0.depth(z);
                for (int y = 0; y < m.h; y++)
                {
                    collect_act_row(stat, m1.row(y));
                }
            }
        }
        return;
    }
}

static void resolve_mha_bottom_blob_index(const ncnn::MultiHeadAttention* mha, int bottom_blob_count, int& q_blob_i, int& k_blob_i, int& v_blob_i, int& attn_mask_i, int& cached_xk_i, int& cached_xv_i)
{
    if (mha->kv_cache)
    {
        if (mha->attn_mask)
        {
            // assert bottom_blob_count == 4/5/6
            if (bottom_blob_count == 4)
            {
                q_blob_i = 0;
                k_blob_i = 0;
                v_blob_i = 0;
                attn_mask_i = 1;
                cached_xk_i = 2;
                cached_xv_i = 3;
            }
            if (bottom_blob_count == 5)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 1;
                attn_mask_i = 2;
                cached_xk_i = 3;
                cached_xv_i = 4;
            }
            if (bottom_blob_count == 6)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 2;
                attn_mask_i = 3;
                cached_xk_i = 4;
                cached_xv_i = 5;
            }
        }
        else
        {
            // assert bottom_blob_count == 3/4/5
            if (bottom_blob_count == 3)
            {
                q_blob_i = 0;
                k_blob_i = 0;
                v_blob_i = 0;
                cached_xk_i = 1;
                cached_xv_i = 2;
            }
            if (bottom_blob_count == 4)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 1;
                cached_xk_i = 2;
                cached_xv_i = 3;
            }
            if (bottom_blob_count == 5)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 2;
                cached_xk_i = 3;
                cached_xv_i = 4;
            }
        }
    }
    else
    {
        if (mha->attn_mask)
        {
            // assert bottom_blob_count == 2/3/4
            if (bottom_blob_count == 2)
            {
                q_blob_i = 0;
                k_blob_i = 0;
                v_blob_i = 0;
                attn_mask_i = 1;
            }
            if (bottom_blob_count == 3)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 1;
                attn_mask_i = 2;
            }
            if (bottom_blob_count == 4)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 2;
                attn_mask_i = 3;
            }
        }
        else
        {
            // assert bottom_blob_count == 1/2/3
            if (bottom_blob_count == 1)
            {
                q_blob_i = 0;
                k_blob_i = 0;
                v_blob_i = 0;
            }
            if (bottom_blob_count == 2)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 1;
            }
            if (bottom_blob_count == 3)
            {
                q_blob_i = 0;
                k_blob_i = 1;
                v_blob_i = 2;
            }
        }
    }
}

static int collect_mha_out_act_rows(const ncnn::MultiHeadAttention* mha, const ncnn::Mat& q_blob, const ncnn::Mat& k_blob, const ncnn::Mat& v_blob, const ncnn::Mat& attn_mask_blob, const ncnn::Mat& cached_xk_blob, const ncnn::Mat& cached_xv_blob, int q_blob_i, int k_blob_i, int v_blob_i, QuantActStat& stat, int num_threads)
{
    const int embed_dim = mha->embed_dim;
    const int num_heads = mha->num_heads;
    const int embed_dim_per_head = embed_dim / num_heads;
    const int qdim = mha->weight_data_size / embed_dim;

    const int src_seqlen = q_blob.h;
    const int cur_seqlen = k_blob.h;
    const int past_seqlen = mha->kv_cache && !cached_xk_blob.empty() ? cached_xk_blob.w : 0;
    const int dst_seqlen = past_seqlen > 0 ? (q_blob_i == k_blob_i ? (past_seqlen + cur_seqlen) : past_seqlen) : cur_seqlen;

    ncnn::Mat q_affine;
    q_affine.create(src_seqlen, embed_dim);
    if (q_affine.empty())
        return -100;

    #pragma omp parallel for num_threads(num_threads)
    for (int i = 0; i < src_seqlen; i++)
    {
        const float* kptr = (const float*)mha->q_weight_data;

        for (int j = 0; j < embed_dim; j++)
        {
            const float* ptr = q_blob.row(i);

            float sum = mha->q_bias_data[j];
            for (int k = 0; k < qdim; k++)
                sum += *ptr++ * *kptr++;

            q_affine.row(j)[i] = sum * mha->scale;
        }
    }

    ncnn::Mat k_affine;
    if (past_seqlen > 0 && q_blob_i != k_blob_i)
    {
        k_affine = cached_xk_blob;
    }
    else
    {
        k_affine.create(dst_seqlen, embed_dim);
        if (k_affine.empty())
            return -100;

        if (past_seqlen > 0)
        {
            #pragma omp parallel for num_threads(num_threads)
            for (int i = 0; i < embed_dim; i++)
                memcpy(k_affine.row(i), cached_xk_blob.row(i), past_seqlen * sizeof(float));
        }

        #pragma omp parallel for num_threads(num_threads)
        for (int i = 0; i < cur_seqlen; i++)
        {
            const float* kptr = (const float*)mha->k_weight_data;

            for (int j = 0; j < embed_dim; j++)
            {
                const float* ptr = k_blob.row(i);

                float sum = mha->k_bias_data[j];
                for (int k = 0; k < mha->kdim; k++)
                    sum += *ptr++ * *kptr++;

                k_affine.row(j)[past_seqlen + i] = sum;
            }
        }
    }

    ncnn::Mat v_affine;
    if (past_seqlen > 0 && q_blob_i != v_blob_i)
    {
        v_affine = cached_xv_blob;
    }
    else
    {
        v_affine.create(dst_seqlen, embed_dim);
        if (v_affine.empty())
            return -100;

        if (past_seqlen > 0)
        {
            #pragma omp parallel for num_threads(num_threads)
            for (int i = 0; i < embed_dim; i++)
                memcpy(v_affine.row(i), cached_xv_blob.row(i), past_seqlen * sizeof(float));
        }

        #pragma omp parallel for num_threads(num_threads)
        for (int i = 0; i < cur_seqlen; i++)
        {
            const float* kptr = (const float*)mha->v_weight_data;

            for (int j = 0; j < embed_dim; j++)
            {
                const float* ptr = v_blob.row(i);

                float sum = mha->v_bias_data[j];
                for (int k = 0; k < mha->vdim; k++)
                    sum += *ptr++ * *kptr++;

                v_affine.row(j)[past_seqlen + i] = sum;
            }
        }
    }

    ncnn::Mat qk_cross;
    qk_cross.create(dst_seqlen, src_seqlen, num_heads);
    if (qk_cross.empty())
        return -100;

    #pragma omp parallel for num_threads(num_threads)
    for (int q = 0; q < num_heads; q++)
    {
        const ncnn::Mat q_affine_head = q_affine.row_range(q * embed_dim_per_head, embed_dim_per_head);
        const ncnn::Mat k_affine_head = k_affine.row_range(q * embed_dim_per_head, embed_dim_per_head);
        ncnn::Mat qk_cross_head = qk_cross.channel(q);

        for (int i = 0; i < src_seqlen; i++)
        {
            float* outptr = qk_cross_head.row(i);

            for (int j = 0; j < dst_seqlen; j++)
            {
                float sum = 0.f;
                for (int l = 0; l < embed_dim_per_head; l++)
                    sum += q_affine_head.row(l)[i] * k_affine_head.row(l)[j];

                outptr[j] = sum;
            }
        }
    }

    if (mha->attn_mask)
    {
        #pragma omp parallel for num_threads(num_threads)
        for (int q = 0; q < num_heads; q++)
        {
            const ncnn::Mat& maskm = attn_mask_blob.dims == 3 ? attn_mask_blob.channel(q) : attn_mask_blob;
            ncnn::Mat qk_cross_head = qk_cross.channel(q);

            for (int i = 0; i < src_seqlen; i++)
            {
                const float* mptr = maskm.row(i);
                float* outptr = qk_cross_head.row(i);

                for (int j = 0; j < dst_seqlen; j++)
                    outptr[j] += mptr[j];
            }
        }
    }

    #pragma omp parallel for num_threads(num_threads)
    for (int q = 0; q < num_heads; q++)
    {
        ncnn::Mat qk_cross_head = qk_cross.channel(q);

        for (int i = 0; i < src_seqlen; i++)
        {
            float* ptr = qk_cross_head.row(i);

            float max = -FLT_MAX;
            for (int j = 0; j < dst_seqlen; j++)
                max = std::max(max, ptr[j]);

            float sum = 0.f;
            for (int j = 0; j < dst_seqlen; j++)
            {
                ptr[j] = (float)expf(ptr[j] - max);
                sum += ptr[j];
            }

            for (int j = 0; j < dst_seqlen; j++)
                ptr[j] /= sum;
        }
    }

    ncnn::Mat qkv_cross;
    qkv_cross.create(src_seqlen, embed_dim);
    if (qkv_cross.empty())
        return -100;

    #pragma omp parallel for num_threads(num_threads)
    for (int q = 0; q < num_heads; q++)
    {
        const ncnn::Mat qk_cross_head = qk_cross.channel(q);
        const ncnn::Mat v_affine_head = v_affine.row_range(q * embed_dim_per_head, embed_dim_per_head);
        ncnn::Mat qkv_cross_head = qkv_cross.row_range(q * embed_dim_per_head, embed_dim_per_head);

        for (int i = 0; i < src_seqlen; i++)
        {
            for (int j = 0; j < embed_dim_per_head; j++)
            {
                const float* qkptr = qk_cross_head.row(i);
                const float* vptr = v_affine_head.row(j);

                float sum = 0.f;
                for (int k = 0; k < dst_seqlen; k++)
                    sum += *qkptr++ * *vptr++;

                qkv_cross_head.row(j)[i] = sum;
            }
        }
    }

    std::vector<float> row(embed_dim);
    for (int i = 0; i < src_seqlen; i++)
    {
        for (int k = 0; k < embed_dim; k++)
            row[k] = qkv_cross.row(k)[i];

        collect_act_row(stat, row.data());
    }

    return 0;
}

static int make_input_scaled_weight(const ncnn::Mat& weight_data, const ncnn::Mat& input_scales, ncnn::Mat& weight_data_scaled)
{
    const int K = weight_data.w;
    const int N = weight_data.h;

    weight_data_scaled.create(K, N);
    if (weight_data_scaled.empty())
        return -100;

    const float* input_scale_ptr = input_scales;
    for (int n = 0; n < N; n++)
    {
        const float* ptr = weight_data.row(n);
        float* outptr = weight_data_scaled.row(n);

        for (int k = 0; k < K; k++)
            outptr[k] = ptr[k] / input_scale_ptr[k];
    }

    return 0;
}

static int make_dequant_weight(const ncnn::Mat& weight_data, const ncnn::Mat& weight_data_quantize_scales, int block_size, int weight_bits, ncnn::Mat& weight_data_dequantized)
{
    const int K = weight_data.w;
    const int N = weight_data.h;
    const int block_count = (K + block_size - 1) / block_size;

    weight_data_dequantized.create(K, N);
    if (weight_data_dequantized.empty())
        return -100;

    for (int n = 0; n < N; n++)
    {
        const float* ptr = weight_data.row(n);
        const float* scale_ptr = weight_data_quantize_scales.row(n);
        float* outptr = weight_data_dequantized.row(n);

        for (int b = 0; b < block_count; b++)
        {
            const int k0 = b * block_size;
            const int max_kk = std::min(block_size, K - k0);
            const float scale = scale_ptr[b];

            for (int k = 0; k < max_kk; k++)
            {
                const int q = float2int_weight(ptr[k0 + k] * scale, weight_bits);
                outptr[k0 + k] = q / scale;
            }
        }
    }

    return 0;
}

static double calc_awq_reconstruction_error(const ncnn::Mat& weight_data, const ncnn::Mat& input_scales, const ncnn::Mat& weight_data_dequantized, const QuantActStat& stat)
{
    const int K = weight_data.w;
    const int N = weight_data.h;
    const int sample_count = (int)(stat.samples.size() / K);

    const float* input_scale_ptr = input_scales;
    double error = 0.0;

    for (int i = 0; i < sample_count; i++)
    {
        const float* sample_ptr = &stat.samples[(size_t)i * K];

        for (int n = 0; n < N; n++)
        {
            const float* wptr = weight_data.row(n);
            const float* deqptr = weight_data_dequantized.row(n);

            double ref = 0.0;
            double out = 0.0;
            for (int k = 0; k < K; k++)
            {
                const float x = sample_ptr[k];
                ref += (double)x * wptr[k];
                out += (double)(x * input_scale_ptr[k]) * deqptr[k];
            }

            const double diff = ref - out;
            error += diff * diff;
        }
    }

    return error;
}

static int make_awq_input_scales(const ncnn::Mat& weight_data, const QuantActStat& stat, int block_size, int weight_bits, int inner_method, int awq_steps, float awq_max_scale, ncnn::Mat& input_scales, ncnn::Mat& weight_data_quantize_scales, int num_threads)
{
    const int K = weight_data.w;
    const int N = weight_data.h;

    ncnn::Mat act_mean(K);
    ncnn::Mat weight_mean(K);
    if (act_mean.empty() || weight_mean.empty())
        return -100;

    const float* sum_abs_ptr = stat.sum_abs;
    float* act_mean_ptr = act_mean;
    for (int k = 0; k < K; k++)
        act_mean_ptr[k] = sum_abs_ptr[k] / stat.count;

    weight_mean.fill(0.f);
    float* weight_mean_ptr = weight_mean;
    for (int n = 0; n < N; n++)
    {
        const float* ptr = weight_data.row(n);
        for (int k = 0; k < K; k++)
            weight_mean_ptr[k] += (float)fabs(ptr[k]);
    }
    for (int k = 0; k < K; k++)
        weight_mean_ptr[k] /= N;

    ncnn::Mat best_input_scales;
    ncnn::Mat best_quantize_scales;
    double best_error = DBL_MAX;

    const int search_steps = awq_steps;
    for (int s = 0; s <= search_steps; s++)
    {
        ncnn::Mat input_scales1(K);
        if (input_scales1.empty())
            return -100;

        float* input_scale_ptr = input_scales1;
        if (s == 0)
        {
            for (int k = 0; k < K; k++)
                input_scale_ptr[k] = 1.f;
        }
        else
        {
            const float alpha = (float)s / search_steps;
            double logsum = 0.0;
            for (int k = 0; k < K; k++)
            {
                float raw = powf((weight_mean_ptr[k] + 1e-6f) / (act_mean_ptr[k] + 1e-6f), alpha);
                if (raw < 1e-12f)
                    raw = 1e-12f;
                input_scale_ptr[k] = raw;
                logsum += ::log((double)raw);
            }

            const float geomean = (float)::exp(logsum / K);
            for (int k = 0; k < K; k++)
            {
                float v = input_scale_ptr[k] / geomean;
                if (v < 1.f / awq_max_scale)
                    v = 1.f / awq_max_scale;
                if (v > awq_max_scale)
                    v = awq_max_scale;

                input_scale_ptr[k] = v;
            }
        }

        ncnn::Mat weight_data_scaled;
        int ret = make_input_scaled_weight(weight_data, input_scales1, weight_data_scaled);
        if (ret != 0)
            return ret;

        ncnn::Mat quantize_scales1;
        ret = make_weight_scales(weight_data_scaled, block_size, weight_bits, inner_method, quantize_scales1, num_threads);
        if (ret != 0)
            return ret;

        ncnn::Mat weight_data_dequantized;
        ret = make_dequant_weight(weight_data_scaled, quantize_scales1, block_size, weight_bits, weight_data_dequantized);
        if (ret != 0)
            return ret;

        const double error = calc_awq_reconstruction_error(weight_data, input_scales1, weight_data_dequantized, stat);
        if (error < best_error)
        {
            best_error = error;
            best_input_scales = input_scales1;
            best_quantize_scales = quantize_scales1;
        }
    }

    input_scales = best_input_scales.clone();
    weight_data_quantize_scales = best_quantize_scales.clone();
    if (input_scales.empty() || weight_data_quantize_scales.empty())
        return -100;

    return 0;
}

static int cholesky_inverse_gptq_hessian(std::vector<double>& h, int n)
{
    for (int i = n - 1; i >= 0; i--)
    {
        double diag = h[(size_t)i * n + i];
        for (int k = i + 1; k < n; k++)
            diag -= h[(size_t)i * n + k] * h[(size_t)i * n + k];

        if (!(diag > 0.0))
            return -1;

        diag = sqrt(diag);
        h[(size_t)i * n + i] = diag;

        for (int j = 0; j < i; j++)
        {
            double v = h[(size_t)j * n + i];
            for (int k = i + 1; k < n; k++)
                v -= h[(size_t)j * n + k] * h[(size_t)i * n + k];

            h[(size_t)j * n + i] = v / diag;
        }
    }

    for (int i = n - 1; i >= 0; i--)
    {
        const double diag = h[(size_t)i * n + i];

        for (int j = n - 1; j > i; j--)
        {
            double v = 0.0;
            for (int k = i + 1; k <= j; k++)
                v += h[(size_t)i * n + k] * h[(size_t)k * n + j];

            h[(size_t)i * n + j] = -v / diag;
        }

        h[(size_t)i * n + i] = 1.0 / diag;
    }

    return 0;
}

static int make_gptq_qweight(const ncnn::Mat& weight_data, const QuantActStat& stat, int block_size, int weight_bits, float damp, ncnn::Mat& weight_data_quantize_scales, ncnn::Mat& weight_data_quantized, int num_threads)
{
    const int K = weight_data.w;
    const int N = weight_data.h;
    const int block_count = (K + block_size - 1) / block_size;
    const int packed_k_bytes = llm_weight_quantize_packed_k_bytes(K, weight_bits);
    const int sample_count = (int)(stat.samples.size() / K);

    weight_data_quantize_scales.create(block_count, N);
    weight_data_quantized.create(packed_k_bytes, N, (size_t)1u);
    if (weight_data_quantize_scales.empty() || weight_data_quantized.empty())
        return -100;

    memset(weight_data_quantized.data, 0, weight_data_quantized.total() * weight_data_quantized.elemsize);

    #pragma omp parallel for num_threads(num_threads)
    for (int n = 0; n < N; n++)
    {
        const float* ptr = weight_data.row(n);
        float* scale_ptr = weight_data_quantize_scales.row(n);
        for (int b = 0; b < block_count; b++)
        {
            const int k0 = b * block_size;
            const int max_kk = std::min(block_size, K - k0);
            scale_ptr[b] = choose_weight_scale(ptr + k0, max_kk, weight_bits, LLM_QUANT_METHOD_MINMAX);
        }
    }

    std::vector<double> H((size_t)K * K, 0.0);

    #pragma omp parallel for num_threads(num_threads)
    for (int i = 0; i < K; i++)
    {
        for (int j = i; j < K; j++)
        {
            double v = 0.0;
            for (int s = 0; s < sample_count; s++)
            {
                const float* x = &stat.samples[(size_t)s * K];
                v += (double)x[i] * x[j];
            }

            H[(size_t)i * K + j] = v;
        }
    }

    double diag_mean = 0.0;
    for (int i = 0; i < K; i++)
        diag_mean += H[(size_t)i * K + i];
    diag_mean /= K;

    for (int i = 0; i < K; i++)
        H[(size_t)i * K + i] += damp * diag_mean;

    if (cholesky_inverse_gptq_hessian(H, K) != 0)
    {
        fprintf(stderr, "gptq hessian decomposition failed\n");
        return -1;
    }

    // H is the upper cholesky factor of the inverse hessian
    #pragma omp parallel num_threads(num_threads)
    {
        std::vector<double> work(K);
        std::vector<double> error(block_size);

        #pragma omp for
        for (int n = 0; n < N; n++)
        {
            const float* ptr = weight_data.row(n);
            for (int k = 0; k < K; k++)
                work[k] = ptr[k];

            const float* scale_ptr = weight_data_quantize_scales.row(n);
            unsigned char* qptr = weight_data_quantized.row<unsigned char>(n);

            for (int k0 = 0; k0 < K; k0 += block_size)
            {
                const int max_kk = std::min(block_size, K - k0);
                const int k1 = k0 + max_kk;

                for (int kk = 0; kk < max_kk; kk++)
                {
                    const int k = k0 + kk;
                    const float scale = scale_ptr[k / block_size];
                    const double w = work[k];
                    const int q = float2int_weight((float)(w * scale), weight_bits);
                    const double deq = (double)q / scale;
                    const double err = (w - deq) / H[(size_t)k * K + k];

                    pack_signed_weight(qptr, k, weight_bits, q);
                    error[kk] = err;

                    for (int j = k; j < k1; j++)
                        work[j] -= err * H[(size_t)k * K + j];
                }

                for (int j = k1; j < K; j++)
                {
                    double err = 0.0;
                    for (int kk = 0; kk < max_kk; kk++)
                        err += error[kk] * H[(size_t)(k0 + kk) * K + j];

                    work[j] -= err;
                }
            }
        }
    }

    return 0;
}

int QuantNet::init(int method)
{
    if (method != LLM_QUANT_METHOD_AWQ && method != LLM_QUANT_METHOD_GPTQ)
        return 0;

    for (size_t i = 0; i < layers.size(); i++)
    {
        const ncnn::Layer* layer = layers[i];
        if (layer->type == "Input")
            input_blobs.push_back(layer->tops[0]);
    }

    const int max_sample_count = method == LLM_QUANT_METHOD_AWQ ? awq_samples : gptq_samples;

    for (size_t i = 0; i < layers.size(); i++)
    {
        if (layers[i]->type == "Gemm")
        {
            const ncnn::Gemm* gemm = (const ncnn::Gemm*)layers[i];

            if (!is_supported_llm_gemm(gemm))
                continue;

            QuantActStat stat;
            int ret = init_act_stat(stat, gemm->constantK, max_sample_count);
            if (ret != 0)
                return ret;

            gemm_act_stats.push_back(stat);
        }
        else if (layers[i]->type == "MultiHeadAttention")
        {
            const ncnn::MultiHeadAttention* mha = (const ncnn::MultiHeadAttention*)layers[i];

            if (!is_supported_llm_multiheadattention(mha))
                continue;

            const size_t mha_act_offset = mha_act_stats.size();
            mha_act_stats.resize(mha_act_offset + 4);

            const int qdim = mha->weight_data_size / mha->embed_dim;
            const int Ks[4] = {qdim, mha->kdim, mha->vdim, mha->embed_dim};

            for (int j = 0; j < 4; j++)
            {
                int ret = init_act_stat(mha_act_stats[mha_act_offset + j], Ks[j], max_sample_count);
                if (ret != 0)
                    return ret;
            }
        }
    }

    return 0;
}

static int check_calibration_input(const QuantNet& table)
{
    const int input_blob_count = (int)table.input_blobs.size();
    if (input_blob_count == 0)
    {
        fprintf(stderr, "no input blob found\n");
        return -1;
    }

    if ((int)table.listspaths.size() != input_blob_count)
    {
        fprintf(stderr, "calibration list count mismatch expected=%d got=%d\n", input_blob_count, (int)table.listspaths.size());
        return -1;
    }

    if ((int)table.shapes.size() != input_blob_count)
    {
        fprintf(stderr, "shape count mismatch expected=%d got=%d\n", input_blob_count, (int)table.shapes.size());
        return -1;
    }

    if (table.file_type != 1)
    {
        fprintf(stderr, "ncnnllm2table calibration supports type=1 only\n");
        return -1;
    }

    if (table.listspaths[0].empty())
    {
        fprintf(stderr, "calibration list is empty\n");
        return -1;
    }

    const int file_count = (int)table.listspaths[0].size();
    for (int i = 1; i < input_blob_count; i++)
    {
        if ((int)table.listspaths[i].size() != file_count)
        {
            fprintf(stderr, "calibration list size mismatch input=%d expected=%d got=%d\n", i, file_count, (int)table.listspaths[i].size());
            return -1;
        }
    }

    return 0;
}

int QuantNet::collect_activation_stats()
{
    if (check_calibration_input(*this) != 0)
        return -1;

    const int input_blob_count = (int)input_blobs.size();
    const int file_count = (int)listspaths[0].size();

    for (int i = 0; i < file_count; i++)
    {
        if (i % 100 == 0)
            fprintf(stderr, "collect_activation_stats %.2f%% [ %d / %d ]\n", i * 100.f / file_count, i, file_count);

        ncnn::Extractor ex = create_extractor();
        ex.set_light_mode(false);

        for (int j = 0; j < input_blob_count; j++)
        {
            ncnn::Mat in = read_npy(shapes[j], listspaths[j][i]);
            if (in.empty())
                return -1;

            ex.input(input_blobs[j], in);
        }

        size_t gemm_act_index = 0;
        size_t mha_act_index = 0;

        for (size_t j = 0; j < layers.size(); j++)
        {
            if (layers[j]->type == "Gemm")
            {
                const ncnn::Gemm* gemm = (const ncnn::Gemm*)layers[j];

                if (!is_supported_llm_gemm(gemm))
                    continue;

                ncnn::Mat bottom_blob;
                ex.extract(gemm->bottoms[0], bottom_blob);

                collect_act_rows(gemm_act_stats[gemm_act_index], bottom_blob);

                gemm_act_index++;
            }
            else if (layers[j]->type == "MultiHeadAttention")
            {
                const ncnn::MultiHeadAttention* mha = (const ncnn::MultiHeadAttention*)layers[j];

                if (!is_supported_llm_multiheadattention(mha))
                    continue;

                int q_blob_i = 0;
                int k_blob_i = 0;
                int v_blob_i = 0;
                int attn_mask_i = 0;
                int cached_xk_i = 0;
                int cached_xv_i = 0;
                resolve_mha_bottom_blob_index(mha, (int)mha->bottoms.size(), q_blob_i, k_blob_i, v_blob_i, attn_mask_i, cached_xk_i, cached_xv_i);

                ncnn::Mat q_blob;
                ncnn::Mat k_blob;
                ncnn::Mat v_blob;
                ncnn::Mat attn_mask_blob;
                ncnn::Mat cached_xk_blob;
                ncnn::Mat cached_xv_blob;

                ex.extract(mha->bottoms[q_blob_i], q_blob);
                ex.extract(mha->bottoms[k_blob_i], k_blob);
                ex.extract(mha->bottoms[v_blob_i], v_blob);
                if (mha->attn_mask)
                    ex.extract(mha->bottoms[attn_mask_i], attn_mask_blob);
                if (mha->kv_cache)
                {
                    ex.extract(mha->bottoms[cached_xk_i], cached_xk_blob);
                    ex.extract(mha->bottoms[cached_xv_i], cached_xv_blob);
                }

                const size_t mha_act_offset = mha_act_index * 4;
                collect_act_rows(mha_act_stats[mha_act_offset], q_blob);
                collect_act_rows(mha_act_stats[mha_act_offset + 1], k_blob);
                collect_act_rows(mha_act_stats[mha_act_offset + 2], v_blob);

                int ret = collect_mha_out_act_rows(mha, q_blob, k_blob, v_blob, attn_mask_blob, cached_xk_blob, cached_xv_blob, q_blob_i, k_blob_i, v_blob_i, mha_act_stats[mha_act_offset + 3], quantize_num_threads);
                if (ret != 0)
                    return ret;

                mha_act_index++;
            }
        }
    }

    fprintf(stderr, "collect_activation_stats 100.00%% [ %d / %d ]\n", file_count, file_count);

    return 0;
}

int QuantNet::save_table(const char* tablepath, int block_size, int weight_bits, int method) const
{
    FILE* fp = fopen(tablepath, "wb");
    if (!fp)
    {
        fprintf(stderr, "fopen %s failed\n", tablepath);
        return -1;
    }

    int table_count = 0;
    size_t gemm_act_index = 0;
    size_t mha_act_index = 0;

    for (size_t i = 0; i < layers.size(); i++)
    {
        if (layers[i]->type == "Gemm")
        {
            const ncnn::Gemm* gemm = (const ncnn::Gemm*)layers[i];

            if (!is_supported_llm_gemm(gemm))
            {
                print_skip_gemm(gemm);
                continue;
            }

            char key[256];
            snprintf(key, 256, "%s_param_1", gemm_name(gemm));

            ncnn::Mat B_data_quantize_scales;
            ncnn::Mat B_data_input_scales;
            ncnn::Mat B_data_quantized;
            int ret = 0;
            if (method == LLM_QUANT_METHOD_AWQ)
            {
                ret = make_awq_input_scales(gemm->B_data, gemm_act_stats[gemm_act_index], block_size, weight_bits, awq_inner_method, awq_steps, awq_max_scale, B_data_input_scales, B_data_quantize_scales, quantize_num_threads);
            }
            else if (method == LLM_QUANT_METHOD_GPTQ)
            {
                ret = make_gptq_qweight(gemm->B_data, gemm_act_stats[gemm_act_index], block_size, weight_bits, gptq_damp, B_data_quantize_scales, B_data_quantized, quantize_num_threads);
            }
            else
            {
                ret = make_weight_scales(gemm->B_data, block_size, weight_bits, method, B_data_quantize_scales, quantize_num_threads);
            }
            if (ret != 0)
            {
                fclose(fp);
                return ret;
            }

            if (method == LLM_QUANT_METHOD_GPTQ)
            {
                std::string qweight_path;
                const std::string qweight_name = make_qweight_filename(tablepath, key, qweight_path);
                ret = write_raw_mat_file(qweight_path.c_str(), B_data_quantized);
                if (ret != 0)
                {
                    fclose(fp);
                    return ret;
                }

                if (write_llm_qweight_table_row(fp, key, weight_bits, block_size, method, qweight_name.c_str(), B_data_quantize_scales) != 0)
                {
                    fclose(fp);
                    return -1;
                }
            }
            else if (write_llm_table_row(fp, key, weight_bits, block_size, method, B_data_quantize_scales) != 0)
            {
                fclose(fp);
                return -1;
            }

            fprintf(stderr, "write_llm_table %s\n", key);
            table_count++;

            if (method == LLM_QUANT_METHOD_AWQ)
            {
                char input_scale_key[512];
                snprintf(input_scale_key, 512, "%s_input_scale", key);

                if (write_llm_input_scale_row(fp, input_scale_key, method, B_data_input_scales) != 0)
                {
                    fclose(fp);
                    return -1;
                }

                fprintf(stderr, "write_llm_table %s\n", input_scale_key);
                table_count++;
            }

            gemm_act_index++;
        }
        else if (layers[i]->type == "MultiHeadAttention")
        {
            const ncnn::MultiHeadAttention* mha = (const ncnn::MultiHeadAttention*)layers[i];

            if (!is_supported_llm_multiheadattention(mha))
            {
                print_skip_multiheadattention(mha);
                continue;
            }

            const int qdim = mha->weight_data_size / mha->embed_dim;

            const ncnn::Mat q_weight_data = mha->q_weight_data.reshape(qdim, mha->embed_dim);
            const ncnn::Mat k_weight_data = mha->k_weight_data.reshape(mha->kdim, mha->embed_dim);
            const ncnn::Mat v_weight_data = mha->v_weight_data.reshape(mha->vdim, mha->embed_dim);
            const ncnn::Mat out_weight_data = mha->out_weight_data.reshape(mha->embed_dim, qdim);

            const ncnn::Mat weights[4] = {q_weight_data, k_weight_data, v_weight_data, out_weight_data};

            const size_t mha_act_offset = mha_act_index * 4;
            for (int j = 0; j < 4; j++)
            {
                char key[256];
                snprintf(key, 256, "%s_param_%d", multiheadattention_name(mha), j);

                ncnn::Mat weight_data_quantize_scales;
                ncnn::Mat weight_data_input_scales;
                ncnn::Mat weight_data_quantized;
                int ret = 0;
                if (method == LLM_QUANT_METHOD_AWQ)
                {
                    ret = make_awq_input_scales(weights[j], mha_act_stats[mha_act_offset + j], block_size, weight_bits, awq_inner_method, awq_steps, awq_max_scale, weight_data_input_scales, weight_data_quantize_scales, quantize_num_threads);
                }
                else if (method == LLM_QUANT_METHOD_GPTQ)
                {
                    ret = make_gptq_qweight(weights[j], mha_act_stats[mha_act_offset + j], block_size, weight_bits, gptq_damp, weight_data_quantize_scales, weight_data_quantized, quantize_num_threads);
                }
                else
                {
                    ret = make_weight_scales(weights[j], block_size, weight_bits, method, weight_data_quantize_scales, quantize_num_threads);
                }
                if (ret != 0)
                {
                    fclose(fp);
                    return ret;
                }

                if (method == LLM_QUANT_METHOD_GPTQ)
                {
                    std::string qweight_path;
                    const std::string qweight_name = make_qweight_filename(tablepath, key, qweight_path);
                    ret = write_raw_mat_file(qweight_path.c_str(), weight_data_quantized);
                    if (ret != 0)
                    {
                        fclose(fp);
                        return ret;
                    }

                    if (write_llm_qweight_table_row(fp, key, weight_bits, block_size, method, qweight_name.c_str(), weight_data_quantize_scales) != 0)
                    {
                        fclose(fp);
                        return -1;
                    }
                }
                else if (write_llm_table_row(fp, key, weight_bits, block_size, method, weight_data_quantize_scales) != 0)
                {
                    fclose(fp);
                    return -1;
                }

                fprintf(stderr, "write_llm_table %s\n", key);
                table_count++;

                if (method == LLM_QUANT_METHOD_AWQ)
                {
                    char input_scale_key[512];
                    snprintf(input_scale_key, 512, "%s_input_scale", key);

                    if (write_llm_input_scale_row(fp, input_scale_key, method, weight_data_input_scales) != 0)
                    {
                        fclose(fp);
                        return -1;
                    }

                    fprintf(stderr, "write_llm_table %s\n", input_scale_key);
                    table_count++;
                }
            }

            mha_act_index++;
        }
    }

    fclose(fp);

    fprintf(stderr, "ncnn llm quant table created %d rows\n", table_count);

    return 0;
}

int main(int argc, char** argv)
{
    if (argc < 4)
    {
        show_usage(argv[0]);
        return -1;
    }

    for (int i = 1; i < argc; i++)
    {
        if (argv[i][0] == '-')
        {
            show_usage(argv[0]);
            return -1;
        }
    }

    char* inparam = argv[1];
    char* inbin = argv[2];
    char* caliblist = 0;
    char* outtable = 0;
    int kv_start = 4;

    if (argc >= 5 && strchr(argv[3], '=') == NULL && strchr(argv[4], '=') == NULL)
    {
        caliblist = argv[3];
        outtable = argv[4];
        kv_start = 5;
    }
    else
    {
        outtable = argv[3];
        kv_start = 4;
    }

    int weight_bits = 6;
    int block_size = 64;
    const char* method = "minmax";
    const char* awq_inner = "minmax";
    int thread = 1;

    QuantNet table;
    table.use_calibration_dataset = caliblist != 0;
    if (caliblist)
        table.listspaths = parse_comma_path_list(caliblist);

    for (int i = kv_start; i < argc; i++)
    {
        char* kv = argv[i];
        char* eqs = strchr(kv, '=');
        if (eqs == NULL)
        {
            fprintf(stderr, "unrecognized arg %s\n", kv);
            return -1;
        }

        eqs[0] = '\0';
        const char* key = kv;
        char* value = eqs + 1;

        if (strcmp(key, "method") == 0)
            method = value;
        else if (strcmp(key, "bits") == 0)
            weight_bits = atoi(value);
        else if (strcmp(key, "block") == 0)
            block_size = atoi(value);
        else if (strcmp(key, "thread") == 0)
            thread = atoi(value);
        else if (strcmp(key, "type") == 0)
            table.file_type = atoi(value);
        else if (strcmp(key, "shape") == 0)
            table.shapes = parse_comma_int_array_list(value);
        else if (strcmp(key, "awq_steps") == 0)
            table.awq_steps = atoi(value);
        else if (strcmp(key, "awq_samples") == 0)
            table.awq_samples = atoi(value);
        else if (strcmp(key, "awq_max_scale") == 0)
            table.awq_max_scale = (float)atof(value);
        else if (strcmp(key, "awq_inner") == 0)
            awq_inner = value;
        else if (strcmp(key, "gptq_samples") == 0)
            table.gptq_samples = atoi(value);
        else if (strcmp(key, "gptq_damp") == 0)
            table.gptq_damp = (float)atof(value);
        else
        {
            fprintf(stderr, "unrecognized arg %s\n", key);
            return -1;
        }
    }

    const int quantize_method = llm_quant_method_from_string(method);
    if (quantize_method < 0)
    {
        fprintf(stderr, "unsupported method=%s\n", method);
        return -1;
    }

    if ((quantize_method == LLM_QUANT_METHOD_AWQ || quantize_method == LLM_QUANT_METHOD_GPTQ) && !table.use_calibration_dataset)
    {
        fprintf(stderr, "method=%s requires calibration list\n", method);
        return -1;
    }

    if (quantize_method != LLM_QUANT_METHOD_AWQ && quantize_method != LLM_QUANT_METHOD_GPTQ && table.use_calibration_dataset)
    {
        fprintf(stderr, "calibration list is only used by method=awq/gptq\n");
        return -1;
    }

    if (llm_weight_block_quantize_term(weight_bits, block_size) == 0)
    {
        show_usage(argv[0]);
        return -1;
    }

    if (thread < 1)
    {
        fprintf(stderr, "malformed thread %d\n", thread);
        return -1;
    }

    if (quantize_method == LLM_QUANT_METHOD_AWQ)
    {
        table.awq_inner_method = llm_quant_method_from_string(awq_inner);
        if (table.awq_inner_method != LLM_QUANT_METHOD_MINMAX && table.awq_inner_method != LLM_QUANT_METHOD_MSECLIP)
        {
            fprintf(stderr, "unsupported awq_inner=%s\n", awq_inner);
            return -1;
        }

        if (table.awq_steps < 0)
        {
            fprintf(stderr, "malformed awq_steps %d\n", table.awq_steps);
            return -1;
        }

        if (table.awq_samples < 1)
        {
            fprintf(stderr, "malformed awq_samples %d\n", table.awq_samples);
            return -1;
        }

        if (!(table.awq_max_scale > 1.f))
        {
            fprintf(stderr, "malformed awq_max_scale %f\n", table.awq_max_scale);
            return -1;
        }
    }

    if (quantize_method == LLM_QUANT_METHOD_GPTQ)
    {
        if (table.gptq_samples < 1)
        {
            fprintf(stderr, "malformed gptq_samples %d\n", table.gptq_samples);
            return -1;
        }

        if (!(table.gptq_damp >= 0.f))
        {
            fprintf(stderr, "malformed gptq_damp %f\n", table.gptq_damp);
            return -1;
        }
    }

    table.quantize_num_threads = thread;
    table.opt.num_threads = thread;
    table.opt.use_packing_layout = false;
    table.opt.use_fp16_packed = false;
    table.opt.use_fp16_storage = false;
    table.opt.use_fp16_arithmetic = false;
    table.opt.use_bf16_storage = false;

    if (table.load_param(inparam) != 0)
        return -1;

    if (strcmp(inbin, "null") == 0)
    {
        DataReaderFromEmpty dr;
        if (table.load_model(dr) != 0)
            return -1;
    }
    else
    {
        if (table.load_model(inbin) != 0)
            return -1;
    }

    if (table.init(quantize_method) != 0)
        return -1;

    if (quantize_method == LLM_QUANT_METHOD_AWQ || quantize_method == LLM_QUANT_METHOD_GPTQ)
    {
        if (table.collect_activation_stats() != 0)
            return -1;
    }

    if (table.save_table(outtable, block_size, weight_bits, quantize_method) != 0)
        return -1;

    return 0;
}