GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 91.1% 102 / 0 / 112
Functions: 100.0% 5 / 0 / 5
Branches: 34.8% 32 / 0 / 92
src/cpu/diffusion_samplers.cpp
Line Branch Exec Source
1 // ─── CPU diffusion sampler steps + timestep embedding (CHUNK 4) ────────────
2 //
3 // FP32 scalar host implementations. Ports the elementwise sampler kernels:
4 // src/cuda/ddim_step.cu, euler_step.cu, dpmpp_2m_step.cu,
5 // src/cuda/timestep_embedding.cu
6 //
7 // IMPORTANT — dtype: the GPU sampler kernels (ddim/euler/dpmpp_2m) run FP16
8 // internally (their tensors must be FP16). The CPU backend is FP32-only, so
9 // the CPU impls require FP32 tensors. CPU↔GPU parity for these three ops
10 // therefore feeds FP16 to the GPU and FP32 to the CPU and compares with a
11 // loose FP16-driven tolerance (see tests/test_diffusion_parity.cpp). The
12 // internal arithmetic is identical FP32 math in both backends — the GPU just
13 // rounds inputs/outputs through FP16 storage. timestep_embedding is FP32 on
14 // both backends, so it gets a tight tolerance.
15 //
16 // ACCUMULATION: every op fully OVERWRITES its outputs (x_prev / x0_out / Y).
17 //
18 // ── ddim_step ──
19 // x0_pred = (x_t - sqrt(1-alpha_t) * eps_pred) / sqrt(alpha_t)
20 // dir = sqrt(max(0, 1 - alpha_prev - sigma_t^2)) * eps_pred
21 // x_prev = sqrt(max(0, alpha_prev)) * x0_pred + dir
22 // (inv_sqrt_alpha_t is 0 when sqrt(alpha_t) <= 0, matching the GPU.)
23 //
24 // ── euler_step ──
25 // x_prev = x_t + (sigma_prev - sigma_t) * eps_pred
26 //
27 // ── dpmpp_2m_step ──
28 // x0_t = x_t - sigma_t * eps_pred
29 // x_prev = c_xt * x_t + c_x0t * x0_t + c_x0prev * x0_prev
30 // x0_out = x0_t
31 //
32 // ── timestep_embedding ── (diffusers get_timestep_embedding,
33 // flip_sin_to_cos=True, downscale_freq_shift=0):
34 // half = dim / 2
35 // freqs[k] = exp(-log(max_period) * k / half)
36 // args[i,j] = timesteps[i] * freqs[k], k = j (j<half) else j-half
37 // Y[i, 0:half] = cos(args)
38 // Y[i, half:2*half] = sin(args)
39 // if dim is odd: Y[i, dim-1] = 0
40
41 #include <brotensor/tensor.h>
42
43 #include <algorithm>
44 #include <cmath>
45 #include <stdexcept>
46 #include <string>
47
48 namespace brotensor::detail::cpu {
49
50 namespace {
51
52 48 inline void check_fp32(const ::brotensor::Tensor& t,
53 const char* op, const char* name) {
54
1/2
✓ Branch 0 taken 48 times.
✗ Branch 1 not taken.
 48 if (t.dtype != Dtype::FP32) {
55 ✗ throw std::runtime_error(std::string(op) + ": " + name +
56 " must be FP32 (CPU backend is FP32-only)");
57 }
58 48 }
59
60 } // namespace
61
62 7 void ddim_step(const ::brotensor::Tensor& x_t,
63 const ::brotensor::Tensor& eps_pred,
64 float alpha_t, float alpha_prev, float sigma_t,
65 ::brotensor::Tensor& x_prev) {
66 7 check_fp32(x_t, "ddim_step", "x_t");
67 7 check_fp32(eps_pred, "ddim_step", "eps_pred");
68
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (x_t.rows != eps_pred.rows || x_t.cols != eps_pred.cols) {
69 ✗ throw std::runtime_error("ddim_step: shape mismatch between x_t and eps_pred");
70 }
71
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 7 if (x_prev.rows != x_t.rows || x_prev.cols != x_t.cols ||
72 ✗ x_prev.dtype != Dtype::FP32) {
73 7 x_prev.resize(x_t.rows, x_t.cols, Dtype::FP32);
74 7 }
75 7 const int total = x_t.size();
76
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (total == 0) return;
77
78 // Scalar coefficients precomputed in FP32 — identical to the GPU host code.
79 7 const float sqrt_alpha_t = std::sqrt(alpha_t);
80
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 const float inv_sqrt_alpha_t = sqrt_alpha_t > 0.0f ? 1.0f / sqrt_alpha_t : 0.0f;
81 7 const float sqrt_1m_alpha_t = std::sqrt(std::max(0.0f, 1.0f - alpha_t));
82 7 const float sqrt_alpha_prev = std::sqrt(std::max(0.0f, alpha_prev));
83 7 const float dir_inner = 1.0f - alpha_prev - sigma_t * sigma_t;
84 7 const float dir_coef = std::sqrt(std::max(0.0f, dir_inner));
85
86 7 const float* xtp = x_t.host_f32();
87 7 const float* epsp = eps_pred.host_f32();
88 7 float* xpp = x_prev.host_f32_mut();
89
90
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 480 times.
 487 for (int i = 0; i < total; ++i) {
91 480 const float xt = xtp[i];
92 480 const float eps = epsp[i];
93 480 const float x0_pred = (xt - sqrt_1m_alpha_t * eps) * inv_sqrt_alpha_t;
94 480 const float dir = dir_coef * eps;
95 480 xpp[i] = sqrt_alpha_prev * x0_pred + dir;
96 480 }
97 7 }
98
99 6 void euler_step(const ::brotensor::Tensor& x_t,
100 const ::brotensor::Tensor& eps_pred,
101 float sigma_t, float sigma_prev,
102 ::brotensor::Tensor& x_prev) {
103 6 check_fp32(x_t, "euler_step", "x_t");
104 6 check_fp32(eps_pred, "euler_step", "eps_pred");
105
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (x_t.rows != eps_pred.rows || x_t.cols != eps_pred.cols) {
106 ✗ throw std::runtime_error("euler_step: shape mismatch between x_t and eps_pred");
107 }
108
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 (x_prev.rows != x_t.rows || x_prev.cols != x_t.cols ||
109 ✗ x_prev.dtype != Dtype::FP32) {
110 6 x_prev.resize(x_t.rows, x_t.cols, Dtype::FP32);
111 6 }
112 6 const int total = x_t.size();
113
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (total == 0) return;
114
115 6 const float dsigma = sigma_prev - sigma_t;
116
117 6 const float* xtp = x_t.host_f32();
118 6 const float* epsp = eps_pred.host_f32();
119 6 float* xpp = x_prev.host_f32_mut();
120
121
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 416 times.
 422 for (int i = 0; i < total; ++i) {
122 416 xpp[i] = xtp[i] + dsigma * epsp[i];
123 416 }
124 6 }
125
126 6 void dpmpp_2m_step(const ::brotensor::Tensor& x_t,
127 const ::brotensor::Tensor& eps_pred,
128 const ::brotensor::Tensor& x0_prev,
129 float sigma_t,
130 float c_xt, float c_x0t, float c_x0prev,
131 ::brotensor::Tensor& x_prev,
132 ::brotensor::Tensor& x0_out) {
133 6 check_fp32(x_t, "dpmpp_2m_step", "x_t");
134 6 check_fp32(eps_pred, "dpmpp_2m_step", "eps_pred");
135 6 check_fp32(x0_prev, "dpmpp_2m_step", "x0_prev");
136
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 12 if (x_t.rows != eps_pred.rows || x_t.cols != eps_pred.cols ||
137 6 x_t.rows != x0_prev.rows || x_t.cols != x0_prev.cols) {
138 ✗ throw std::runtime_error("dpmpp_2m_step: shape mismatch");
139 }
140
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 (x_prev.rows != x_t.rows || x_prev.cols != x_t.cols ||
141 ✗ x_prev.dtype != Dtype::FP32) {
142 6 x_prev.resize(x_t.rows, x_t.cols, Dtype::FP32);
143 6 }
144
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 (x0_out.rows != x_t.rows || x0_out.cols != x_t.cols ||
145 ✗ x0_out.dtype != Dtype::FP32) {
146 6 x0_out.resize(x_t.rows, x_t.cols, Dtype::FP32);
147 6 }
148 6 const int total = x_t.size();
149
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (total == 0) return;
150
151 6 const float* xtp = x_t.host_f32();
152 6 const float* epsp = eps_pred.host_f32();
153 6 const float* x0pp = x0_prev.host_f32();
154 6 float* xpp = x_prev.host_f32_mut();
155 6 float* x0op = x0_out.host_f32_mut();
156
157
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 512 times.
 518 for (int i = 0; i < total; ++i) {
158 512 const float xt = xtp[i];
159 512 const float eps = epsp[i];
160 512 const float x0p = x0pp[i];
161 512 const float x0t = xt - sigma_t * eps;
162 512 xpp[i] = c_xt * xt + c_x0t * x0t + c_x0prev * x0p;
163 512 x0op[i] = x0t;
164 512 }
165 6 }
166
167 4 void timestep_embedding(const ::brotensor::Tensor& timesteps,
168 int dim, float max_period,
169 ::brotensor::Tensor& Y) {
170 4 check_fp32(timesteps, "timestep_embedding", "timesteps");
171
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (timesteps.cols != 1) {
172 ✗ throw std::runtime_error("timestep_embedding: timesteps must be (N,1)");
173 }
174
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (dim <= 0) {
175 ✗ throw std::runtime_error("timestep_embedding: dim must be positive");
176 }
177 4 const int N = timesteps.rows;
178
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 4 if (Y.rows != N || Y.cols != dim || Y.dtype != Dtype::FP32) {
179 4 Y.resize(N, dim, Dtype::FP32);
180 4 }
181
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (N == 0) return;
182
183 4 const int half = dim / 2;
184 4 const float log_max_period = std::log(max_period);
185
186 4 const float* tsp = timesteps.host_f32();
187 4 float* Yp = Y.host_f32_mut();
188
189
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 4 times.
 17 for (int i = 0; i < N; ++i) {
190 13 const float ts = tsp[i];
191
2/2
✓ Branch 0 taken 1972 times.
✓ Branch 1 taken 13 times.
 1985 for (int j = 0; j < dim; ++j) {
192
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 1968 times.
 1972 if (j >= 2 * half) {
193 // Odd-dim tail slot.
194 4 Yp[i * dim + j] = 0.0f;
195 4 continue;
196 }
197
2/2
✓ Branch 0 taken 984 times.
✓ Branch 1 taken 984 times.
 1968 const int k = j < half ? j : j - half;
198 5904 const float freq = std::exp(-log_max_period *
199 3936 static_cast<float>(k) /
200 1968 static_cast<float>(half));
201 1968 const float arg = ts * freq;
202
2/2
✓ Branch 0 taken 984 times.
✓ Branch 1 taken 984 times.
 1968 Yp[i * dim + j] = j < half ? std::cos(arg) : std::sin(arg);
203 1968 }
204 13 }
205 4 }
206
207 } // namespace brotensor::detail::cpu
208