src/cpu/codec_quant.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU codec quantization ops (brosoundml CHUNK 5, family D) ────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementations of the quantization bottlenecks of neural | ||
| 4 | // audio codecs: | ||
| 5 | // vq_encode_forward / vq_encode_backward — vector quantization (EnCodec / | ||
| 6 | // DAC residual-VQ encode step) | ||
| 7 | // fsq_quantize_forward / fsq_quantize_backward — finite scalar quantization | ||
| 8 | // (NanoCodec) | ||
| 9 | // | ||
| 10 | // ── INT32 outputs ─────────────────────────────────────────────────────────── | ||
| 11 | // vq_encode_forward.indices — (N, 1) INT32 codeword indices. | ||
| 12 | // fsq_quantize_forward.packed_indices — (N, 1) INT32 mixed-radix codes. | ||
| 13 | // INT32 is a pure storage carrier; these tensors are resized AND dtype-set | ||
| 14 | // to INT32, and accessed via the dtype-agnostic host_raw / host_raw_mut | ||
| 15 | // accessors cast to int32_t* (host_f32 throws on a non-FP32 dtype). | ||
| 16 | // | ||
| 17 | // ── Accumulation ──────────────────────────────────────────────────────────── | ||
| 18 | // *_forward — all outputs OVERWRITTEN. | ||
| 19 | // vq_encode_backward — dX OVERWRITTEN (straight-through identity). | ||
| 20 | // fsq_quantize_backward — dX OVERWRITTEN (straight-through identity). | ||
| 21 | // Neither backward accumulates and neither produces a codebook gradient — | ||
| 22 | // they are purely the encoder STE passthrough. See ops.h for the rationale. | ||
| 23 | // | ||
| 24 | // CPU is FP32-only; all arithmetic is FP32. | ||
| 25 | |||
| 26 | #include <brotensor/tensor.h> | ||
| 27 | |||
| 28 | #include <cmath> | ||
| 29 | #include <cstddef> | ||
| 30 | #include <cstdint> | ||
| 31 | #include <stdexcept> | ||
| 32 | #include <string> | ||
| 33 | |||
| 34 | namespace brotensor::detail::cpu { | ||
| 35 | |||
| 36 | namespace { | ||
| 37 | |||
| 38 | ✗ | [[noreturn]] void fail(const char* op, const std::string& reason) { | |
| 39 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason); | |
| 40 | ✗ | } | |
| 41 | |||
| 42 | 25 | void require_fp32(const char* op, const ::brotensor::Tensor& t, | |
| 43 | const char* name) { | ||
| 44 |
1/2✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
|
25 | if (t.dtype != ::brotensor::Dtype::FP32) { |
| 45 | ✗ | fail(op, std::string(name) + " must be FP32 (CPU backend is FP32-only)"); | |
| 46 | } | ||
| 47 | 25 | } | |
| 48 | |||
| 49 | 5 | void require_int32(const char* op, const ::brotensor::Tensor& t, | |
| 50 | const char* name) { | ||
| 51 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (t.dtype != ::brotensor::Dtype::INT32) { |
| 52 | ✗ | fail(op, std::string(name) + " must be INT32"); | |
| 53 | } | ||
| 54 | 5 | } | |
| 55 | |||
| 56 | } // namespace | ||
| 57 | |||
| 58 | // ─── vq_encode ────────────────────────────────────────────────────────────── | ||
| 59 | |||
| 60 | 6 | void vq_encode_forward(const ::brotensor::Tensor& x, | |
| 61 | const ::brotensor::Tensor& codebook, | ||
| 62 | ::brotensor::Tensor& indices, | ||
| 63 | ::brotensor::Tensor& quantized) { | ||
| 64 | 6 | require_fp32("vq_encode_forward", x, "x"); | |
| 65 | 6 | require_fp32("vq_encode_forward", codebook, "codebook"); | |
| 66 | 6 | const int N = x.rows; | |
| 67 | 6 | const int D = x.cols; | |
| 68 | 6 | const int K = codebook.rows; | |
| 69 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (codebook.cols != D) { |
| 70 | ✗ | fail("vq_encode_forward", "codebook must have the same column count as x"); | |
| 71 | } | ||
| 72 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
6 | if (K == 0 && N != 0) { |
| 73 | ✗ | fail("vq_encode_forward", "codebook must have at least one codeword"); | |
| 74 | } | ||
| 75 | |||
| 76 | // indices: (N, 1) INT32 — resize AND dtype-set. | ||
| 77 |
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 (indices.rows != N || indices.cols != 1 || |
| 78 | ✗ | indices.dtype != ::brotensor::Dtype::INT32) { | |
| 79 | 6 | indices.resize(N, 1, ::brotensor::Dtype::INT32); | |
| 80 | 6 | } | |
| 81 | // quantized: (N, D) FP32. | ||
| 82 |
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 (quantized.rows != N || quantized.cols != D || |
| 83 | ✗ | quantized.dtype != ::brotensor::Dtype::FP32) { | |
| 84 | 6 | quantized.resize(N, D, ::brotensor::Dtype::FP32); | |
| 85 | 6 | } | |
| 86 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | if (N == 0) return; |
| 87 | |||
| 88 | 6 | const float* xp = x.host_f32(); | |
| 89 | 6 | const float* cp = codebook.host_f32(); | |
| 90 | 6 | int32_t* ip = static_cast<int32_t*>(indices.host_raw_mut()); | |
| 91 | 6 | float* qp = quantized.host_f32_mut(); | |
| 92 | |||
| 93 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 73 times.
|
79 | for (int n = 0; n < N; ++n) { |
| 94 | 73 | const float* x_row = xp + static_cast<std::size_t>(n) * D; | |
| 95 | 73 | float best_d2 = 3.4028235e38f; // +FLT_MAX | |
| 96 | 73 | int best_k = 0; | |
| 97 |
2/2✓ Branch 0 taken 8845 times.
✓ Branch 1 taken 73 times.
|
8918 | for (int k = 0; k < K; ++k) { |
| 98 | 8845 | const float* c_row = cp + static_cast<std::size_t>(k) * D; | |
| 99 | 8845 | float d2 = 0.0f; | |
| 100 |
2/2✓ Branch 0 taken 137249 times.
✓ Branch 1 taken 8845 times.
|
146094 | for (int j = 0; j < D; ++j) { |
| 101 | 137249 | const float diff = x_row[j] - c_row[j]; | |
| 102 | 137249 | d2 += diff * diff; | |
| 103 | 137249 | } | |
| 104 | // Strict `<` keeps the lowest index on ties. | ||
| 105 |
2/2✓ Branch 0 taken 8537 times.
✓ Branch 1 taken 308 times.
|
8845 | if (d2 < best_d2) { best_d2 = d2; best_k = k; } |
| 106 | 8845 | } | |
| 107 | 73 | ip[n] = static_cast<int32_t>(best_k); | |
| 108 | 73 | const float* c_best = cp + static_cast<std::size_t>(best_k) * D; | |
| 109 | 73 | float* q_row = qp + static_cast<std::size_t>(n) * D; | |
| 110 |
2/2✓ Branch 0 taken 750 times.
✓ Branch 1 taken 73 times.
|
823 | for (int j = 0; j < D; ++j) q_row[j] = c_best[j]; |
| 111 | 73 | } | |
| 112 | 6 | } | |
| 113 | |||
| 114 | 4 | void vq_encode_backward(const ::brotensor::Tensor& dQuantized, | |
| 115 | ::brotensor::Tensor& dX) { | ||
| 116 | 4 | require_fp32("vq_encode_backward", dQuantized, "dQuantized"); | |
| 117 | // Straight-through estimator: the argmin is non-differentiable, so the | ||
| 118 | // gradient passes through unchanged. dX is OVERWRITTEN, not accumulated. | ||
| 119 |
4/6✓ Branch 0 taken 1 time.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 1 time.
|
4 | if (dX.rows != dQuantized.rows || dX.cols != dQuantized.cols || |
| 120 | 1 | dX.dtype != ::brotensor::Dtype::FP32) { | |
| 121 | 3 | dX.resize(dQuantized.rows, dQuantized.cols, ::brotensor::Dtype::FP32); | |
| 122 | 3 | } | |
| 123 | 4 | const int total = dQuantized.size(); | |
| 124 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | if (total == 0) return; |
| 125 | 4 | const float* dqp = dQuantized.host_f32(); | |
| 126 | 4 | float* dxp = dX.host_f32_mut(); | |
| 127 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | if (dxp == dqp) return; // dX aliases dQuantized |
| 128 |
2/2✓ Branch 0 taken 4 times.
✓ Branch 1 taken 181 times.
|
185 | for (int i = 0; i < total; ++i) dxp[i] = dqp[i]; |
| 129 | 4 | } | |
| 130 | |||
| 131 | // ─── fsq_quantize ─────────────────────────────────────────────────────────── | ||
| 132 | |||
| 133 | 5 | void fsq_quantize_forward(const ::brotensor::Tensor& x, | |
| 134 | const ::brotensor::Tensor& levels, | ||
| 135 | ::brotensor::Tensor& quantized, | ||
| 136 | ::brotensor::Tensor& packed_indices) { | ||
| 137 | 5 | require_fp32("fsq_quantize_forward", x, "x"); | |
| 138 | 5 | require_int32("fsq_quantize_forward", levels, "levels"); | |
| 139 | 5 | const int N = x.rows; | |
| 140 | 5 | const int D = x.cols; | |
| 141 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (levels.size() != D) { |
| 142 | ✗ | fail("fsq_quantize_forward", "levels must have D elements (one per column of x)"); | |
| 143 | } | ||
| 144 | |||
| 145 | // quantized: (N, D) FP32. | ||
| 146 |
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 (quantized.rows != N || quantized.cols != D || |
| 147 | ✗ | quantized.dtype != ::brotensor::Dtype::FP32) { | |
| 148 | 5 | quantized.resize(N, D, ::brotensor::Dtype::FP32); | |
| 149 | 5 | } | |
| 150 | // packed_indices: (N, 1) INT32 — resize AND dtype-set. | ||
| 151 |
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 (packed_indices.rows != N || packed_indices.cols != 1 || |
| 152 | ✗ | packed_indices.dtype != ::brotensor::Dtype::INT32) { | |
| 153 | 5 | packed_indices.resize(N, 1, ::brotensor::Dtype::INT32); | |
| 154 | 5 | } | |
| 155 | |||
| 156 | 5 | const int32_t* Lp = static_cast<const int32_t*>(levels.host_raw()); | |
| 157 | // Validate level counts up front (every L_d >= 2). | ||
| 158 |
2/2✓ Branch 0 taken 24 times.
✓ Branch 1 taken 5 times.
|
29 | for (int d = 0; d < D; ++d) { |
| 159 |
1/2✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
|
24 | if (Lp[d] < 2) { |
| 160 | ✗ | fail("fsq_quantize_forward", "every level count must be >= 2"); | |
| 161 | } | ||
| 162 | 24 | } | |
| 163 |
2/4✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 5 times.
|
5 | if (N == 0 || D == 0) return; |
| 164 | |||
| 165 | 5 | const float* xp = x.host_f32(); | |
| 166 | 5 | float* qp = quantized.host_f32_mut(); | |
| 167 | 5 | int32_t* pp = static_cast<int32_t*>(packed_indices.host_raw_mut()); | |
| 168 | |||
| 169 |
2/2✓ Branch 0 taken 64 times.
✓ Branch 1 taken 5 times.
|
69 | for (int n = 0; n < N; ++n) { |
| 170 | 64 | const float* x_row = xp + static_cast<std::size_t>(n) * D; | |
| 171 | 64 | float* q_row = qp + static_cast<std::size_t>(n) * D; | |
| 172 | // Mixed-radix pack: dimension 0 is the least-significant digit. | ||
| 173 | // packed = i_0 + L_0 * (i_1 + L_1 * (i_2 + ...)) | ||
| 174 | // Build from the most-significant digit down via Horner's scheme. | ||
| 175 | 64 | long long packed = 0; | |
| 176 |
2/2✓ Branch 0 taken 451 times.
✓ Branch 1 taken 64 times.
|
515 | for (int d = D - 1; d >= 0; --d) { |
| 177 | 451 | const int L = Lp[d]; | |
| 178 | 451 | const float h = static_cast<float>(L - 1) * 0.5f; // half-width | |
| 179 | // 1. clamp into [-1, 1]. | ||
| 180 | 451 | float v = x_row[d]; | |
| 181 |
2/2✓ Branch 0 taken 40 times.
✓ Branch 1 taken 411 times.
|
451 | if (v < -1.0f) v = -1.0f; |
| 182 |
2/2✓ Branch 0 taken 370 times.
✓ Branch 1 taken 41 times.
|
411 | else if (v > 1.0f) v = 1.0f; |
| 183 | // 2. map to a level index in [0, L-1]. | ||
| 184 | 451 | float idx_f = std::round((v + 1.0f) * 0.5f * static_cast<float>(L - 1)); | |
| 185 | 451 | int idx = static_cast<int>(idx_f); | |
| 186 |
1/2✓ Branch 0 taken 451 times.
✗ Branch 1 not taken.
|
451 | if (idx < 0) idx = 0; |
| 187 |
1/2✓ Branch 0 taken 451 times.
✗ Branch 1 not taken.
|
451 | else if (idx > L - 1) idx = L - 1; |
| 188 | // 3. dequantize back into [-1, 1]. | ||
| 189 | 451 | q_row[d] = static_cast<float>(idx) / h - 1.0f; | |
| 190 | 902 | packed = packed * static_cast<long long>(L) + | |
| 191 | 451 | static_cast<long long>(idx); | |
| 192 | 451 | } | |
| 193 | 64 | pp[n] = static_cast<int32_t>(packed); | |
| 194 | 64 | } | |
| 195 | 5 | } | |
| 196 | |||
| 197 | 4 | void fsq_quantize_backward(const ::brotensor::Tensor& dQuantized, | |
| 198 | ::brotensor::Tensor& dX) { | ||
| 199 | 4 | require_fp32("fsq_quantize_backward", dQuantized, "dQuantized"); | |
| 200 | // Straight-through estimator: round is non-differentiable, so the gradient | ||
| 201 | // passes through unchanged. dX is OVERWRITTEN, not accumulated. | ||
| 202 |
4/6✓ Branch 0 taken 1 time.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 1 time.
|
4 | if (dX.rows != dQuantized.rows || dX.cols != dQuantized.cols || |
| 203 | 1 | dX.dtype != ::brotensor::Dtype::FP32) { | |
| 204 | 3 | dX.resize(dQuantized.rows, dQuantized.cols, ::brotensor::Dtype::FP32); | |
| 205 | 3 | } | |
| 206 | 4 | const int total = dQuantized.size(); | |
| 207 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | if (total == 0) return; |
| 208 | 4 | const float* dqp = dQuantized.host_f32(); | |
| 209 | 4 | float* dxp = dX.host_f32_mut(); | |
| 210 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | if (dxp == dqp) return; // dX aliases dQuantized |
| 211 |
2/2✓ Branch 0 taken 4 times.
✓ Branch 1 taken 187 times.
|
191 | for (int i = 0; i < total; ++i) dxp[i] = dqp[i]; |
| 212 | 4 | } | |
| 213 | |||
| 214 | } // namespace brotensor::detail::cpu | ||
| 215 |