GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 95.9% 94 / 0 / 98
Functions: 100.0% 5 / 0 / 5
Branches: 53.6% 45 / 0 / 84
src/cpu/window_partition.cpp
Line Branch Exec Source
1 // ─── CPU window partition / reverse ─────────────────────────────────────────
2 //
3 // FP32 scalar host implementations of the SAM-style windowed-attention
4 // layout pair. Both ops are pure rearrangements (no math) and are exact
5 // inverses of each other — neither has a separate _backward op; callers
6 // apply the other op to map gradients back through the layout change.
7 //
8 // ── Layout ─────────────────────────────────────────────────────────────────
9 // Input NCHW: X(N, C*H*W) at ((n*C + c)*H + h)*W + w
10 // Windowed: Y(N*nw_h*nw_w, C*window*window)
11 // row index = n*nw_h*nw_w + nh*nw_w + nw
12 // within-row = (c*window + lh)*window + lw
13 // With nw_h = H/window, nw_w = W/window, and (h, w) = (nh*window + lh,
14 // nw*window + lw).
15 //
16 // This is the NCHW analogue of SAM's NHWC window_partition. We keep
17 // channels-contiguous within each window so subsequent attention sees the
18 // expected (B*nw, C, window, window) layout without an extra transpose.
19 //
20 // ── ACCUMULATION ────────────────────────────────────────────────────────────
21 // Both ops OVERWRITE the output. No accumulation — they're pure copies.
22
23 #include <brotensor/tensor.h>
24
25 #include <stdexcept>
26 #include <string>
27
28 namespace brotensor::detail::cpu {
29
30 namespace {
31
32 1 [[noreturn]] inline void fail(const char* op, const std::string& reason) {
33
7/14
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✓ Branch 6 taken 1 time.
✗ Branch 7 not taken.
✓ Branch 8 taken 1 time.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 1 time.
✓ Branch 12 taken 1 time.
✗ Branch 13 not taken.
 1 throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason);
34 1 }
35
36 15 inline void check_fp32(const ::brotensor::Tensor& t,
37 const char* op, const char* name) {
38
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (t.dtype != Dtype::FP32) {
39 ✗ fail(op, std::string(name) +
40 " must be FP32 (CPU backend is FP32-only)");
41 }
42 15 }
43
44 15 inline void check_args(const char* op, int N, int C, int H, int W, int window) {
45
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (N < 0 || C < 1 || H < 1 || W < 1) {
46
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 1 fail(op, "C/H/W must be >=1 and N >=0");
47 }
48
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (window < 1) {
49 ✗ fail(op, "window must be >=1");
50 }
51
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 1 time.
 15 if (H % window != 0 || W % window != 0) {
52
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 fail(op, "H and W must be multiples of window (use pad2d first if "
53 "the input doesn't align)");
54 }
55 15 }
56
57 } // namespace
58
59 9 void window_partition_forward(const ::brotensor::Tensor& X,
60 int N, int C, int H, int W, int window,
61 ::brotensor::Tensor& Y) {
62 9 const char* op = "window_partition_forward";
63 9 check_fp32(X, op, "X");
64 9 check_args(op, N, C, H, W, window);
65
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (X.rows != N || X.cols != C * H * W) {
66 ✗ fail(op, "X shape must be (N, C*H*W)");
67 }
68 9 const int nw_h = H / window;
69 9 const int nw_w = W / window;
70 9 const int B_out = N * nw_h * nw_w;
71 9 const int cols_out = C * window * window;
72
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 9 if (Y.rows != B_out || Y.cols != cols_out || Y.dtype != Dtype::FP32) {
73 9 Y.resize(B_out, cols_out, Dtype::FP32);
74 9 }
75
2/4
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
 9 if (N == 0 || cols_out == 0) return;
76
77 9 const float* Xp = X.host_f32();
78 9 float* Yp = Y.host_f32_mut();
79
80
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 15 times.
 24 for (int n = 0; n < N; ++n) {
81
2/2
✓ Branch 0 taken 31 times.
✓ Branch 1 taken 15 times.
 46 for (int nh = 0; nh < nw_h; ++nh) {
82
2/2
✓ Branch 0 taken 74 times.
✓ Branch 1 taken 31 times.
 105 for (int nw = 0; nw < nw_w; ++nw) {
83 74 const int b_out = (n * nw_h + nh) * nw_w + nw;
84 74 float* y_row =
85 74 Yp + static_cast<long>(b_out) * cols_out;
86
2/2
✓ Branch 0 taken 224 times.
✓ Branch 1 taken 74 times.
 298 for (int c = 0; c < C; ++c) {
87 224 const float* x_chan =
88 224 Xp + (static_cast<long>(n) * C + c) * H * W;
89
2/2
✓ Branch 0 taken 860 times.
✓ Branch 1 taken 224 times.
 1084 for (int lh = 0; lh < window; ++lh) {
90 860 const int h = nh * window + lh;
91 860 const float* x_row =
92 860 x_chan + static_cast<long>(h) * W + nw * window;
93 860 float* y_block =
94 1720 y_row + (static_cast<long>(c) * window + lh)
95 860 * window;
96
2/2
✓ Branch 0 taken 3964 times.
✓ Branch 1 taken 860 times.
 4824 for (int lw = 0; lw < window; ++lw) {
97 3964 y_block[lw] = x_row[lw];
98 3964 }
99 860 }
100 224 }
101 74 }
102 31 }
103 15 }
104 9 }
105
106 5 void window_reverse_forward(const ::brotensor::Tensor& X,
107 int N, int C, int H, int W, int window,
108 ::brotensor::Tensor& Y) {
109 5 const char* op = "window_reverse_forward";
110 5 check_fp32(X, op, "X");
111 5 check_args(op, N, C, H, W, window);
112 5 const int nw_h = H / window;
113 5 const int nw_w = W / window;
114 5 const int B_in = N * nw_h * nw_w;
115 5 const int cols_in = C * window * window;
116
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (X.rows != B_in || X.cols != cols_in) {
117 ✗ fail(op, "X shape must be (N*nw_h*nw_w, C*window*window)");
118 }
119 5 const int cols_out = C * H * W;
120
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 5 if (Y.rows != N || Y.cols != cols_out || Y.dtype != Dtype::FP32) {
121 5 Y.resize(N, cols_out, Dtype::FP32);
122 5 }
123
2/4
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 5 times.
✗ Branch 3 not taken.
 5 if (N == 0 || cols_out == 0) return;
124
125 5 const float* Xp = X.host_f32();
126 5 float* Yp = Y.host_f32_mut();
127
128
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 8 times.
 13 for (int n = 0; n < N; ++n) {
129
2/2
✓ Branch 0 taken 21 times.
✓ Branch 1 taken 8 times.
 29 for (int nh = 0; nh < nw_h; ++nh) {
130
2/2
✓ Branch 0 taken 54 times.
✓ Branch 1 taken 21 times.
 75 for (int nw = 0; nw < nw_w; ++nw) {
131 54 const int b_in = (n * nw_h + nh) * nw_w + nw;
132 54 const float* x_row =
133 54 Xp + static_cast<long>(b_in) * cols_in;
134
2/2
✓ Branch 0 taken 170 times.
✓ Branch 1 taken 54 times.
 224 for (int c = 0; c < C; ++c) {
135 170 float* y_chan =
136 170 Yp + (static_cast<long>(n) * C + c) * H * W;
137
2/2
✓ Branch 0 taken 636 times.
✓ Branch 1 taken 170 times.
 806 for (int lh = 0; lh < window; ++lh) {
138 636 const int h = nh * window + lh;
139 636 float* y_row =
140 636 y_chan + static_cast<long>(h) * W + nw * window;
141 636 const float* x_block =
142 1272 x_row + (static_cast<long>(c) * window + lh)
143 636 * window;
144
2/2
✓ Branch 0 taken 2988 times.
✓ Branch 1 taken 636 times.
 3624 for (int lw = 0; lw < window; ++lw) {
145 2988 y_row[lw] = x_block[lw];
146 2988 }
147 636 }
148 170 }
149 54 }
150 21 }
151 8 }
152 5 }
153
154 } // namespace brotensor::detail::cpu
155