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 | #include "cpu_jit.h" | ||
| 15 | |||
| 16 | #include <stdexcept> | ||
| 17 | #include <string> | ||
| 18 | |||
| 19 | namespace brotensor::detail::cpu { | ||
| 20 | |||
| 21 | namespace { | ||
| 22 | |||
| 23 | 1228 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 24 | const char* op, const char* name) { | ||
| 25 |
1/2✓ Branch 0 taken 1228 times.
✗ Branch 1 not taken.
|
1228 | if (t.dtype != Dtype::FP32) { |
| 26 | ✗ | throw std::runtime_error(std::string(op) + ": " + name + | |
| 27 | " must be FP32 (CPU backend is FP32-only)"); | ||
| 28 | } | ||
| 29 | 1228 | } | |
| 30 | |||
| 31 | } // namespace | ||
| 32 | |||
| 33 | 246 | void modulate(const ::brotensor::Tensor& X, const ::brotensor::Tensor& scale, | |
| 34 | const ::brotensor::Tensor& shift, ::brotensor::Tensor& Y) { | ||
| 35 | 246 | check_fp32(X, "modulate", "X"); | |
| 36 | 246 | check_fp32(scale, "modulate", "scale"); | |
| 37 | 246 | check_fp32(shift, "modulate", "shift"); | |
| 38 | 246 | const int L = X.rows; | |
| 39 | 246 | const int D = X.cols; | |
| 40 |
1/2✓ Branch 0 taken 246 times.
✗ Branch 1 not taken.
|
246 | if (scale.size() != D || shift.size() != D) { |
| 41 | ✗ | throw std::runtime_error("modulate: scale and shift must each have X.cols elements"); | |
| 42 | } | ||
| 43 |
4/6✓ Branch 0 taken 240 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 240 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 240 times.
|
246 | if (Y.rows != L || Y.cols != D || Y.dtype != Dtype::FP32) { |
| 44 | 6 | Y.resize(L, D, Dtype::FP32); | |
| 45 | 6 | } | |
| 46 |
1/2✓ Branch 0 taken 246 times.
✗ Branch 1 not taken.
|
246 | if (L * D == 0) return; |
| 47 | |||
| 48 | 246 | const float* Xp = X.host_f32(); | |
| 49 | 246 | const float* sc = scale.host_f32(); | |
| 50 | 246 | const float* sh = shift.host_f32(); | |
| 51 | 246 | float* Yp = Y.host_f32_mut(); | |
| 52 | |||
| 53 |
1/2✓ Branch 0 taken 246 times.
✗ Branch 1 not taken.
|
246 | if (jit::is_jit_available()) { |
| 54 | 246 | jit::modulate(Xp, sc, sh, Yp, L, D); | |
| 55 | 246 | return; | |
| 56 | } | ||
| 57 | |||
| 58 | ✗ | for (int l = 0; l < L; ++l) { | |
| 59 | ✗ | const int base = l * D; | |
| 60 | ✗ | for (int d = 0; d < D; ++d) { | |
| 61 | ✗ | Yp[base + d] = Xp[base + d] * (1.0f + sc[d]) + sh[d]; | |
| 62 | ✗ | } | |
| 63 | ✗ | } | |
| 64 | 246 | } | |
| 65 | |||
| 66 | 245 | void broadcast_mul(const ::brotensor::Tensor& X, const ::brotensor::Tensor& v, | |
| 67 | ::brotensor::Tensor& Y) { | ||
| 68 | 245 | check_fp32(X, "broadcast_mul", "X"); | |
| 69 | 245 | check_fp32(v, "broadcast_mul", "v"); | |
| 70 | 245 | const int L = X.rows; | |
| 71 | 245 | const int D = X.cols; | |
| 72 |
1/2✓ Branch 0 taken 245 times.
✗ Branch 1 not taken.
|
245 | if (v.size() != D) { |
| 73 | ✗ | throw std::runtime_error("broadcast_mul: v must have X.cols elements"); | |
| 74 | } | ||
| 75 |
4/6✓ Branch 0 taken 240 times.
✓ Branch 1 taken 5 times.
✓ Branch 2 taken 240 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 240 times.
|
245 | if (Y.rows != L || Y.cols != D || Y.dtype != Dtype::FP32) { |
| 76 | 5 | Y.resize(L, D, Dtype::FP32); | |
| 77 | 5 | } | |
| 78 |
1/2✓ Branch 0 taken 245 times.
✗ Branch 1 not taken.
|
245 | if (L * D == 0) return; |
| 79 | |||
| 80 | 245 | const float* Xp = X.host_f32(); | |
| 81 | 245 | const float* vp = v.host_f32(); | |
| 82 | 245 | float* Yp = Y.host_f32_mut(); | |
| 83 | |||
| 84 |
1/2✓ Branch 0 taken 245 times.
✗ Branch 1 not taken.
|
245 | if (jit::is_jit_available()) { |
| 85 | 245 | jit::broadcast_mul(Xp, vp, Yp, L, D); | |
| 86 | 245 | return; | |
| 87 | } | ||
| 88 | |||
| 89 | ✗ | for (int l = 0; l < L; ++l) { | |
| 90 | ✗ | const int base = l * D; | |
| 91 | ✗ | for (int d = 0; d < D; ++d) { | |
| 92 | ✗ | Yp[base + d] = Xp[base + d] * vp[d]; | |
| 93 | ✗ | } | |
| 94 | ✗ | } | |
| 95 | 245 | } | |
| 96 | |||
| 97 | } // namespace brotensor::detail::cpu | ||
| 98 |