#include "ggml.h"
#include "ggml-cpu.h"
#include "ggml-backend.h"
#include "ggml-alloc.h"
#include <cstdio>
#include <cmath>
#include <vector>
#include <random>
#include <cstring>
int main(){
    ggml_backend_t be = ggml_backend_cpu_init(); ggml_backend_cpu_set_n_threads(be,16);
    std::mt19937 g(7); std::normal_distribution<float> nd(0,1);
    ggml_type types[]={GGML_TYPE_Q4_K,GGML_TYPE_Q5_K,GGML_TYPE_Q6_K,GGML_TYPE_IQ4_XS,GGML_TYPE_Q4_1,GGML_TYPE_Q8_0};
    const int K=2048,N=512,E=64,NU=6;
    for(ggml_type t:types) for(int mode=0;mode<2;mode++) for(int T: {1,64}){
        ggml_quantize_init(t);
        size_t ctxsz=ggml_tensor_overhead()*16+ggml_graph_overhead()+(1<<20);
        ggml_init_params ip={ctxsz,nullptr,true}; ggml_context*c=ggml_init(ip);
        int ne=(mode==0)?1:E;
        ggml_tensor*w=ggml_new_tensor_3d(c,t,K,N,ne);
        ggml_tensor*x; ggml_tensor*ids=nullptr; ggml_tensor*y;
        if(mode==0){ x=ggml_new_tensor_2d(c,GGML_TYPE_F32,K,T); y=ggml_mul_mat(c,ggml_view_2d(c,w,K,N,w->nb[1],0),x);} 
        else { x=ggml_new_tensor_3d(c,GGML_TYPE_F32,K,NU,T); ids=ggml_new_tensor_2d(c,GGML_TYPE_I32,NU,T); y=ggml_mul_mat_id(c,w,x,ids);} 
        ggml_cgraph*gr=ggml_new_graph(c); ggml_build_forward_expand(gr,y);
        ggml_backend_buffer_t buf=ggml_backend_alloc_ctx_tensors(c,be);
        // data
        size_t rows=(size_t)N*ne; std::vector<float> wf((size_t)K*rows), im(K,1.f);
        for(auto&v:wf)v=nd(g)*0.05f;
        std::vector<char> wq(ggml_row_size(t,K)*rows);
        ggml_quantize_chunk(t,wf.data(),wq.data(),0,rows,K,im.data());
        ggml_backend_tensor_set(w,wq.data(),0,wq.size());
        std::vector<float> xf((size_t)ggml_nelements(x)); for(auto&v:xf)v=nd(g);
        ggml_backend_tensor_set(x,xf.data(),0,xf.size()*4);
        std::vector<int32_t> idv; 
        if(ids){ idv.resize((size_t)NU*T); for(int tt=0;tt<T;tt++){ for(int u=0;u<NU;u++) idv[tt*NU+u]=(tt*7+u*11+3)%E; } ggml_backend_tensor_set(ids,idv.data(),0,idv.size()*4);}        
        ggml_backend_graph_compute(be,gr);
        std::vector<float> out(ggml_nelements(y)); ggml_backend_tensor_get(y,out.data(),0,out.size()*4);
        // reference: dequantize weights, float dot
        std::vector<float> wd((size_t)K*rows);
        auto tr=ggml_get_type_traits(t); tr->to_float(wq.data(),wd.data(),(int64_t)K*rows);
        double num=0,den=0;
        for(int tt=0;tt<T;tt++){
            int nu=(mode==0)?1:NU;
            for(int u=0;u<nu;u++){
                int e=(mode==0)?0:idv[tt*NU+u];
                const float* xv=&xf[(size_t)(mode==0? tt*K : (tt*NU+u)*K)];
                for(int n=0;n<N;n++){
                    double r=0; const float*wr=&wd[((size_t)e*N+n)*K]; for(int k=0;k<K;k++) r+=(double)wr[k]*xv[k];
                    double o=out[mode==0? (size_t)tt*N+n : ((size_t)tt*NU+u)*N+n];
                    num+=(o-r)*(o-r); den+=r*r;
                }
            }
        }
        printf("%-7s %s T=%-3d rel_L2_err=%.3e\n",ggml_type_name(t),mode?"mul_mat_id":"mul_mat   ",T,sqrt(num/den));
        ggml_backend_buffer_free(buf); ggml_free(c);
    }
}