export const H = 2048,
  F = 6144,
  VOCAB = 130560,
  LAYERS = 42,
  KV = 256,
  HEADS = 16,
  D = 128,
  MAX_CHUNKS = 32,
  PREFILL_BATCH = 32;
const STEP = `struct Step { position:u32, count:u32, chunk:u32, chunks:u32 }; @group(0) @binding(0) var<uniform> step:Step;\n`;
const bind = (index, type, access = "read") =>
  `@group(1) @binding(${index}) var<storage,${access}> ${type};\n`;
const Q8_CODE = `
fn weightAt(i:u32)->f32 {
 let block=i>>5u;let b=block*34u;let scale=unpack2x16float(weights[b>>2u])[(b&2u)>>1u];
 let at=b+2u+(i&31u);let word=weights[at>>2u];let q=bitcast<i32>(word<<((3u-(at&3u))*8u))>>24;
 return f32(q)*scale;
}
fn weight4(block:u32,lane:u32)->vec4<f32>{
 let b=block*34u;let scale=unpack2x16float(weights[b>>2u])[(b&2u)>>1u];
 let a=b+2u+lane*4u;let shift=(a&3u)*8u;var word=weights[a>>2u];
 if(shift!=0u){word=(word>>shift)|(weights[(a>>2u)+1u]<<(32u-shift));}
 let q=vec4<i32>(bitcast<i32>(word<<24u)>>24,bitcast<i32>(word<<16u)>>24,bitcast<i32>(word<<8u)>>24,bitcast<i32>(word)>>24);
 return vec4<f32>(q)*scale;
}
`;
function sumDecl(subgroups, n = 256) {
  return subgroups
    ? `fn teamSum(v:f32,tid:u32)->f32{var x=v;for(var s=1u;s<32u;s<<=1u){x+=subgroupShuffleXor(x,s);}return x;}`
    : `var<workgroup> red:array<f32,${n}>;fn teamSum(v:f32,tid:u32)->f32{red[tid]=v;workgroupBarrier();let lane=tid&31u;let base=tid-lane;for(var s=16u;s>0u;s>>=1u){if(lane<s){red[tid]+=red[tid+s];}workgroupBarrier();}let x=red[base];workgroupBarrier();return x;}`;
}
export function buildMiniCPMKernels({
  subgroups = false,
  kv = "f16",
  weightFormat = "q8",
  draft = false,
  smallWorkgroups = false,
} = {}) {
  const Q8 =
    weightFormat === "bf16"
      ? `fn weightAt(i:u32)->f32{let w=weights[i>>1u];return bitcast<f32>(select(w<<16u,w&0xffff0000u,(i&1u)==1u));}fn weight4(block:u32,lane:u32)->vec4<f32>{let i=block*32u+lane*4u;let a=weights[i>>1u];let b=weights[(i>>1u)+1u];return vec4<f32>(bitcast<f32>(a<<16u),bitcast<f32>(a&0xffff0000u),bitcast<f32>(b<<16u),bitcast<f32>(b&0xffff0000u));}`
      : Q8_CODE;
  const prefix = (subgroups ? "enable subgroups;\n" : "") + STEP,
    K = {},
    sum = sumDecl(subgroups);
  const storageType = kv === "f32" ? "f32" : "u32";
  const readKV = (name) =>
    kv === "f32"
      ? `return vec4<f32>(${name}[i],${name}[i+1u],${name}[i+2u],${name}[i+3u]);`
      : kv === "f16"
        ? `return vec4<f32>(unpack2x16float(${name}[i>>1u]),unpack2x16float(${name}[(i>>1u)+1u]));`
        : kv === "q8_0"
          ? `let b=i>>5u;let base=b*9u;let w=${name}[base+1u+((i&31u)>>2u)];return unpack4x8snorm(w)*bitcast<f32>(${name}[base]);`
          : `let b=i>>5u;let base=b*5u;let w=${name}[base+1u+((i&31u)>>3u)]>>((i&4u)*4u);return (vec4<f32>(f32(w&15u),f32((w>>4u)&15u),f32((w>>8u)&15u),f32((w>>12u)&15u))-7.0)*bitcast<f32>(${name}[base]);`;
  const readers = `fn readK(i:u32)->vec4<f32>{${readKV("keys")}} fn readV(i:u32)->vec4<f32>{${readKV("values")}}`;
  K.embed =
    prefix +
    bind(0, "weights:array<u32>") +
    bind(1, "tokens:array<u32>") +
    bind(2, "x:array<f32>", "read_write") +
    Q8 +
    `
 @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g:vec3<u32>){if(g.x>=${H}u||g.y>=step.count){return;}x[g.y*${H}u+g.x]=weightAt(tokens[step.position+g.y]*${H}u+g.x);}`;
  K.norm =
    prefix +
    bind(0, "x:array<f32>", "read_write") +
    bind(1, "residual:array<f32>") +
    bind(2, "weights:array<f32>") +
    bind(3, "out:array<f32>", "read_write") +
    `
 var<workgroup> squares:array<f32,256>;
 @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l:vec3<u32>,@builtin(workgroup_id) g:vec3<u32>){
 let base=g.x*${H}u;var ss=0.0;for(var i=l.x;i<${H}u;i+=256u){let v=x[base+i]+residual[base+i];x[base+i]=v;ss+=v*v;}squares[l.x]=ss;workgroupBarrier();
 for(var s=128u;s>0u;s>>=1u){if(l.x<s){squares[l.x]+=squares[l.x+s];}workgroupBarrier();}
 let scale=inverseSqrt(squares[0]/${H}.0+0.000001);for(var i=l.x;i<${H}u;i+=256u){out[base+i]=x[base+i]*scale*weights[i];}}
 `;
  const gemv = (columns, rows, gate = false) =>
    prefix +
    bind(0, "input:array<f32>") +
    bind(1, "weights:array<u32>") +
    bind(2, "out:array<f32>", "read_write") +
    Q8 +
    sum +
    `
 var<workgroup> sharedX:array<vec4<f32>,${columns / 4}>;
 @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l:vec3<u32>,@builtin(workgroup_id) g:vec3<u32>){
 for(var i=l.x;i<${columns / 4}u;i+=256u){let b=i*4u;sharedX[i]=vec4<f32>(input[b],input[b+1u],input[b+2u],input[b+3u]);}workgroupBarrier();
 let row=g.x*8u+(l.x>>5u);let lane=l.x&31u;var acc=0.0;${gate ? "var other=0.0;" : ""}
 for(var b=lane;b<${columns / 32}u;b+=32u){for(var w=0u;w<8u;w++){let v=sharedX[b*8u+w];acc+=dot(v,weight4(row*${columns / 32}u+b,w));${gate ? `other+=dot(v,weight4((row+${rows}u)*${columns / 32}u+b,w));` : ""}}}
 let total=teamSum(acc,l.x);${gate ? "let up=teamSum(other,l.x);" : ""}
 if(lane==0u&&row<${rows}u){out[row]=${gate ? "(total/(1.0+exp(-total)))*up" : "total"};}}
 `;
  const gemm = (columns, rows) =>
    prefix +
    bind(0, "input:array<f32>") +
    bind(1, "weights:array<u32>") +
    bind(2, "out:array<f32>", "read_write") +
    Q8 +
    `
 var<workgroup> a:array<f32,1024>;var<workgroup> b:array<f32,1024>;
 @compute @workgroup_size(16,16) fn main(@builtin(local_invocation_id) l:vec3<u32>,@builtin(workgroup_id) g:vec3<u32>){
 let tid=l.y*16u+l.x;let row=g.x*32u+l.x*2u;let token=g.y*32u+l.y*2u;var acc=vec4<f32>(0.0);
 for(var kk=0u;kk<${columns}u;kk+=32u){
 for(var j=tid;j<1024u;j+=256u){let r=j>>5u;let c=j&31u;let t=g.y*32u+r;var av=0.0;if(t<step.count){av=input[t*${columns}u+kk+c];}a[j]=av;b[j]=weightAt((g.x*32u+r)*${columns}u+kk+c);}workgroupBarrier();
 for(var k=0u;k<32u;k++){let aa=vec2<f32>(a[(l.y*2u)*32u+k],a[(l.y*2u+1u)*32u+k]);let bb=vec2<f32>(b[(l.x*2u)*32u+k],b[(l.x*2u+1u)*32u+k]);acc+=vec4<f32>(aa.x*bb.x,aa.x*bb.y,aa.y*bb.x,aa.y*bb.y);}workgroupBarrier();}
 if(token<step.count){out[token*${rows}u+row]=acc.x;out[token*${rows}u+row+1u]=acc.y;}
 if(token+1u<step.count){out[(token+1u)*${rows}u+row]=acc.z;out[(token+1u)*${rows}u+row+1u]=acc.w;}}
 `;
  const block = (columns, rows) =>
    prefix +
    bind(0, "input:array<vec4<f32>>") +
    bind(1, "weights:array<u32>") +
    bind(2, "out:array<f32>", "read_write") +
    Q8 +
    (subgroups
      ? `fn teamSum4(v:vec4<f32>,tid:u32)->vec4<f32>{var x=v;for(var s=1u;s<32u;s<<=1u){x+=subgroupShuffleXor(x,s);}return x;}`
      : `var<workgroup> red4:array<vec4<f32>,256>;fn teamSum4(v:vec4<f32>,tid:u32)->vec4<f32>{red4[tid]=v;workgroupBarrier();let lane=tid&31u;for(var s=16u;s>0u;s>>=1u){if(lane<s){red4[tid]+=red4[tid+s];}workgroupBarrier();}let r=red4[tid-lane];workgroupBarrier();return r;}`) +
    `
 @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let row=g.x*8u+(l.x>>5u);let lane=l.x&31u;var a=vec4<f32>(0.0);var b=vec4<f32>(0.0);
 for(var j=lane;j<${columns / 32}u;j+=32u){for(var w=0u;w<8u;w++){let v=weight4(row*${columns / 32}u+j,w);let at=j*8u+w;
 ${Array.from({ length: 8 }, (_, i) => `if(${i}u<step.count){${i < 4 ? "a" : "b"}[${i % 4}]+=dot(v,input[${(i * columns) / 4}u+at]);}`).join("\n")}
 }}let aa=teamSum4(a,l.x);let bb=teamSum4(b,l.x);if(lane==0u){${Array.from({ length: 8 }, (_, i) => `if(${i}u<step.count){out[${i * rows}u+row]=${i < 4 ? "aa" : "bb"}[${i % 4}];}`).join("\n")}}}
 `;
  for (const [name, cols, rows] of [
    ["qkv", H, H + 2 * KV],
    ["output", H, H],
    ["down", F, H],
    ["head", H, VOCAB],
  ]) {
    K[name] = gemv(cols, rows);
    if (name !== "head") K[name + "Batch"] = gemm(cols, rows);
    K[name + "Block"] = block(cols, rows);
  }
  K.upGate = gemv(H, F, true);
  K.upGateBatch = gemm(H, F * 2);
  K.upGateBlock = block(H, F * 2);
  if (draft) {
    K.projectBatch = gemm(H * 5, H);
    K.projectBlock = block(H * 5, H);
    K.contextBatch = gemm(H, KV * 2);
    K.contextBlock = block(H, KV * 2);
    K.markov = gemv(256, VOCAB);
  }
  K.silu =
    prefix +
    bind(0, "input:array<f32>") +
    bind(1, "out:array<f32>", "read_write") +
    `@compute @workgroup_size(256)fn main(@builtin(global_invocation_id)g:vec3<u32>){if(g.x>=${F}u||g.y>=step.count){return;}let x=input[g.y*${F * 2}u+g.x];out[g.y*${F}u+g.x]=(x/(1.0+exp(-x)))*input[g.y*${F * 2}u+${F}u+g.x];}`;
  K.rope =
    prefix +
    bind(0, "qkv:array<f32>") +
    bind(1, "queries:array<f32>", "read_write") +
    bind(2, "keyStage:array<f32>", "read_write") +
    bind(3, "valueStage:array<f32>", "read_write") +
    `
 @compute @workgroup_size(128)fn main(@builtin(global_invocation_id)g:vec3<u32>){let i=g.x*2u;let t=g.y;if(i>=${H + KV}u||t>=step.count){return;}
 let base=t*${H + 2 * KV}u;let angle=f32(step.position+t)*pow(5000000.0,-f32(i%128u)/128.0);let c=cos(angle);let s=sin(angle);let a=qkv[base+i];let b=qkv[base+i+1u];
 if(i<${H}u){queries[t*${H}u+i]=a*c-b*s;queries[t*${H}u+i+1u]=a*s+b*c;}
 else{let p=t*${KV}u+i-${H}u;keyStage[p]=a*c-b*s;keyStage[p+1u]=a*s+b*c;valueStage[p]=qkv[base+i+${KV}u];valueStage[p+1u]=qkv[base+i+${KV}u+1u];}}
 `;
  K.store =
    prefix +
    bind(0, "keyStage:array<f32>") +
    bind(1, "valueStage:array<f32>") +
    bind(2, `keys:array<${storageType}>`, "read_write") +
    bind(3, `values:array<${storageType}>`, "read_write") +
    `
 var<workgroup> ranges:array<f32,8>;
 @compute @workgroup_size(8)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let i=g.x*32u+l.x*4u;let t=g.y;let which=g.z;let b=t*${KV}u+i;var v:vec4<f32>;
 if(which==0u){v=vec4<f32>(keyStage[b],keyStage[b+1u],keyStage[b+2u],keyStage[b+3u]);}else{v=vec4<f32>(valueStage[b],valueStage[b+1u],valueStage[b+2u],valueStage[b+3u]);}
 let slot=(step.position+t)*${KV}u+i;
 ${
   kv === "f32"
     ? `if(which==0u){keys[slot]=v.x;keys[slot+1u]=v.y;keys[slot+2u]=v.z;keys[slot+3u]=v.w;}else{values[slot]=v.x;values[slot+1u]=v.y;values[slot+2u]=v.z;values[slot+3u]=v.w;}`
     : kv === "f16"
       ? `if(which==0u){keys[slot>>1u]=pack2x16float(v.xy);keys[(slot>>1u)+1u]=pack2x16float(v.zw);}else{values[slot>>1u]=pack2x16float(v.xy);values[(slot>>1u)+1u]=pack2x16float(v.zw);}`
       : `
 ranges[l.x]=max(max(abs(v.x),abs(v.y)),max(abs(v.z),abs(v.w)));workgroupBarrier();for(var s=4u;s>0u;s>>=1u){if(l.x<s){ranges[l.x]=max(ranges[l.x],ranges[l.x+s]);}workgroupBarrier();}
 ${kv === "q8_0" ? `let scale=max(ranges[0],1e-12);let at=(slot>>5u)*9u;if(l.x==0u){if(which==0u){keys[at]=bitcast<u32>(scale);}else{values[at]=bitcast<u32>(scale);}}let word=pack4x8snorm(v/scale);if(which==0u){keys[at+1u+l.x]=word;}else{values[at+1u+l.x]=word;}` : `let scale=max(ranges[0]/7.0,1e-12);let q=vec4<u32>(clamp(round(v/scale)+7.0,vec4<f32>(0.0),vec4<f32>(14.0)));workgroupBarrier();ranges[l.x]=bitcast<f32>(q.x|(q.y<<4u)|(q.z<<8u)|(q.w<<12u));workgroupBarrier();let at=(slot>>5u)*5u;if(l.x==0u){if(which==0u){keys[at]=bitcast<u32>(scale);}else{values[at]=bitcast<u32>(scale);}}if((l.x&1u)==0u){let word=bitcast<u32>(ranges[l.x])|(bitcast<u32>(ranges[l.x+1u])<<16u);if(which==0u){keys[at+1u+(l.x>>1u)]=word;}else{values[at+1u+(l.x>>1u)]=word;}}`}`
 }}
 `;
  // Four 32-lane teams each accumulate an independent key range, then merge.
  K.attention =
    prefix +
    bind(0, "queries:array<f32>") +
    bind(1, `keys:array<${storageType}>`) +
    bind(2, `values:array<${storageType}>`) +
    bind(3, "partial:array<f32>", "read_write") +
    bind(4, "md:array<vec2<f32>>", "read_write") +
    readers +
    sumDecl(subgroups, 128) +
    `
 var<workgroup> accs:array<vec4<f32>,128>;var<workgroup> stats:array<vec2<f32>,4>;
 @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let team=l.x>>5u;let lane=l.x&31u;let head=g.y;let token=g.z;let pos=step.position+token;let begin=g.x*step.chunk;let end=min(pos+1u,begin+step.chunk);let qb=token*${H}u+head*128u+lane*4u;
 let q=vec4<f32>(queries[qb],queries[qb+1u],queries[qb+2u],queries[qb+3u]);var m=-3.402823e38;var den=0.0;var acc=vec4<f32>(0.0);let n=select(0u,(end-begin+3u)/4u,end>begin);
 for(var j=0u;j<n;j++){let t=begin+j*4u+team;let valid=t<end;let i=select(0u,t*${KV}u+(head/8u)*128u+lane*4u,valid);let s=teamSum(dot(q,readK(i)),l.x)*0.08838834764831845;
 if(valid){let nm=max(m,s);let a=exp(m-nm);let b=exp(s-nm);acc=acc*a+readV(i)*b;den=den*a+b;m=nm;}}
 accs[l.x]=acc;if(lane==0u){stats[team]=vec2<f32>(m,den);}workgroupBarrier();
 let mm=max(max(stats[0].x,stats[1].x),max(stats[2].x,stats[3].x));var result=vec4<f32>(0.0);var total=0.0;
 for(var j=0u;j<4u;j++){let scale=exp(stats[j].x-mm);result+=accs[j*32u+lane]*scale;total+=stats[j].y*scale;}
 if(team==0u){let dst=((token*16u+head)*${MAX_CHUNKS}u+g.x)*128u+lane*4u;partial[dst]=result.x;partial[dst+1u]=result.y;partial[dst+2u]=result.z;partial[dst+3u]=result.w;if(lane==0u){md[(token*16u+head)*${MAX_CHUNKS}u+g.x]=vec2<f32>(mm,total);}}}
 `;
  K.combine =
    prefix +
    bind(0, "partial:array<f32>") +
    bind(1, "md:array<vec2<f32>>") +
    bind(2, "out:array<f32>", "read_write") +
    `
 @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let base=(g.y*16u+g.x)*${MAX_CHUNKS}u;var m=-3.402823e38;for(var j=0u;j<step.chunks;j++){m=max(m,md[base+j].x);}var a=0.0;var den=0.0;for(var j=0u;j<step.chunks;j++){let v=md[base+j];let s=exp(v.x-m);a+=partial[(base+j)*128u+l.x]*s;den+=v.y*s;}out[g.y*${H}u+g.x*128u+l.x]=a/max(den,1e-20);}
 `;
  const sorter = `var<workgroup> vals:array<f32,1024>;var<workgroup> ids:array<u32,1024>;
 fn sort(tid:u32){for(var k=2u;k<=1024u;k<<=1u){for(var j=k>>1u;j>0u;j>>=1u){for(var i=tid;i<1024u;i+=256u){let other=i^j;if(other>i){let descending=(i&k)==0u;let greater=vals[i]>vals[other]||(vals[i]==vals[other]&&ids[i]<ids[other]);if(greater!=descending){let v=vals[i];vals[i]=vals[other];vals[other]=v;let id=ids[i];ids[i]=ids[other];ids[other]=id;}}}workgroupBarrier();}}}
 `;
  for (const keep of [64, 256]) {
    K["sortFirst" + keep] =
      prefix +
      bind(0, "logits:array<f32>") +
      bind(1, "out:array<vec2<u32>>", "read_write") +
      bind(2, "tokens:array<u32>") +
      `@group(1) @binding(3) var<uniform> penalty:vec4<f32>;` +
      sorter +
      `
 @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){for(var i=l.x;i<1024u;i+=256u){let id=g.x*1024u+i;var v=-3.402823e38;if(id<${VOCAB}u){v=logits[id];if(penalty.x!=1.0){let start=select(0u,step.position-63u,step.position>63u);for(var p=start;p<=step.position;p++){if(tokens[p]==id){v=select(v/penalty.x,v*penalty.x,v<0.0);break;}}}}vals[i]=v;ids[i]=id;}workgroupBarrier();sort(l.x);if(l.x<${keep}u){out[g.x*${keep}u+l.x]=vec2<u32>(bitcast<u32>(vals[l.x]),ids[l.x]);}}
 `;
    K["sortMerge" + keep] =
      prefix +
      bind(0, "input:array<vec2<u32>>") +
      bind(1, "out:array<vec2<u32>>", "read_write") +
      `@group(1) @binding(2) var<uniform> size:vec4<u32>;` +
      sorter +
      `
 @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){for(var i=l.x;i<1024u;i+=256u){let at=g.x*1024u+i;var v=vec2<u32>(bitcast<u32>(-3.402823e38),${VOCAB}u);if(at<size.x){v=input[at];}vals[i]=bitcast<f32>(v.x);ids[i]=v.y;}workgroupBarrier();sort(l.x);if(l.x<${keep}u){out[g.x*${keep}u+l.x]=vec2<u32>(bitcast<u32>(vals[l.x]),ids[l.x]);}}
 `;
  }
  K.saveCandidates =
    prefix +
    bind(0, "input:array<vec2<u32>>") +
    bind(1, "out:array<vec2<u32>>", "read_write") +
    `@compute @workgroup_size(64)fn main(@builtin(global_invocation_id)g:vec3<u32>){out[g.x]=input[g.x];}`;
  K.sample =
    prefix +
    bind(0, "sorted:array<vec2<u32>>") +
    bind(1, "tokens:array<u32>", "read_write") +
    `
 struct Sampler{temperature:f32,topP:f32,topK:u32,seed:u32};@group(1)@binding(2)var<uniform> cfg:Sampler;
 fn hash(x:u32)->u32{var v=x*747796405u+2891336453u;v=((v>>((v>>28u)+4u))^v)*277803737u;return(v>>22u)^v;}
 @compute @workgroup_size(1)fn main(){var token=sorted[0].y;if(cfg.temperature>0.0){let mx=bitcast<f32>(sorted[0].x);var sum=0.0;for(var i=0u;i<cfg.topK;i++){sum+=exp((bitcast<f32>(sorted[i].x)-mx)/cfg.temperature);}var retained=0.0;var count=0u;for(var i=0u;i<cfg.topK;i++){retained+=exp((bitcast<f32>(sorted[i].x)-mx)/cfg.temperature);count++;if(retained>=cfg.topP*sum){break;}}let random=(f32(hash(cfg.seed^(step.position*2654435761u)))+0.5)/4294967296.0;let threshold=random*retained;var cum=0.0;for(var i=0u;i<count;i++){cum+=exp((bitcast<f32>(sorted[i].x)-mx)/cfg.temperature);token=sorted[i].y;if(cum>=threshold){break;}}}tokens[step.position+1u]=token;}
 `;
  if (smallWorkgroups) {
    K.down = K.down
      .replace(/var<workgroup> sharedX:[^;]+;/, "")
      .replace(/for\(var i=l.x;i<1536u;[\s\S]*?workgroupBarrier\(\);/, "")
      .replace(
        "let v=sharedX[b*8u+w];",
        "let at=(b*8u+w)*4u;let v=vec4<f32>(input[at],input[at+1u],input[at+2u],input[at+3u]);",
      );
  }
  return K;
}

// Small binding limits: keep vocabulary weights in four independently bound parts.
export function shardVocabularyKernel(code, embedding = false) {
  const binding = embedding ? 0 : 1;
  code = code.replace(/@binding\((\d+)\)/g, (match, n) =>
    Number(n) > binding ? `@binding(${Number(n) + 3})` : match,
  );
  const declaration = `@group(1) @binding(${binding}) var<storage,read> weights:array<u32>;`;
  code = code.replace(
    declaration,
    Array.from(
      { length: 4 },
      (_, i) =>
        `@group(1) @binding(${binding + i}) var<storage,read> weights${i}:array<u32>;`,
    ).join("\n") +
      `
 fn weightWord(i:u32)->u32{let part=i/17756160u;let at=i%17756160u;switch part{case 0u:{return weights0[at];}case 1u:{return weights1[at];}case 2u:{return weights2[at];}default:{return weights3[at];}}}`,
  );
  return code.replace(/weights\[([^\]]+)\]/g, "weightWord($1)");
}