nx_conv_transpose1d.nx source
↩ module page · 83 lines · 6036 B
1// nx_conv_transpose1d.nx -- BRICK 2 of the neural vocoder forward: transposed 1-D convolution (ConvTranspose1d),
2// the UPSAMPLER at the heart of HiFi-GAN's generator (4 stages upsample the 80-band mel to raw audio rate). Plus
3// leaky_relu (the activation every conv is wrapped in). Sovereign f32, same substrate as the FIXED Qwen forward.
4// ConvTranspose1d = the transpose/"scatter" of conv: each input sample spreads (weight * value) into a K-wide
5// output window at stride spacing. weight layout [C_in, C_out, K]. Known-answer tested byte-exact. 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// LeakyReLU: x>=0 -> x ; x<0 -> slope*x. HiFi-GAN uses slope 0.1. (sign bit test; -0.0 * slope = 0, fine.)
14func leaky_relu(x: i64, slope: i64) -> i64 { if ((x>>31)&1)==1 { return nx_f32_mul(x, slope) } return x }
15
16// ConvTranspose1d: input[C_in,L_in], weight[C_in,C_out,K], bias[C_out], stride/pad/out_pad/dil -> output[C_out,L_out].
17// L_out = (L_in-1)*stride - 2*pad + dil*(K-1) + out_pad + 1. Each input sample SCATTERS into the output.
18func conv_transpose1d(inp: *i64, w: *i64, b: *i64, out: *i64,
19 C_in: i64, L_in: i64, C_out: i64, K: i64, stride: i64, pad: i64, out_pad: i64, dil: i64) -> i64 {
20 let L_out: i64 = (L_in-1)*stride - 2*pad + dil*(K-1) + out_pad + 1
21 // init every output to its channel bias
22 var co: i64 = 0
23 while co < C_out { var t: i64=0; while t < L_out { out[co*L_out + t] = b[co]; t=t+1 } co=co+1 }
24 // scatter each input sample into the output window
25 var ci: i64 = 0
26 while ci < C_in {
27 var i: i64 = 0
28 while i < L_in {
29 let xv: i64 = inp[ci*L_in + i]
30 var c2: i64 = 0
31 while c2 < C_out {
32 var k: i64 = 0
33 while k < K {
34 let opos: i64 = i*stride - pad + k*dil
35 if opos >= 0 { if opos < L_out {
36 let idx: i64 = c2*L_out + opos
37 out[idx] = nx_f32_add(out[idx], nx_f32_mul(w[(ci*C_out+c2)*K + k], xv))
38 } }
39 k = k + 1
40 }
41 c2 = c2 + 1
42 }
43 i = i + 1
44 }
45 ci = ci + 1
46 }
47 return L_out
48}
49
50func main() -> i64 {
51 cw("=== nx_conv_transpose1d -- BRICK 2: the HiFi-GAN upsampler (ConvTranspose1d) + leaky_relu ===\n" as *u8)
52 var pass: i64=0; var tot: i64=0
53
54 // ---- KAT 1: upsample-by-2, input [1,2], weight [1,1] (K=2), stride 2 -> [1,1,2,2] ----
55 let in1: *i64 = sys_mmap(8*8) as *i64; in1[0]=0x3F800000; in1[1]=0x40000000 // 1,2
56 let w1: *i64 = sys_mmap(8*8) as *i64; w1[0]=0x3F800000; w1[1]=0x3F800000 // [1,1]
57 let b0: *i64 = sys_mmap(8) as *i64; b0[0]=0
58 let o1: *i64 = sys_mmap(16*8) as *i64
59 let lo1: i64 = conv_transpose1d(in1, w1, b0, o1, 1, 2, 1, 2, 2, 0, 0, 1)
60 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 1 1 2 2)\n" as *u8)
61 tot=tot+1
62 if lo1==4 { if (o1[0]&0xFFFFFFFF)==0x3F800000 { if (o1[1]&0xFFFFFFFF)==0x3F800000 { if (o1[2]&0xFFFFFFFF)==0x40000000 { if (o1[3]&0xFFFFFFFF)==0x40000000 { pass=pass+1; cw("PASS T1 upsample-by-2 (stride scatter) byte-exact\n" as *u8) } else {cw("FAIL T1e\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)}
63
64 // ---- KAT 2: stride 2, K=4, pad 1 (HiFi-GAN style K=2*stride). input [1,2], weight [1,2,3,4] -> [2,5,8,6] ----
65 let w2: *i64 = sys_mmap(8*8) as *i64; w2[0]=0x3F800000; w2[1]=0x40000000; w2[2]=0x40400000; w2[3]=0x40800000 // 1,2,3,4
66 let o2: *i64 = sys_mmap(16*8) as *i64
67 let lo2: i64 = conv_transpose1d(in1, w2, b0, o2, 1, 2, 1, 4, 2, 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 40000000 40a00000 41000000 40c00000 = 2 5 8 6)\n" as *u8)
69 tot=tot+1
70 if lo2==4 { if (o2[0]&0xFFFFFFFF)==0x40000000 { if (o2[1]&0xFFFFFFFF)==0x40A00000 { if (o2[2]&0xFFFFFFFF)==0x41000000 { if (o2[3]&0xFFFFFFFF)==0x40C00000 { pass=pass+1; cw("PASS T2 strided/padded/wide-kernel upsample (HiFi-GAN style) byte-exact\n" as *u8) } else {cw("FAIL T2e\n" as *u8)} } else {cw("FAIL T2d\n" as *u8)} } else {cw("FAIL T2c\n" as *u8)} } else {cw("FAIL T2b\n" as *u8)} } else {cw("FAIL T2a L_out="); cn(lo2); cw("\n" as *u8)}
71
72 // ---- KAT 3: leaky_relu(slope 0.5): +2.0 -> 2.0 ; -2.0 -> -1.0 ----
73 let lp: i64 = 0x3F000000 // 0.5
74 let r_pos: i64 = leaky_relu(0x40000000, lp) // +2.0
75 let r_neg: i64 = leaky_relu(0xC0000000, lp) // -2.0
76 cw("KAT3 leaky_relu(2.0)="); chx(r_pos&0xFFFFFFFF); cw(" leaky_relu(-2.0,0.5)="); chx(r_neg&0xFFFFFFFF); cw(" (expect 40000000, bf800000 = 2.0, -1.0)\n" as *u8)
77 tot=tot+1
78 if (r_pos&0xFFFFFFFF)==0x40000000 { if (r_neg&0xFFFFFFFF)==0xBF800000 { pass=pass+1; cw("PASS T3 leaky_relu (identity on +, slope on -) byte-exact\n" as *u8) } else {cw("FAIL T3b\n" as *u8)} } else {cw("FAIL T3a\n" as *u8)}
79
80 cw("nx_conv_transpose1d pass="); cn(pass); cw("/"); cn(tot)
81 if pass==tot { cw(" GREEN -- upsampler + activation byte-exact. Brick 2/6 of the HiFi-GAN forward (conv1d done; next: MRF resblock, tanh, weight-norm, graph).\n" as *u8); sys_exit(0); return 0 }
82 cw(" RED\n" as *u8); sys_exit(1); return 1
83}