src/cpu/embedding.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU embedding lookup ops (CHUNK 1) ──────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementations. Ports the GPU kernels in | ||
| 4 | // src/cuda/embedding.cu — FP32 path only. | ||
| 5 | // | ||
| 6 | // `d_idx` is a host int32 pointer on the CPU backend (the dispatcher keeps | ||
| 7 | // raw-pointer args on the host). | ||
| 8 | // | ||
| 9 | // embedding_lookup_forward — gather B rows from `table` by index. | ||
| 10 | // embedding_lookup_backward — scatter-accumulate dOut rows into dTable. | ||
| 11 | // ACCUMULATES (+=); caller zeros dTable. | ||
| 12 | |||
| 13 | #include <brotensor/tensor.h> | ||
| 14 | |||
| 15 | #include <cstddef> | ||
| 16 | #include <cstdint> | ||
| 17 | #include <stdexcept> | ||
| 18 | |||
| 19 | namespace brotensor::detail::cpu { | ||
| 20 | |||
| 21 | 8 | void embedding_lookup_forward(const ::brotensor::Tensor& table, | |
| 22 | const int32_t* d_idx, int B, | ||
| 23 | ::brotensor::Tensor& out) { | ||
| 24 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (table.dtype != ::brotensor::Dtype::FP32) { |
| 25 | ✗ | throw std::runtime_error( | |
| 26 | "brotensor: embedding_lookup_forward: CPU backend requires " | ||
| 27 | "FP32 table (quantized tables must be dequantized first)"); | ||
| 28 | } | ||
| 29 | 8 | const int D = table.cols; | |
| 30 |
4/6✓ Branch 0 taken 7 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 7 times.
|
8 | if (out.rows != B || out.cols != D || |
| 31 | 7 | out.dtype != ::brotensor::Dtype::FP32) { | |
| 32 | 1 | out.resize(B, D, ::brotensor::Dtype::FP32); | |
| 33 | 1 | } | |
| 34 | 8 | const int total = B * D; | |
| 35 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (total == 0) return; |
| 36 | |||
| 37 | 8 | const float* tp = table.host_f32(); | |
| 38 | 8 | float* op = out.host_f32_mut(); | |
| 39 |
2/2✓ Branch 0 taken 94 times.
✓ Branch 1 taken 8 times.
|
102 | for (int b = 0; b < B; ++b) { |
| 40 | 94 | const int row = d_idx[b]; | |
| 41 | 94 | const float* src = tp + static_cast<std::size_t>(row) * D; | |
| 42 | 94 | float* dst = op + static_cast<std::size_t>(b) * D; | |
| 43 |
2/2✓ Branch 0 taken 1149 times.
✓ Branch 1 taken 94 times.
|
1243 | for (int j = 0; j < D; ++j) dst[j] = src[j]; |
| 44 | 94 | } | |
| 45 | 8 | } | |
| 46 | |||
| 47 | 10 | void embedding_lookup_backward(const ::brotensor::Tensor& dOut, | |
| 48 | const int32_t* d_idx, int B, | ||
| 49 | ::brotensor::Tensor& dTable) { | ||
| 50 | 10 | const int D = dTable.cols; | |
| 51 | 10 | const int total = B * D; | |
| 52 |
1/2✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
|
10 | if (total == 0) return; |
| 53 | |||
| 54 | 10 | const float* dop = dOut.host_f32(); | |
| 55 | 10 | float* dtp = dTable.host_f32_mut(); | |
| 56 | // Scatter-accumulate: caller zeros dTable beforehand. Repeated indices | ||
| 57 | // accumulate, matching the GPU atomicAdd. | ||
| 58 |
2/2✓ Branch 0 taken 127 times.
✓ Branch 1 taken 10 times.
|
137 | for (int b = 0; b < B; ++b) { |
| 59 | 127 | const int row = d_idx[b]; | |
| 60 | 127 | const float* src = dop + static_cast<std::size_t>(b) * D; | |
| 61 | 127 | float* dst = dtp + static_cast<std::size_t>(row) * D; | |
| 62 |
2/2✓ Branch 0 taken 1648 times.
✓ Branch 1 taken 127 times.
|
1775 | for (int j = 0; j < D; ++j) dst[j] += src[j]; |
| 63 | 127 | } | |
| 64 | 10 | } | |
| 65 | |||
| 66 | } // namespace brotensor::detail::cpu | ||
| 67 |