GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 100.0% 26 / 0 / 26
Functions: 100.0% 5 / 0 / 5
Branches: -% 0 / 0 / 0
include/brotensor/ops/conv1d.h
Line Branch Exec Source
1 #pragma once
2
3 // brotensor ops/conv1d.h — 1D convolution family (audio): conv1d, pad1d, conv_transpose1d, causal_conv1d.
4
5 #include "../tensor.h"
6 #include "conv.h"
7 #include <cstdint>
8
9 namespace brotensor {
10
11
12 // ─── 1D convolution family (audio) ─────────────────────────────────────────
13 //
14 // The audio counterpart of the conv2d family — the building block of
15 // WaveNet / Conformer / vocoder stacks.
16 //
17 // Layout (NCL): a 1D-conv activation is (N, C*L) — N signals folded into rows,
18 // each row a flat C-major / L-minor buffer X[(n*C+c)*L + l] (the NCHW
19 // convention with the height axis dropped). Weights are OIL:
20 // Wt[(c_out*(C_in/groups)+c_in_local)*kL + kl].
21 //
22 // conv1d, its three backward halves, and conv1d_int8w_fp16 are header-only
23 // inline wrappers over the conv2d ops (a 1D conv is a 2D conv with H=kH=1), so
24 // every backend that implements conv2d gets conv1d for free; conv1d_int8w_fp16
25 // therefore throws on the CPU backend (conv2d_int8w is a null CPU slot).
26 // causal_conv1d is likewise a wrapper (left-pad, then a valid conv1d).
27 // pad1d, conv_transpose1d, and causal_conv1d_update are genuinely new ops with
28 // their own vtable rows, implemented on all three backends (CPU / CUDA / Metal).
29
30 // Pad the length axis of an NCL tensor by pad_left / pad_right samples — the
31 // temporal analogue of an image pad (causal-conv left padding, "same" padding,
32 // reflect padding). `mode`: 0 zero, 1 reflect (mirror without repeating the
33 // edge sample; requires pad < L), 2 replicate (clamp to the edge sample).
34 // X: (N, C*L). Y: (N, C*(L+pad_left+pad_right)), resized + dtype-set to X.
35 void pad1d_forward(const Tensor& X, int N, int C, int L,
36 int pad_left, int pad_right, int mode, Tensor& Y);
37
38
39 // Backward (adjoint) of pad1d: each input sample sums the gradients of the
40 // output samples that read it. dX overwritten (resized + dtype-set to dY).
41 // dY: (N, C*(L+pad_left+pad_right)). dX: (N, C*L).
42 // N, C, L and the pad / mode args match the forward call.
43 void pad1d_backward(const Tensor& dY, int N, int C, int L,
44 int pad_left, int pad_right, int mode, Tensor& dX);
45
46
47 // 1D convolution, NCL. Header-only wrapper over conv2d_forward (H=kH=1).
48 // X: (N,C_in*L). Wt: (C_out, (C_in/groups)*kL) OIL. bias: (C_out,1) or null.
49 // Y: (N,C_out*L_out); L_out = (L+2*padding-dilation*(kL-1)-1)/stride + 1.
50 7 inline void conv1d(const Tensor& X, const Tensor& Wt, const Tensor* bias,
51 int N, int C_in, int L, int C_out, int kL,
52 int stride, int padding, int dilation, int groups,
53 Tensor& Y) {
54 14 conv2d_forward(X, Wt, bias, N, C_in, /*H=*/1, /*W=*/L, C_out,
55 7 /*kH=*/1, /*kW=*/kL, /*stride_h=*/1, /*stride_w=*/stride,
56 7 /*pad_h=*/0, /*pad_w=*/padding, /*dil_h=*/1, /*dil_w=*/dilation,
57 7 groups, Y);
58 7 }
59
60 // Convenience overload: groups defaults to 1.
61 inline void conv1d(const Tensor& X, const Tensor& Wt, const Tensor* bias,
62 int N, int C_in, int L, int C_out, int kL,
63 int stride, int padding, int dilation, Tensor& Y) {
64 conv1d(X, Wt, bias, N, C_in, L, C_out, kL, stride, padding, dilation,
65 /*groups=*/1, Y);
66 }
67
68
69 // conv1d backward w.r.t. input — wrapper over conv2d_backward_input (H=kH=1).
70 // dX overwritten.
71 3 inline void conv1d_backward_input(const Tensor& Wt, const Tensor& dY,
72 int N, int C_in, int L, int C_out, int kL,
73 int stride, int padding, int dilation,
74 int groups, Tensor& dX) {
75 6 conv2d_backward_input(Wt, dY, N, C_in, /*H=*/1, /*W=*/L, C_out,
76 3 /*kH=*/1, /*kW=*/kL, 1, stride, 0, padding,
77 3 1, dilation, groups, dX);
78 3 }
79
80 inline void conv1d_backward_input(const Tensor& Wt, const Tensor& dY,
81 int N, int C_in, int L, int C_out, int kL,
82 int stride, int padding, int dilation,
83 Tensor& dX) {
84 conv1d_backward_input(Wt, dY, N, C_in, L, C_out, kL, stride, padding,
85 dilation, /*groups=*/1, dX);
86 }
87
88
89 // conv1d backward w.r.t. weight — wrapper over conv2d_backward_weight.
90 // dWt accumulated — caller zeros.
91 3 inline void conv1d_backward_weight(const Tensor& X, const Tensor& dY,
92 int N, int C_in, int L, int C_out, int kL,
93 int stride, int padding, int dilation,
94 int groups, Tensor& dWt) {
95 6 conv2d_backward_weight(X, dY, N, C_in, /*H=*/1, /*W=*/L, C_out,
96 3 /*kH=*/1, /*kW=*/kL, 1, stride, 0, padding,
97 3 1, dilation, groups, dWt);
98 3 }
99
100 inline void conv1d_backward_weight(const Tensor& X, const Tensor& dY,
101 int N, int C_in, int L, int C_out, int kL,
102 int stride, int padding, int dilation,
103 Tensor& dWt) {
104 conv1d_backward_weight(X, dY, N, C_in, L, C_out, kL, stride, padding,
105 dilation, /*groups=*/1, dWt);
106 }
107
108
109 // conv1d backward w.r.t. bias — wrapper over conv2d_backward_bias.
110 // dB accumulated — caller zeros.
111 3 inline void conv1d_backward_bias(const Tensor& dY, int N, int C_out, int L_out,
112 Tensor& dB) {
113 3 conv2d_backward_bias(dY, N, C_out, /*H_out=*/1, /*W_out=*/L_out, dB);
114 3 }
115
116
117 // W8A16 1D convolution. Header-only wrapper over conv2d_int8w_fp16_forward
118 // (H=kH=1). FP16 activations, INT8 per-output-row weights. Throws on the CPU
119 // backend (no CPU W8A16 conv slot).
120 // X: (N,C_in*L) FP16. W_int8: (C_out, (C_in/groups)*kL) INT8 OIL.
121 // scales: (C_out,1) FP32. bias: (C_out,1) FP16 or null. Y: (N,C_out*L_out) FP16.
122 inline void conv1d_int8w_fp16(const Tensor& X, const Tensor& W_int8,
123 const Tensor& scales, const Tensor* bias,
124 int N, int C_in, int L, int C_out, int kL,
125 int stride, int padding, int dilation, int groups,
126 Tensor& Y) {
127 conv2d_int8w_fp16_forward(X, W_int8, scales, bias, N, C_in, /*H=*/1,
128 /*W=*/L, C_out, /*kH=*/1, /*kW=*/kL,
129 /*stride_h=*/1, /*stride_w=*/stride,
130 /*pad_h=*/0, /*pad_w=*/padding,
131 /*dil_h=*/1, /*dil_w=*/dilation, groups, Y);
132 }
133
134 inline void conv1d_int8w_fp16(const Tensor& X, const Tensor& W_int8,
135 const Tensor& scales, const Tensor* bias,
136 int N, int C_in, int L, int C_out, int kL,
137 int stride, int padding, int dilation, Tensor& Y) {
138 conv1d_int8w_fp16(X, W_int8, scales, bias, N, C_in, L, C_out, kL, stride,
139 padding, dilation, /*groups=*/1, Y);
140 }
141
142
143 // Causal 1D convolution. Header-only wrapper: left-pad the length axis by
144 // dilation*(kL-1) (zero), then run a valid (padding=0) conv1d. Output length
145 // equals L when stride==1; every output sample depends only on inputs at or
146 // before its position.
147 // X: (N,C_in*L). Wt: (C_out, (C_in/groups)*kL) OIL. bias: (C_out,1) or null.
148 // scratch: caller-owned, resized to (N, C_in*(L+dilation*(kL-1))) and
149 // overwritten — keeps the wrapper allocation-free across calls.
150 // Y: (N,C_out*L_out), resized by conv1d.
151 4 inline void causal_conv1d(const Tensor& X, const Tensor& Wt, const Tensor* bias,
152 int N, int C_in, int L, int C_out, int kL,
153 int stride, int dilation, int groups,
154 Tensor& scratch, Tensor& Y) {
155 4 const int pad_left = dilation * (kL - 1);
156 8 pad1d_forward(X, N, C_in, L, pad_left, /*pad_right=*/0, /*mode=*/0,
157 4 scratch);
158 8 conv1d(scratch, Wt, bias, N, C_in, L + pad_left, C_out, kL, stride,
159 4 /*padding=*/0, dilation, groups, Y);
160 4 }
161
162 inline void causal_conv1d(const Tensor& X, const Tensor& Wt, const Tensor* bias,
163 int N, int C_in, int L, int C_out, int kL,
164 int stride, int dilation, Tensor& scratch, Tensor& Y) {
165 causal_conv1d(X, Wt, bias, N, C_in, L, C_out, kL, stride, dilation,
166 /*groups=*/1, scratch, Y);
167 }
168
169
170 // 1D transposed convolution, NCL — the upsampling primitive of neural vocoders
171 // (HiFi-GAN, EnCodec/DAC decoders). A genuinely new kernel, CPU FP32-only.
172 // L_out = (L-1)*stride - 2*padding + dilation*(kL-1) + output_padding + 1.
173 // output_padding (< stride) disambiguates the L_out values that map to one L
174 // under a strided forward conv (torch's ConvTranspose1d arg).
175 // Weight layout is input-channel-major: Wt (C_in, (C_out/groups)*kL),
176 // Wt[(c_in*(C_out/groups)+c_out_local)*kL + kl]. groups divides C_in and C_out;
177 // groups==C_in==C_out is depthwise transposed conv.
178 // Forward (scatter): each X[n,c_in,l] is scattered, per kernel tap kl, into
179 // output position l_out = l*stride - padding + kl*dilation.
180 // X: (N,C_in*L). Wt: (C_in,(C_out/groups)*kL). bias: (C_out,1) or null.
181 // Y: (N,C_out*L_out), resized + dtype-set to match X.
182 void conv_transpose1d_forward(const Tensor& X, const Tensor& Wt,
183 const Tensor* bias,
184 int N, int C_in, int L, int C_out, int kL,
185 int stride, int padding, int output_padding,
186 int dilation, int groups, Tensor& Y);
187
188 // Convenience overload: groups defaults to 1.
189 inline void conv_transpose1d_forward(const Tensor& X, const Tensor& Wt,
190 const Tensor* bias,
191 int N, int C_in, int L, int C_out, int kL,
192 int stride, int padding,
193 int output_padding, int dilation,
194 Tensor& Y) {
195 conv_transpose1d_forward(X, Wt, bias, N, C_in, L, C_out, kL, stride,
196 padding, output_padding, dilation, /*groups=*/1, Y);
197 }
198
199
200 // conv_transpose1d backward w.r.t. input — the adjoint is a plain gather conv:
201 // dX[n,c_in,l] gathers dY over every tap and the group's output channels at
202 // l_out = l*stride - padding + kl*dilation. dX overwritten (resized +
203 // dtype-set to dY). All hyperparams match the forward call.
204 void conv_transpose1d_backward_input(const Tensor& Wt, const Tensor& dY,
205 int N, int C_in, int L, int C_out, int kL,
206 int stride, int padding,
207 int output_padding, int dilation,
208 int groups, Tensor& dX);
209
210 inline void conv_transpose1d_backward_input(const Tensor& Wt, const Tensor& dY,
211 int N, int C_in, int L, int C_out,
212 int kL, int stride, int padding,
213 int output_padding, int dilation,
214 Tensor& dX) {
215 conv_transpose1d_backward_input(Wt, dY, N, C_in, L, C_out, kL, stride,
216 padding, output_padding, dilation,
217 /*groups=*/1, dX);
218 }
219
220
221 // conv_transpose1d backward w.r.t. weight:
222 // dWt[c_in,c_out_local,kl] += sum_{n,l} X[n,c_in,l]*dY[n,c_out,l_out],
223 // l_out = l*stride - padding + kl*dilation (skipped when OOB).
224 // dWt accumulated — caller zeros. All hyperparams match the forward call.
225 void conv_transpose1d_backward_weight(const Tensor& X, const Tensor& dY,
226 int N, int C_in, int L, int C_out, int kL,
227 int stride, int padding,
228 int output_padding, int dilation,
229 int groups, Tensor& dWt);
230
231 inline void conv_transpose1d_backward_weight(const Tensor& X, const Tensor& dY,
232 int N, int C_in, int L, int C_out,
233 int kL, int stride, int padding,
234 int output_padding, int dilation,
235 Tensor& dWt) {
236 conv_transpose1d_backward_weight(X, dY, N, C_in, L, C_out, kL, stride,
237 padding, output_padding, dilation,
238 /*groups=*/1, dWt);
239 }
240
241
242 // conv_transpose1d backward w.r.t. bias:
243 // dB[c_out] += sum_{n,l_out} dY[n,c_out,l_out].
244 // dB accumulated — caller zeros.
245 void conv_transpose1d_backward_bias(const Tensor& dY, int N, int C_out,
246 int L_out, Tensor& dB);
247
248
249 // One streaming step of a causal depthwise 1D conv against a rolling state
250 // cache (in the spirit of kv_cache_append) — for autoregressive / streaming
251 // decoders. Forward-only, new vtable row, CPU FP32-only.
252 // Depthwise: C channels in and out, one length-kL filter per channel. With
253 // L_step new samples it produces L_step outputs:
254 // Y[n,c,t] = bias[c] + sum_{kl} W[c,kl] * buf[n,c,t+kl*dilation],
255 // buf = state[n,c,:] ++ X[n,c,:].
256 // `state` is updated in place to the last (kL-1)*dilation samples of buf, so a
257 // sequence of calls reproduces one full causal_conv1d over the concatenated
258 // input (caller zero-initialises state before the first step).
259 // X: (N,C*L_step) new samples. Wt: (C,kL) depthwise filter.
260 // bias: (C,1) or null. state: (N,C*(kL-1)*dilation) — read AND overwritten.
261 // Y: (N,C*L_step), resized + dtype-set to match X.
262 void causal_conv1d_update(const Tensor& X, const Tensor& Wt, const Tensor* bias,
263 int N, int C, int L_step, int kL, int dilation,
264 Tensor& state, Tensor& Y);
265
266 } // namespace brotensor
267