src/cpu/modulate.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU AdaLN modulation ops ────────────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementations of the DiT / SD3 / Flux broadcast-affine | ||
| 4 | // primitives. Ports src/cuda/modulate.cu — FP32 path only (the CPU backend is | ||
| 5 | // FP32-only per CLAUDE.md). | ||
| 6 | // | ||
| 7 | // ── modulate ── Y[l, d] = X[l, d] * (1 + scale[d]) + shift[d] | ||
| 8 | // ── broadcast_mul ── Y[l, d] = X[l, d] * v[d] | ||
| 9 | // | ||
| 10 | // scale / shift / v are length-D vectors broadcast across every token row. | ||
| 11 | // Both ops fully OVERWRITE their output. | ||
| 12 | |||
| 13 | #include <brotensor/tensor.h> | ||
| 14 | |||
| 15 | #include <stdexcept> | ||
| 16 | #include <string> | ||
| 17 | |||
| 18 | namespace brotensor::detail::cpu { | ||
| 19 | |||
| 20 | namespace { | ||
| 21 | |||
| 22 | 28 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 23 | const char* op, const char* name) { | ||
| 24 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | if (t.dtype != Dtype::FP32) { |
| 25 | ✗ | throw std::runtime_error(std::string(op) + ": " + name + | |
| 26 | " must be FP32 (CPU backend is FP32-only)"); | ||
| 27 | } | ||
| 28 | 28 | } | |
| 29 | |||
| 30 | } // namespace | ||
| 31 | |||
| 32 | 6 | void modulate(const ::brotensor::Tensor& X, const ::brotensor::Tensor& scale, | |
| 33 | const ::brotensor::Tensor& shift, ::brotensor::Tensor& Y) { | ||
| 34 | 6 | check_fp32(X, "modulate", "X"); | |
| 35 | 6 | check_fp32(scale, "modulate", "scale"); | |
| 36 | 6 | check_fp32(shift, "modulate", "shift"); | |
| 37 | 6 | const int L = X.rows; | |
| 38 | 6 | const int D = X.cols; | |
| 39 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (scale.size() != D || shift.size() != D) { |
| 40 | ✗ | throw std::runtime_error("modulate: scale and shift must each have X.cols elements"); | |
| 41 | } | ||
| 42 |
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 != L || Y.cols != D || Y.dtype != Dtype::FP32) { |
| 43 | 6 | Y.resize(L, D, Dtype::FP32); | |
| 44 | 6 | } | |
| 45 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (L * D == 0) return; |
| 46 | |||
| 47 | 6 | const float* Xp = X.host_f32(); | |
| 48 | 6 | const float* sc = scale.host_f32(); | |
| 49 | 6 | const float* sh = shift.host_f32(); | |
| 50 | 6 | float* Yp = Y.host_f32_mut(); | |
| 51 | |||
| 52 |
2/2✓ Branch 0 taken 29 times.
✓ Branch 1 taken 6 times.
|
35 | for (int l = 0; l < L; ++l) { |
| 53 | 29 | const int base = l * D; | |
| 54 |
2/2✓ Branch 0 taken 5035 times.
✓ Branch 1 taken 29 times.
|
5064 | for (int d = 0; d < D; ++d) { |
| 55 | 5035 | Yp[base + d] = Xp[base + d] * (1.0f + sc[d]) + sh[d]; | |
| 56 | 5035 | } | |
| 57 | 29 | } | |
| 58 | 6 | } | |
| 59 | |||
| 60 | 5 | void broadcast_mul(const ::brotensor::Tensor& X, const ::brotensor::Tensor& v, | |
| 61 | ::brotensor::Tensor& Y) { | ||
| 62 | 5 | check_fp32(X, "broadcast_mul", "X"); | |
| 63 | 5 | check_fp32(v, "broadcast_mul", "v"); | |
| 64 | 5 | const int L = X.rows; | |
| 65 | 5 | const int D = X.cols; | |
| 66 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (v.size() != D) { |
| 67 | ✗ | throw std::runtime_error("broadcast_mul: v must have X.cols elements"); | |
| 68 | } | ||
| 69 |
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 != L || Y.cols != D || Y.dtype != Dtype::FP32) { |
| 70 | 5 | Y.resize(L, D, Dtype::FP32); | |
| 71 | 5 | } | |
| 72 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (L * D == 0) return; |
| 73 | |||
| 74 | 5 | const float* Xp = X.host_f32(); | |
| 75 | 5 | const float* vp = v.host_f32(); | |
| 76 | 5 | float* Yp = Y.host_f32_mut(); | |
| 77 | |||
| 78 |
2/2✓ Branch 0 taken 28 times.
✓ Branch 1 taken 5 times.
|
33 | for (int l = 0; l < L; ++l) { |
| 79 | 28 | const int base = l * D; | |
| 80 |
2/2✓ Branch 0 taken 5034 times.
✓ Branch 1 taken 28 times.
|
5062 | for (int d = 0; d < D; ++d) { |
| 81 | 5034 | Yp[base + d] = Xp[base + d] * vp[d]; | |
| 82 | 5034 | } | |
| 83 | 28 | } | |
| 84 | 5 | } | |
| 85 | |||
| 86 | } // namespace brotensor::detail::cpu | ||
| 87 |