GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 95.0% 114 / 0 / 120
Functions: 100.0% 3 / 0 / 3
Branches: 61.5% 64 / 0 / 104
src/cpu/spatial_merge.cpp
Line Branch Exec Source
1 // ─── CPU spatial 2x2 pixel-unshuffle (Qwen-VL merger / Flux.2 VAE) ──────────
2 //
3 // Pure gather. Stacks each 2x2 spatial block into the channel axis:
4 // X: NCHW (N, C*H*W) — H, W must be even.
5 // Y: NCHW (N, 4*C*(H/2)*(W/2))
6 // block = dh*2 + dw, (h_in, w_in) = (2*h_out + dh, 2*w_out + dw).
7 // channel_major=false: c_out = block*C + c_in (Qwen-VL, block-major)
8 // channel_major=true: c_out = c_in*4 + block (torch pixel_unshuffle / Flux.2)
9 // FP32 only on CPU (CPU backend is FP32-only by convention).
10
11 #include <brotensor/tensor.h>
12
13 #include <cstddef>
14 #include <stdexcept>
15 #include <string>
16
17 namespace brotensor::detail::cpu {
18
19 23 void spatial_merge_2x2_forward(const ::brotensor::Tensor& X,
20 int N, int C, int H, int W,
21 bool channel_major,
22 ::brotensor::Tensor& Y) {
23
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (X.dtype != Dtype::FP32) {
24 ✗ throw std::runtime_error("spatial_merge_2x2_forward: X must be FP32 "
25 "(CPU backend is FP32-only)");
26 }
27
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (N < 0 || C < 0 || H < 0 || W < 0) {
28 ✗ throw std::runtime_error("spatial_merge_2x2_forward: negative dimension");
29 }
30
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 1 time.
 23 if ((H & 1) != 0 || (W & 1) != 0) {
31
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 throw std::runtime_error("spatial_merge_2x2_forward: H and W must be even");
32 }
33 22 const int H_out = H / 2;
34 22 const int W_out = W / 2;
35 22 const int C_out = 4 * C;
36 22 const int cols = C_out * H_out * W_out;
37
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 22 if (Y.rows != N || Y.cols != cols || Y.dtype != Dtype::FP32) {
38 22 Y.resize(N, cols, Dtype::FP32);
39 22 }
40
4/8
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 22 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 22 times.
✗ Branch 7 not taken.
 22 if (N == 0 || C == 0 || H_out == 0 || W_out == 0) return;
41
42 22 const float* Xp = X.host_f32();
43 22 float* Yp = Y.host_f32_mut();
44
45 22 const int HW = H * W;
46 22 const int HW_out = H_out * W_out;
47
48
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 34 times.
 56 for (int n = 0; n < N; ++n) {
49
2/2
✓ Branch 0 taken 68 times.
✓ Branch 1 taken 34 times.
 102 for (int dh = 0; dh < 2; ++dh) {
50
2/2
✓ Branch 0 taken 136 times.
✓ Branch 1 taken 68 times.
 204 for (int dw = 0; dw < 2; ++dw) {
51 136 const int block = dh * 2 + dw;
52
2/2
✓ Branch 0 taken 1060 times.
✓ Branch 1 taken 136 times.
 1196 for (int c_in = 0; c_in < C; ++c_in) {
53 1060 const int c_out =
54
2/2
✓ Branch 0 taken 592 times.
✓ Branch 1 taken 468 times.
 1060 channel_major ? c_in * 4 + block : block * C + c_in;
55 1060 const float* xc = Xp + (n * C + c_in) * HW;
56 1060 float* yc = Yp + (n * C_out + c_out) * HW_out;
57
2/2
✓ Branch 0 taken 7448 times.
✓ Branch 1 taken 1060 times.
 8508 for (int h_out = 0; h_out < H_out; ++h_out) {
58 7448 const int h_in = 2 * h_out + dh;
59
2/2
✓ Branch 0 taken 70112 times.
✓ Branch 1 taken 7448 times.
 77560 for (int w_out = 0; w_out < W_out; ++w_out) {
60 70112 const int w_in = 2 * w_out + dw;
61 70112 yc[h_out * W_out + w_out] = xc[h_in * W + w_in];
62 70112 }
63 7448 }
64 1060 }
65 136 }
66 68 }
67 34 }
68 22 }
69
70 // ─── CPU DC-AE up-shortcut: repeat_interleave + 2x pixel-shuffle ────────────
71 //
72 // Y[n, c_out, 2h+i, 2w+j] = X[n, (4*c_out + 2i + j)/repeats, h, w],
73 // repeats = 4*C_out/C_in. See ops/spatial.h. FP32-only on CPU.
74
75 29 void pixel_shuffle_upsample_2x_forward(const ::brotensor::Tensor& X,
76 int N, int C_in, int H, int W,
77 int C_out,
78 ::brotensor::Tensor& Y) {
79
1/2
✓ Branch 0 taken 29 times.
✗ Branch 1 not taken.
 29 if (X.dtype != Dtype::FP32) {
80 ✗ throw std::runtime_error("pixel_shuffle_upsample_2x_forward: X must be "
81 "FP32 (CPU backend is FP32-only)");
82 }
83
2/2
✓ Branch 0 taken 25 times.
✓ Branch 1 taken 4 times.
 29 if (N < 0 || C_in <= 0 || H < 0 || W < 0 || C_out <= 0) {
84
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 throw std::runtime_error("pixel_shuffle_upsample_2x_forward: bad dimension");
85 }
86
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 1 time.
 25 if ((4 * C_out) % C_in != 0) {
87
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 throw std::runtime_error("pixel_shuffle_upsample_2x_forward: C_in must "
88 "divide 4*C_out");
89 }
90 24 const int repeats = (4 * C_out) / C_in;
91 24 const int H_out = 2 * H;
92 24 const int W_out = 2 * W;
93 24 const int cols = C_out * H_out * W_out;
94
3/6
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 23 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 24 if (Y.rows != N || Y.cols != cols || Y.dtype != Dtype::FP32) {
95 24 Y.resize(N, cols, Dtype::FP32);
96 24 }
97
5/6
✓ Branch 0 taken 23 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 22 times.
✓ Branch 3 taken 1 time.
✗ Branch 4 not taken.
✓ Branch 5 taken 22 times.
 24 if (N == 0 || H == 0 || W == 0) return;
98
99 22 const float* Xp = X.host_f32();
100 22 float* Yp = Y.host_f32_mut();
101 22 const int HW = H * W;
102 22 const int HW_out = H_out * W_out;
103
104
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 37 times.
 59 for (int n = 0; n < N; ++n) {
105
2/2
✓ Branch 0 taken 395 times.
✓ Branch 1 taken 37 times.
 432 for (int c_out = 0; c_out < C_out; ++c_out) {
106 395 float* yc = Yp + (static_cast<std::size_t>(n) * C_out + c_out) * HW_out;
107
2/2
✓ Branch 0 taken 9998 times.
✓ Branch 1 taken 395 times.
 10393 for (int h_out = 0; h_out < H_out; ++h_out) {
108 9998 const int i = h_out & 1;
109 9998 const int h = h_out >> 1;
110
2/2
✓ Branch 0 taken 383100 times.
✓ Branch 1 taken 9998 times.
 393098 for (int w_out = 0; w_out < W_out; ++w_out) {
111 383100 const int j = w_out & 1;
112 383100 const int w = w_out >> 1;
113 383100 const int src_c = (4 * c_out + 2 * i + j) / repeats;
114 383100 const float* xc =
115 383100 Xp + (static_cast<std::size_t>(n) * C_in + src_c) * HW;
116 383100 yc[h_out * W_out + w_out] = xc[h * W + w];
117 383100 }
118 9998 }
119 395 }
120 37 }
121 24 }
122
123 // ─── CPU DiT unpatchify: token rows -> image ────────────────────────────────
124 //
125 // Y[c, i*P+py, j*P+px] = tokens[i*wp+j, col], c in [0,C_keep). See ops/spatial.h.
126 // FP32-only on CPU.
127
128 6 void patch_unpack_forward(const ::brotensor::Tensor& tokens,
129 int hp, int wp, int P, int C_total, int C_keep,
130 bool channel_major,
131 ::brotensor::Tensor& Y) {
132
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (tokens.dtype != Dtype::FP32) {
133 ✗ throw std::runtime_error("patch_unpack_forward: tokens must be FP32 "
134 "(CPU backend is FP32-only)");
135 }
136
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (hp < 0 || wp < 0 || P <= 0 || C_total <= 0 || C_keep <= 0 ||
137 6 C_keep > C_total) {
138 ✗ throw std::runtime_error("patch_unpack_forward: bad dimension");
139 }
140 6 const int PP = P * P;
141 6 const int N = hp * wp;
142
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (tokens.rows != N || tokens.cols != PP * C_total) {
143 ✗ throw std::runtime_error("patch_unpack_forward: tokens shape mismatch");
144 }
145 6 const int H = hp * P;
146 6 const int W = wp * P;
147 6 const int cols = C_keep * H * W;
148
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (Y.rows != 1 || Y.cols != cols || Y.dtype != Dtype::FP32) {
149 6 Y.resize(1, cols, Dtype::FP32);
150 6 }
151
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (N == 0) return;
152
153 6 const float* Xp = tokens.host_f32();
154 6 float* Yp = Y.host_f32_mut();
155 6 const int row_stride = PP * C_total;
156 6 const std::size_t plane = static_cast<std::size_t>(H) * W;
157
158
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 20 times.
 26 for (int c = 0; c < C_keep; ++c) {
159
2/2
✓ Branch 0 taken 320 times.
✓ Branch 1 taken 20 times.
 340 for (int i = 0; i < hp; ++i) {
160
2/2
✓ Branch 0 taken 640 times.
✓ Branch 1 taken 320 times.
 960 for (int py = 0; py < P; ++py) {
161 640 const int y = i * P + py;
162
2/2
✓ Branch 0 taken 33552 times.
✓ Branch 1 taken 640 times.
 34192 for (int j = 0; j < wp; ++j) {
163 33552 const int tok = i * wp + j;
164 67104 const float* trow = Xp +
165 33552 static_cast<std::size_t>(tok) * row_stride;
166
2/2
✓ Branch 0 taken 67104 times.
✓ Branch 1 taken 33552 times.
 100656 for (int px = 0; px < P; ++px) {
167 67104 const int x = j * P + px;
168 67104 const int block = py * P + px;
169
2/2
✓ Branch 0 taken 272 times.
✓ Branch 1 taken 66832 times.
 67104 const int col = channel_major
170 272 ? c * PP + block
171 66832 : block * C_total + c;
172 201312 Yp[static_cast<std::size_t>(c) * plane +
173 268416 static_cast<std::size_t>(y) * W + x] = trow[col];
174 67104 }
175 33552 }
176 640 }
177 320 }
178 20 }
179 6 }
180
181 } // namespace brotensor::detail::cpu
182