nx_vocops.nx source
↩ module page · 111 lines · 5359 B
1// nx_vocops.nx -- the SOVEREIGN NEURAL-VOCODER OP LIBRARY (no main; imported by the HiFi-GAN generator forward).
2// Consolidates the 1-D neural ops verified as standalone bricks (supersedes nx_conv1d.nx + nx_conv_transpose1d.nx):
3// conv1d -- dilated 1-D convolution (conv_pre, MRF, conv_post)
4// conv_transpose1d -- the upsampler (ConvTranspose1d), 4 stages mel->audio-rate
5// leaky_relu -- the activation every conv is wrapped in (HiFi-GAN slope 0.1)
6// resblock1_1dil -- ONE dilation branch of a HiFi-GAN ResBlock: x + conv(leaky(conv(leaky(x)))) (residual)
7// Same f32/SSE substrate as the FIXED Qwen forward. Tested byte-exact by nx_vocops_test. license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_f32.nx"
10import "nx_f32_div.nx"
11import "nx_f32_cvt.nx"
12
13// out[co][t] = bias[co] + sum_ci sum_k weight[co][ci][k] * inp[ci][ t*stride - pad + k*dil ]. layouts:
14// 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 { acc = nx_f32_add(acc, nx_f32_mul(w[(co*C_in+ci)*K + k], inp[ci*L + ipos])) } }
29 k = k + 1
30 }
31 ci = ci + 1
32 }
33 out[co*L_out + t] = acc
34 t = t + 1
35 }
36 co = co + 1
37 }
38 return L_out
39}
40
41// ConvTranspose1d: input[C_in,L_in] * weight[C_in,C_out,K] + bias[C_out] -> output[C_out,L_out] (scatter).
42// L_out = (L_in-1)*stride - 2*pad + dil*(K-1) + out_pad + 1.
43func conv_transpose1d(inp: *i64, w: *i64, b: *i64, out: *i64,
44 C_in: i64, L_in: i64, C_out: i64, K: i64, stride: i64, pad: i64, out_pad: i64, dil: i64) -> i64 {
45 let L_out: i64 = (L_in-1)*stride - 2*pad + dil*(K-1) + out_pad + 1
46 var co: i64 = 0
47 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 }
48 var ci: i64 = 0
49 while ci < C_in {
50 var i: i64 = 0
51 while i < L_in {
52 let xv: i64 = inp[ci*L_in + i]
53 var c2: i64 = 0
54 while c2 < C_out {
55 var k: i64 = 0
56 while k < K {
57 let opos: i64 = i*stride - pad + k*dil
58 if opos >= 0 { if opos < L_out { let idx: i64 = c2*L_out + opos; out[idx] = nx_f32_add(out[idx], nx_f32_mul(w[(ci*C_out+c2)*K + k], xv)) } }
59 k = k + 1
60 }
61 c2 = c2 + 1
62 }
63 i = i + 1
64 }
65 ci = ci + 1
66 }
67 return L_out
68}
69
70// LeakyReLU elementwise, in place over n f32 values (slope raw bits, e.g. 0.1 = 0x3DCCCCCD).
71func leaky_relu(x: i64, slope: i64) -> i64 { if ((x>>31)&1)==1 { return nx_f32_mul(x, slope) } return x }
72func leaky_relu_vec(v: *i64, n: i64, slope: i64) -> i64 { var i: i64=0; while i<n { v[i]=leaky_relu(v[i], slope); i=i+1 } return 0 }
73
74// WEIGHT-NORM reconstruction: HiFi-GAN convs are stored as weight_v + weight_g (PyTorch weight_norm, dim=0).
75// Effective weight W[co] = g[co] * V[co] / ||V[co]|| where the L2 norm is over all dims but the output channel.
76// v/w laid out [C_out, inner] (inner = C_in*K); g is [C_out]. Fills w_out. This turns checkpoint tensors into
77// the conv weights conv1d/conv_transpose1d consume -- required to load a real HiFi-GAN.
78func weight_norm_apply(v: *i64, g: *i64, w_out: *i64, C_out: i64, inner: i64) -> i64 {
79 var co: i64 = 0
80 while co < C_out {
81 var ss: i64 = 0
82 var j: i64 = 0
83 while j < inner { let x: i64 = v[co*inner + j]; ss = nx_f32_add(ss, nx_f32_mul(x, x)); j = j + 1 }
84 let norm: i64 = nx_f32_sqrt(ss)
85 let scale: i64 = nx_f32_div(g[co], norm) // g / ||V|| (broadcast over inner)
86 j = 0
87 while j < inner { w_out[co*inner + j] = nx_f32_mul(v[co*inner + j], scale); j = j + 1 }
88 co = co + 1
89 }
90 return 0
91}
92
93// ONE dilation branch of a HiFi-GAN ResBlock (single channel C, length L, same-length convs):
94// xt = leaky(x); xt = conv1d(xt, w1, dil); xt = leaky(xt); xt = conv1d(xt, w2, dil=1); out = x + xt (residual).
95// pad chosen for same-length (K odd). Scratch t1/t2 supplied by caller (size C*L each).
96func resblock1_1dil(x: *i64, w1: *i64, b1: *i64, w2: *i64, b2: *i64, out: *i64,
97 C: i64, L: i64, K: i64, dil: i64, slope: i64, t1: *i64, t2: *i64) -> i64 {
98 // t1 = leaky(x)
99 var i: i64=0; while i<C*L { t1[i]=leaky_relu(x[i], slope); i=i+1 }
100 // t2 = conv1d(t1, w1, dil, pad=dil*(K-1)/2) -- same length
101 let p1: i64 = dil*(K-1)/2
102 conv1d(t1, w1, b1, t2, C, L, C, K, 1, p1, dil)
103 // t1 = leaky(t2)
104 i=0; while i<C*L { t1[i]=leaky_relu(t2[i], slope); i=i+1 }
105 // t2 = conv1d(t1, w2, dil=1, pad=(K-1)/2)
106 let p2: i64 = (K-1)/2
107 conv1d(t1, w2, b2, t2, C, L, C, K, 1, p2, 1)
108 // out = x + t2 (residual)
109 i=0; while i<C*L { out[i]=nx_f32_add(x[i], t2[i]); i=i+1 }
110 return 0
111}