nx_conv1d.nx source
↩ module page · 84 lines · 5863 B
1// nx_conv1d.nx -- the ATOM of the neural vocoder forward (HiFi-GAN / WaveNet are conv1d stacks). Sovereign f32
2// 1-D convolution: input [C_in, L] * weight [C_out, C_in, K] + bias [C_out] -> output [C_out, L_out], with
3// stride / padding / dilation (HiFi-GAN's MRF uses dilated convs; its upsampler uses transposed conv = the next
4// brick). This is the FIRST real brick of the neural-realism voice forward (same f32/SSE substrate as the FIXED
5// Qwen forward + Z-Image DiT). Known-answer tested byte-exact. No ML lib. license_tier: ORIGINAL expect_exit: 0
6import "nx_syscalls.nx"
7import "nx_f32.nx"
8
9func cw(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
10func cn(v: i64) -> i64 { let bb: *u8=sys_mmap(28); var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{bb[i]=t[k-1-i];i=i+1} sys_write(1,bb,k); return 0 }
11func chx(v: i64) -> i64 { let b: *u8=sys_mmap(8); var i: i64=0; while i<8 { let n: i64=(v>>((7-i)*4))&0xF; if n<10 {b[i]=(48+n) as u8} else {b[i]=(87+n) as u8} i=i+1 } sys_write(1,b,8); return 0 }
12
13// out[co][t] = bias[co] + sum_ci sum_k weight[co][ci][k] * inp[ci][ t*stride - pad + k*dil ] (0 outside [0,L))
14// layouts: inp[ci*L + t], weight[(co*C_in+ci)*K + k], bias[co], out[co*L_out + t]. returns L_out.
15func conv1d(inp: *i64, w: *i64, b: *i64, out: *i64,
16 C_in: i64, L: i64, C_out: i64, K: i64, stride: i64, pad: i64, dil: i64) -> i64 {
17 let L_out: i64 = (L + 2*pad - dil*(K-1) - 1)/stride + 1
18 var co: i64 = 0
19 while co < C_out {
20 var t: i64 = 0
21 while t < L_out {
22 var acc: i64 = b[co]
23 var ci: i64 = 0
24 while ci < C_in {
25 var k: i64 = 0
26 while k < K {
27 let ipos: i64 = t*stride - pad + k*dil
28 if ipos >= 0 { if ipos < L {
29 acc = nx_f32_add(acc, nx_f32_mul(w[(co*C_in+ci)*K + k], inp[ci*L + ipos]))
30 } }
31 k = k + 1
32 }
33 ci = ci + 1
34 }
35 out[co*L_out + t] = acc
36 co = co // (no-op to keep the loop var clear)
37 t = t + 1
38 }
39 co = co + 1
40 }
41 return L_out
42}
43
44func main() -> i64 {
45 cw("=== nx_conv1d -- sovereign f32 1-D convolution (the atom of the neural vocoder forward) ===\n" as *u8)
46 var pass: i64=0; var tot: i64=0
47
48 // ---- KAT 1: single channel, K=3, pad=1, box filter [1,1,1] over [1,2,3,4,5] -> [3,6,9,12,9] ----
49 let inp1: *i64 = sys_mmap(8*8) as *i64
50 inp1[0]=0x3F800000; inp1[1]=0x40000000; inp1[2]=0x40400000; inp1[3]=0x40800000; inp1[4]=0x40A00000 // 1..5
51 let w1: *i64 = sys_mmap(4*8) as *i64
52 w1[0]=0x3F800000; w1[1]=0x3F800000; w1[2]=0x3F800000 // [1,1,1]
53 let b1: *i64 = sys_mmap(8) as *i64; b1[0]=0 // bias 0.0
54 let o1: *i64 = sys_mmap(8*8) as *i64
55 let lo1: i64 = conv1d(inp1, w1, b1, o1, 1, 5, 1, 3, 1, 1, 1)
56 cw("KAT1 L_out="); cn(lo1); cw(" out=["); var i: i64=0; while i<lo1 { chx(o1[i]&0xFFFFFFFF); if i<lo1-1 {cw(" ")} i=i+1 } cw("] (expect 40400000 40c00000 41100000 41400000 41100000 = 3 6 9 12 9)\n" as *u8)
57 tot=tot+1
58 if lo1==5 { if (o1[0]&0xFFFFFFFF)==0x40400000 { if (o1[3]&0xFFFFFFFF)==0x41400000 { if (o1[4]&0xFFFFFFFF)==0x41100000 { pass=pass+1; cw("PASS T1 box-filter conv (padding + edge handling) byte-exact\n" as *u8) } else {cw("FAIL T1d\n" as *u8)} } else {cw("FAIL T1c\n" as *u8)} } else {cw("FAIL T1b\n" as *u8)} } else {cw("FAIL T1a L_out="); cn(lo1); cw("\n" as *u8)}
59
60 // ---- KAT 2: 2 in-channels, K=1 -> channel summation. ch0=[1,2,3] ch1=[10,20,30], w=[1,2] -> [21,42,63] ----
61 let inp2: *i64 = sys_mmap(8*8) as *i64
62 inp2[0]=0x3F800000; inp2[1]=0x40000000; inp2[2]=0x40400000 // ch0 = 1,2,3
63 inp2[3]=0x41200000; inp2[4]=0x41A00000; inp2[5]=0x41F00000 // ch1 = 10,20,30
64 let w2: *i64 = sys_mmap(4*8) as *i64; w2[0]=0x3F800000; w2[1]=0x40000000 // [1, 2]
65 let b2: *i64 = sys_mmap(8) as *i64; b2[0]=0
66 let o2: *i64 = sys_mmap(8*8) as *i64
67 let lo2: i64 = conv1d(inp2, w2, b2, o2, 2, 3, 1, 1, 1, 0, 1)
68 cw("KAT2 L_out="); cn(lo2); cw(" out=["); i=0; while i<lo2 { chx(o2[i]&0xFFFFFFFF); if i<lo2-1 {cw(" ")} i=i+1 } cw("] (expect 41a80000 42280000 427c0000 = 21 42 63)\n" as *u8)
69 tot=tot+1
70 if lo2==3 { if (o2[0]&0xFFFFFFFF)==0x41A80000 { if (o2[2]&0xFFFFFFFF)==0x427C0000 { pass=pass+1; cw("PASS T2 multi-channel summation byte-exact\n" as *u8) } else {cw("FAIL T2c\n" as *u8)} } else {cw("FAIL T2b\n" as *u8)} } else {cw("FAIL T2a\n" as *u8)}
71
72 // ---- KAT 3: dilation=2, K=3, pad=2 (HiFi-GAN MRF uses dilated convs). [1,2,3,4,5] w=[1,1,1] ----
73 // out[t] = inp[t-2] + inp[t] + inp[t+2] ; L_out = (5+4-2*2-1)/1+1 = 5
74 let o3: *i64 = sys_mmap(8*8) as *i64
75 let lo3: i64 = conv1d(inp1, w1, b1, o3, 1, 5, 1, 3, 1, 2, 2)
76 // t=0: inp[-2]+inp[0]+inp[2]=0+1+3=4 ; t=2: inp[0]+inp[2]+inp[4]=1+3+5=9 ; t=4: inp[2]+inp[4]+inp[6]=3+5+0=8
77 cw("KAT3 (dilated) out=["); i=0; while i<lo3 { chx(o3[i]&0xFFFFFFFF); if i<lo3-1 {cw(" ")} i=i+1 } cw("] (expect t0=4 t2=9 t4=8)\n" as *u8)
78 tot=tot+1
79 if lo3==5 { if (o3[0]&0xFFFFFFFF)==0x40800000 { if (o3[2]&0xFFFFFFFF)==0x41100000 { if (o3[4]&0xFFFFFFFF)==0x41000000 { pass=pass+1; cw("PASS T3 DILATED conv (the HiFi-GAN MRF op) byte-exact\n" as *u8) } else {cw("FAIL T3d\n" as *u8)} } else {cw("FAIL T3c\n" as *u8)} } else {cw("FAIL T3b\n" as *u8)} } else {cw("FAIL T3a\n" as *u8)}
80
81 cw("nx_conv1d pass="); cn(pass); cw("/"); cn(tot)
82 if pass==tot { cw(" GREEN -- conv1d (box/multi-channel/dilated) byte-exact. First brick of the neural vocoder forward.\n" as *u8); sys_exit(0); return 0 }
83 cw(" RED\n" as *u8); sys_exit(1); return 1
84}