GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 92.9% 156 / 0 / 168
Functions: 81.2% 26 / 0 / 32
Branches: 44.0% 103 / 0 / 234
src/cpu/diffusion_samplers.cpp
Line Branch Exec Source
1 // ─── CPU diffusion sampler steps + timestep embedding (CHUNK 4) ────────────
2 //
3 // Scalar host implementations supporting FP32, FP16, and BF16 tensors.
4 // Ports the elementwise sampler kernels:
5 // src/cuda/ddim_step.cu, euler_step.cu, dpmpp_2m_step.cu,
6 // src/cuda/timestep_embedding.cu
7 //
8 // Arithmetic is evaluated in FP32; inputs/outputs are converted to/from FP32
9 // when using FP16 or BF16 storage, matching GPU sampler behavior.
10 //
11 // ACCUMULATION: every op fully OVERWRITES its outputs (x_prev / x0_out / Y).
12 //
13 // ── ddim_step ──
14 // x0_pred = (x_t - sqrt(1-alpha_t) * eps_pred) / sqrt(alpha_t)
15 // dir = sqrt(max(0, 1 - alpha_prev - sigma_t^2)) * eps_pred
16 // x_prev = sqrt(max(0, alpha_prev)) * x0_pred + dir
17 // (inv_sqrt_alpha_t is 0 when sqrt(alpha_t) <= 0, matching the GPU.)
18 //
19 // ── euler_step ──
20 // x_prev = x_t + (sigma_prev - sigma_t) * eps_pred
21 //
22 // ── dpmpp_2m_step ──
23 // x0_t = x_t - sigma_t * eps_pred
24 // x_prev = c_xt * x_t + c_x0t * x0_t + c_x0prev * x0_prev
25 // x0_out = x0_t
26 //
27 // ── timestep_embedding ── (diffusers get_timestep_embedding,
28 // flip_sin_to_cos=True, downscale_freq_shift=0):
29 // half = dim / 2
30 // freqs[k] = exp(-log(max_period) * k / half)
31 // args[i,j] = timesteps[i] * freqs[k], k = j (j<half) else j-half
32 // Y[i, 0:half] = cos(args)
33 // Y[i, half:2*half] = sin(args)
34 // if dim is odd: Y[i, dim-1] = 0
35
36 #include <brotensor/tensor.h>
37
38 #include <algorithm>
39 #include <cmath>
40 #include <cstdint>
41 #include <stdexcept>
42 #include <string>
43
44 namespace brotensor::detail::cpu {
45
46 namespace {
47
48 35 inline void check_sampler_dtype(Dtype dt, const char* op, const char* name) {
49
5/6
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 27 times.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 35 if (dt != Dtype::FP32 && dt != Dtype::FP16 && dt != Dtype::BF16) {
50 ✗ throw std::runtime_error(std::string(op) + ": " + name +
51 " must be FP32, FP16, or BF16");
52 }
53 35 }
54
55 struct Fp32Accessor {
56 3457 static float load(const void* ptr, int i) {
57 3457 return static_cast<const float*>(ptr)[i];
58 }
59 4004 static void store(void* ptr, int i, float v) {
60 4004 static_cast<float*>(ptr)[i] = v;
61 4004 }
62 };
63
64 struct Fp16Accessor {
65 116 static float load(const void* ptr, int i) {
66 116 return fp16_bits_to_fp32(static_cast<const uint16_t*>(ptr)[i]);
67 }
68 64 static void store(void* ptr, int i, float v) {
69 64 static_cast<uint16_t*>(ptr)[i] = fp32_to_fp16_bits(v);
70 64 }
71 };
72
73 struct Bf16Accessor {
74 116 static float load(const void* ptr, int i) {
75 116 return bf16_bits_to_fp32(static_cast<const uint16_t*>(ptr)[i]);
76 }
77 64 static void store(void* ptr, int i, float v) {
78 64 static_cast<uint16_t*>(ptr)[i] = fp32_to_bf16_bits(v);
79 64 }
80 };
81
82 template <typename Acc>
83 10 void ddim_step_impl(const void* xtp, const void* epsp, void* xpp, int total,
84 float sqrt_alpha_prev, float inv_sqrt_alpha_t,
85 float sqrt_1m_alpha_t, float dir_coef) {
86
6/6
✓ Branch 0 taken 496 times.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 16 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 16 times.
✓ Branch 5 taken 1 time.
 538 for (int i = 0; i < total; ++i) {
87 528 const float xt = Acc::load(xtp, i);
88 528 const float eps = Acc::load(epsp, i);
89 528 const float x0_pred = (xt - sqrt_1m_alpha_t * eps) * inv_sqrt_alpha_t;
90 528 const float dir = dir_coef * eps;
91 528 Acc::store(xpp, i, sqrt_alpha_prev * x0_pred + dir);
92 528 }
93 10 }
94
95 template <typename Acc>
96 9 void euler_step_impl(const void* xtp, const void* epsp, void* xpp, int total, float dsigma) {
97
6/6
✓ Branch 0 taken 432 times.
✓ Branch 1 taken 7 times.
✓ Branch 2 taken 16 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 16 times.
✓ Branch 5 taken 1 time.
 473 for (int i = 0; i < total; ++i) {
98 464 const float xt = Acc::load(xtp, i);
99 464 const float eps = Acc::load(epsp, i);
100 464 Acc::store(xpp, i, xt + dsigma * eps);
101 464 }
102 9 }
103
104 template <typename Acc>
105 9 void dpmpp_2m_step_impl(const void* xtp, const void* epsp, const void* x0pp,
106 void* xpp, void* x0op, int total,
107 float sigma_t, float c_xt, float c_x0t, float c_x0prev) {
108
6/6
✓ Branch 0 taken 528 times.
✓ Branch 1 taken 7 times.
✓ Branch 2 taken 16 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 16 times.
✓ Branch 5 taken 1 time.
 569 for (int i = 0; i < total; ++i) {
109 560 const float xt = Acc::load(xtp, i);
110 560 const float eps = Acc::load(epsp, i);
111 560 const float x0p = Acc::load(x0pp, i);
112 560 const float x0t = xt - sigma_t * eps;
113 560 Acc::store(xpp, i, c_xt * xt + c_x0t * x0t + c_x0prev * x0p);
114 560 Acc::store(x0op, i, x0t);
115 560 }
116 9 }
117
118 template <typename InAcc, typename OutAcc>
119 7 void timestep_embedding_impl(const void* tsp, void* Yp, int N, int dim, int half, float log_max_period) {
120
6/18
✓ Branch 0 taken 17 times.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✓ Branch 7 taken 1 time.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✓ Branch 12 taken 4 times.
✓ Branch 13 taken 1 time.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
 32 for (int i = 0; i < N; ++i) {
121 25 const float ts = InAcc::load(tsp, i);
122
6/18
✓ Branch 0 taken 1988 times.
✓ Branch 1 taken 17 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✓ Branch 7 taken 4 times.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✓ Branch 13 taken 4 times.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
 2045 for (int j = 0; j < dim; ++j) {
123
4/18
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 1984 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 16 times.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✓ Branch 13 taken 16 times.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
 2020 if (j >= 2 * half) {
124 4 OutAcc::store(Yp, i * dim + j, 0.0f);
125 4 continue;
126 }
127
6/18
✓ Branch 0 taken 992 times.
✓ Branch 1 taken 992 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
✓ Branch 7 taken 8 times.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✓ Branch 12 taken 8 times.
✓ Branch 13 taken 8 times.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
 2016 const int k = j < half ? j : j - half;
128 6048 const float freq = std::exp(-log_max_period *
129 4032 static_cast<float>(k) /
130 2016 static_cast<float>(half));
131 2016 const float arg = ts * freq;
132
6/18
✓ Branch 0 taken 992 times.
✓ Branch 1 taken 992 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
✓ Branch 7 taken 8 times.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✓ Branch 12 taken 8 times.
✓ Branch 13 taken 8 times.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 17 not taken.
 2016 OutAcc::store(Yp, i * dim + j, j < half ? std::cos(arg) : std::sin(arg));
133 2016 }
134 25 }
135 7 }
136
137 } // namespace
138
139 10 void ddim_step(const ::brotensor::Tensor& x_t,
140 const ::brotensor::Tensor& eps_pred,
141 float alpha_t, float alpha_prev, float sigma_t,
142 ::brotensor::Tensor& x_prev) {
143 10 check_sampler_dtype(x_t.dtype, "ddim_step", "x_t");
144
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (eps_pred.dtype != x_t.dtype) {
145 ✗ throw std::runtime_error("ddim_step: x_t and eps_pred must have the same dtype");
146 }
147
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (x_t.rows != eps_pred.rows || x_t.cols != eps_pred.cols) {
148 ✗ throw std::runtime_error("ddim_step: shape mismatch between x_t and eps_pred");
149 }
150 10 const Dtype dtype = x_t.dtype;
151
4/6
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 7 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 10 if (x_prev.rows != x_t.rows || x_prev.cols != x_t.cols || x_prev.dtype != dtype) {
152 7 x_prev.resize(x_t.rows, x_t.cols, dtype);
153 7 }
154 10 const int total = x_t.size();
155
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (total == 0) return;
156
157 // Scalar coefficients precomputed in FP32 — identical to the GPU host code.
158 10 const float sqrt_alpha_t = std::sqrt(alpha_t);
159
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 const float inv_sqrt_alpha_t = sqrt_alpha_t > 0.0f ? 1.0f / sqrt_alpha_t : 0.0f;
160 10 const float sqrt_1m_alpha_t = std::sqrt(std::max(0.0f, 1.0f - alpha_t));
161 10 const float sqrt_alpha_prev = std::sqrt(std::max(0.0f, alpha_prev));
162 10 const float dir_inner = 1.0f - alpha_prev - sigma_t * sigma_t;
163 10 const float dir_coef = std::sqrt(std::max(0.0f, dir_inner));
164
165
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 8 times.
 10 if (dtype == Dtype::FP32) {
166 16 ddim_step_impl<Fp32Accessor>(x_t.data, eps_pred.data, x_prev.data, total,
167 8 sqrt_alpha_prev, inv_sqrt_alpha_t, sqrt_1m_alpha_t, dir_coef);
168
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 10 } else if (dtype == Dtype::FP16) {
169 2 ddim_step_impl<Fp16Accessor>(x_t.data, eps_pred.data, x_prev.data, total,
170 1 sqrt_alpha_prev, inv_sqrt_alpha_t, sqrt_1m_alpha_t, dir_coef);
171 1 } else {
172 2 ddim_step_impl<Bf16Accessor>(x_t.data, eps_pred.data, x_prev.data, total,
173 1 sqrt_alpha_prev, inv_sqrt_alpha_t, sqrt_1m_alpha_t, dir_coef);
174 }
175 10 }
176
177 9 void euler_step(const ::brotensor::Tensor& x_t,
178 const ::brotensor::Tensor& eps_pred,
179 float sigma_t, float sigma_prev,
180 ::brotensor::Tensor& x_prev) {
181 9 check_sampler_dtype(x_t.dtype, "euler_step", "x_t");
182
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (eps_pred.dtype != x_t.dtype) {
183 ✗ throw std::runtime_error("euler_step: x_t and eps_pred must have the same dtype");
184 }
185
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (x_t.rows != eps_pred.rows || x_t.cols != eps_pred.cols) {
186 ✗ throw std::runtime_error("euler_step: shape mismatch between x_t and eps_pred");
187 }
188 9 const Dtype dtype = x_t.dtype;
189
4/6
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 9 if (x_prev.rows != x_t.rows || x_prev.cols != x_t.cols || x_prev.dtype != dtype) {
190 6 x_prev.resize(x_t.rows, x_t.cols, dtype);
191 6 }
192 9 const int total = x_t.size();
193
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (total == 0) return;
194
195 9 const float dsigma = sigma_prev - sigma_t;
196
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 7 times.
 9 if (dtype == Dtype::FP32) {
197 7 euler_step_impl<Fp32Accessor>(x_t.data, eps_pred.data, x_prev.data, total, dsigma);
198
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 9 } else if (dtype == Dtype::FP16) {
199 1 euler_step_impl<Fp16Accessor>(x_t.data, eps_pred.data, x_prev.data, total, dsigma);
200 1 } else {
201 1 euler_step_impl<Bf16Accessor>(x_t.data, eps_pred.data, x_prev.data, total, dsigma);
202 }
203 9 }
204
205 9 void dpmpp_2m_step(const ::brotensor::Tensor& x_t,
206 const ::brotensor::Tensor& eps_pred,
207 const ::brotensor::Tensor& x0_prev,
208 float sigma_t,
209 float c_xt, float c_x0t, float c_x0prev,
210 ::brotensor::Tensor& x_prev,
211 ::brotensor::Tensor& x0_out) {
212 9 check_sampler_dtype(x_t.dtype, "dpmpp_2m_step", "x_t");
213
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (eps_pred.dtype != x_t.dtype || x0_prev.dtype != x_t.dtype) {
214 ✗ throw std::runtime_error("dpmpp_2m_step: all inputs must have the same dtype");
215 }
216
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 18 if (x_t.rows != eps_pred.rows || x_t.cols != eps_pred.cols ||
217 9 x_t.rows != x0_prev.rows || x_t.cols != x0_prev.cols) {
218 ✗ throw std::runtime_error("dpmpp_2m_step: shape mismatch");
219 }
220 9 const Dtype dtype = x_t.dtype;
221
4/6
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 9 if (x_prev.rows != x_t.rows || x_prev.cols != x_t.cols || x_prev.dtype != dtype) {
222 6 x_prev.resize(x_t.rows, x_t.cols, dtype);
223 6 }
224
4/6
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 9 if (x0_out.rows != x_t.rows || x0_out.cols != x_t.cols || x0_out.dtype != dtype) {
225 6 x0_out.resize(x_t.rows, x_t.cols, dtype);
226 6 }
227 9 const int total = x_t.size();
228
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (total == 0) return;
229
230
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 7 times.
 9 if (dtype == Dtype::FP32) {
231 14 dpmpp_2m_step_impl<Fp32Accessor>(x_t.data, eps_pred.data, x0_prev.data,
232 7 x_prev.data, x0_out.data, total,
233 7 sigma_t, c_xt, c_x0t, c_x0prev);
234
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 9 } else if (dtype == Dtype::FP16) {
235 2 dpmpp_2m_step_impl<Fp16Accessor>(x_t.data, eps_pred.data, x0_prev.data,
236 1 x_prev.data, x0_out.data, total,
237 1 sigma_t, c_xt, c_x0t, c_x0prev);
238 1 } else {
239 2 dpmpp_2m_step_impl<Bf16Accessor>(x_t.data, eps_pred.data, x0_prev.data,
240 1 x_prev.data, x0_out.data, total,
241 1 sigma_t, c_xt, c_x0t, c_x0prev);
242 }
243 9 }
244
245 7 void timestep_embedding(const ::brotensor::Tensor& timesteps,
246 int dim, float max_period,
247 ::brotensor::Tensor& Y) {
248 7 check_sampler_dtype(timesteps.dtype, "timestep_embedding", "timesteps");
249
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (timesteps.cols != 1) {
250 ✗ throw std::runtime_error("timestep_embedding: timesteps must be (N,1)");
251 }
252
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (dim <= 0) {
253 ✗ throw std::runtime_error("timestep_embedding: dim must be positive");
254 }
255 7 const int N = timesteps.rows;
256
3/6
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 7 times.
✗ Branch 5 not taken.
 7 Dtype out_dt = (Y.dtype == Dtype::FP16 || Y.dtype == Dtype::BF16 || Y.dtype == Dtype::FP32)
257 7 ? Y.dtype : timesteps.dtype;
258
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 (Y.rows != N || Y.cols != dim || Y.dtype != out_dt) {
259 7 Y.resize(N, dim, out_dt);
260 7 }
261
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (N == 0) return;
262
263 7 const int half = dim / 2;
264 7 const float log_max_period = std::log(max_period);
265
266 14 auto dispatch_out = [&](auto in_acc) {
267 using InT = decltype(in_acc);
268
3/6
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
✗ Branch 4 not taken.
✓ Branch 5 taken 1 time.
 7 if (Y.dtype == Dtype::FP32) {
269 7 timestep_embedding_impl<InT, Fp32Accessor>(timesteps.data, Y.data, N, dim, half, log_max_period);
270
0/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 7 } else if (Y.dtype == Dtype::FP16) {
271 ✗ timestep_embedding_impl<InT, Fp16Accessor>(timesteps.data, Y.data, N, dim, half, log_max_period);
272 ✗ } else {
273 ✗ timestep_embedding_impl<InT, Bf16Accessor>(timesteps.data, Y.data, N, dim, half, log_max_period);
274 }
275 7 };
276
277
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 5 times.
 7 if (timesteps.dtype == Dtype::FP32) {
278 5 dispatch_out(Fp32Accessor{});
279
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 7 } else if (timesteps.dtype == Dtype::FP16) {
280 1 dispatch_out(Fp16Accessor{});
281 1 } else {
282 1 dispatch_out(Bf16Accessor{});
283 }
284 7 }
285
286 } // namespace brotensor::detail::cpu
287