| //===- compact.cc -----------------------------------------------*- C++ -*-===// | |
| // SPDX-License-Identifier: AGPL-3.0-or-later | |
| // Copyright (C) 2026 Scott Boudreaux / Elyan Labs. Commercial license: COMMERCIAL.md | |
| // | |
| // open-xdna :: full STREAM COMPACTION on XDNA1 — dense survivor pack. | |
| // | |
| // out = [ all x where x >= tau, in order ] ++ [ zeros ] | |
| // | |
| // This is the top-k pack that makes the prune actually shrink downstream work: survivors are | |
| // packed contiguously at the front so the next matmul runs on a dense [0:count] slice. | |
| // | |
| // NOTE ON METHOD: compaction is NOT elementwise — out[k] depends on the running survivor count | |
| // (a scan + gather), so it can't tile-stream. It runs on the AIE *scalar* unit over the whole | |
| // in-tile vector. The vectorized form (Hillis-Steele prefix-sum of the mask + aie::shuffle | |
| // runtime gather) is the performance optimization; this is the correct baseline. | |
| //===----------------------------------------------------------------------===// | |
| void compact(bfloat16 *restrict a, bfloat16 *restrict c, int32_t *tau_q, | |
| const int N) { | |
| const float tau = ((float)(*tau_q)) / 256.0f; | |
| event0(); | |
| int k = 0; | |
| for (int i = 0; i < N; i++) { // scan + gather (AIE scalar unit) | |
| if ((float)a[i] >= tau) { | |
| c[k] = a[i]; // pack survivor to the dense front | |
| k++; | |
| } | |
| } | |
| for (int j = k; j < N; j++) { // zero-pad the tail | |
| c[j] = (bfloat16)0.0f; | |
| } | |
| event1(); | |
| return; | |
| } | |
| extern "C" { | |
| void bf16_compact(bfloat16 *a_in, bfloat16 *c_out, int32_t *tau_q, int32_t N) { | |
| compact(a_in, c_out, tau_q, N); | |
| } | |
| } // extern "C" | |