#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <algorithm>
#include <array>
#include <chrono>
#include <cstdint>
#include <limits>
#include <queue>
#include <set>
#include <stdexcept>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>
#include "parallel.hpp"
#ifdef _OPENMP
#include <omp.h>
#endif
#if defined(_OPENMP) && defined(__GLIBCXX__)
#include <parallel/algorithm>
#define AC2_PARALLEL_SORT __gnu_parallel::sort
#else
#define AC2_PARALLEL_SORT std::sort
#endif
namespace py = pybind11;
namespace {
constexpr int kMaxLookahead = 32;
constexpr size_t kCertifyBatch = 32;
constexpr size_t kApplyGrain = 256; // minimum candidate pairs per parallel apply group
constexpr size_t kShards = 64; // rule-table shards updated concurrently after a merge
inline uint64_t mix(uint64_t key) { return key * 0x9E3779B97F4A7C15ull; }
inline size_t shard_of(uint64_t key) { return static_cast<size_t>(mix(key) >> 58); }
// Small open-addressing accumulator for rule-count deltas; clearing costs O(used).
class DeltaMap {
public:
void add(uint64_t key, int64_t delta) {
if ((filled_.size() + 1) * 2 > keys_.size()) grow();
size_t h = slot(key);
while (keys_[h] != kEmpty && keys_[h] != key) h = (h + 1) & (keys_.size() - 1);
if (keys_[h] == kEmpty) {
keys_[h] = key;
filled_.push_back(static_cast<uint32_t>(h));
}
values_[h] += delta;
}
template <class F>
void each(F&& visit) const {
for (uint32_t h : filled_) visit(keys_[h], values_[h]);
}
void clear() {
for (uint32_t h : filled_) {
keys_[h] = kEmpty;
values_[h] = 0;
}
filled_.clear();
}
private:
static constexpr uint64_t kEmpty = ~0ull;
size_t slot(uint64_t key) const {
return static_cast<size_t>(mix(key) >> 20) & (keys_.size() - 1);
}
void grow() {
std::vector<uint64_t> keys = std::move(keys_);
std::vector<int64_t> values = std::move(values_);
std::vector<uint32_t> filled = std::move(filled_);
const size_t capacity = std::max<size_t>(64, keys.size() * 2);
keys_.assign(capacity, kEmpty);
values_.assign(capacity, 0);
filled_.clear();
for (uint32_t h : filled) add(keys[h], values[h]);
}
std::vector<uint64_t> keys_;
std::vector<int64_t> values_;
std::vector<uint32_t> filled_;
};
inline uint64_t rule_key(int32_t left, int32_t right, int32_t offset) {
return (static_cast<uint64_t>(left) << 40) | (static_cast<uint64_t>(right) << 16) |
static_cast<uint64_t>(offset);
}
inline void unpack(uint64_t key, int32_t& left, int32_t& right, int32_t& offset) {
left = static_cast<int32_t>(key >> 40);
right = static_cast<int32_t>((key >> 16) & 0xffffffull);
offset = static_cast<int32_t>(key & 0xffffull);
}
// Immutable-per-token geometry shared by the trainer, encoder and decoder.
struct Tables {
int32_t dims = 0; // physical slots including E
int32_t e_dim = 0; // index of E
int32_t lookahead = 0; // maximum token span in frames
std::vector<uint32_t> masks; // token * lookahead -> support mask
std::vector<uint32_t> emask; // token -> bit per offset holding E
std::vector<int8_t> first_from; // token * lookahead + k -> first support offset >= k
std::vector<int32_t> span, rdim;
std::vector<std::vector<std::pair<int32_t, int32_t>>> pattern; // (offset, spatial token)
int32_t base = 0;
void init(const std::vector<uint32_t>& spatial_support, int32_t dims_, int32_t lookahead_) {
dims = dims_;
e_dim = dims_ - 1;
lookahead = lookahead_;
if (lookahead < 1 || lookahead > kMaxLookahead) throw std::invalid_argument("lookahead");
if (dims < 2 || dims > 16) throw std::invalid_argument("dims must be in [2, 16]");
base = static_cast<int32_t>(spatial_support.size());
for (int32_t t = 0; t < base; ++t) {
uint32_t s = spatial_support[t];
if (!s) throw std::invalid_argument("spatial token without support");
push({{0, t}}, std::vector<uint32_t>{s});
}
}
void push(std::vector<std::pair<int32_t, int32_t>> pat, const std::vector<uint32_t>& m) {
const size_t token = span.size();
masks.resize((token + 1) * lookahead, 0);
uint32_t e = 0;
int32_t sp = 0;
for (size_t o = 0; o < m.size(); ++o) {
masks[token * lookahead + o] = m[o];
if (m[o]) sp = static_cast<int32_t>(o) + 1;
if (m[o] >> e_dim & 1u) e |= 1u << o;
}
emask.push_back(e);
first_from.resize((token + 1) * lookahead, -1);
for (int32_t k = sp - 1, next = -1; k >= 0; --k) {
if (masks[token * lookahead + k]) next = k;
first_from[token * lookahead + k] = static_cast<int8_t>(next);
}
span.push_back(sp);
rdim.push_back(__builtin_ctz(m[0]));
pattern.push_back(std::move(pat));
}
int32_t size() const { return static_cast<int32_t>(span.size()); }
uint32_t mask(int32_t token, int32_t o) const {
return masks[static_cast<size_t>(token) * lookahead + o];
}
// Merge geometry; returns false when the union would exceed the lookahead or overlap.
bool compose(int32_t a, int32_t b, int32_t off, std::vector<uint32_t>& out) const {
if (off < 0 || off >= lookahead) return false;
const int32_t sp = std::max(span[a], off + span[b]);
if (sp > lookahead) return false;
out.assign(sp, 0);
for (int32_t o = 0; o < span[a]; ++o) out[o] = mask(a, o);
for (int32_t o = 0; o < span[b]; ++o) {
if (out[o + off] & mask(b, o)) return false;
out[o + off] |= mask(b, o);
}
return true;
}
int32_t add_merge(int32_t a, int32_t b, int32_t off) {
std::vector<uint32_t> m;
if (!compose(a, b, off, m)) throw std::invalid_argument("illegal temporal merge");
auto pat = pattern[a];
for (auto [o, t] : pattern[b]) pat.push_back({o + off, t});
std::sort(pat.begin(), pat.end());
push(std::move(pat), m);
return size() - 1;
}
};
using Key = std::tuple<int64_t, int32_t, int64_t, int32_t>;
using InstanceId = uint32_t;
constexpr InstanceId kNoInstance = UINT32_MAX;
// Mutable instance state over a global frame axis. Used per chunk by the encoder and
// globally by the trainer.
struct Instances {
const Tables* tb = nullptr;
int64_t frames = 0;
std::vector<int32_t> tok;
std::vector<int64_t> root;
std::vector<uint8_t> alive;
std::vector<InstanceId> member; // frame * dims + dim -> instance
void reset(const Tables* tables, int64_t frame_count) {
tb = tables;
frames = frame_count;
tok.clear();
root.clear();
alive.clear();
member.assign(static_cast<size_t>(frame_count) * tables->dims, kNoInstance);
}
InstanceId cell(int64_t f, int32_t d) const {
return member[static_cast<size_t>(f) * tb->dims + d];
}
Key key(InstanceId i) const {
const uint32_t e = tb->emask[tok[i]];
return Key(root[i], e ? 1 : 0, e ? root[i] + __builtin_ctz(e) : 0, tb->rdim[tok[i]]);
}
InstanceId add(int32_t token, int64_t r) {
if (tok.size() >= kNoInstance) throw std::length_error("instance ids");
const InstanceId i = static_cast<InstanceId>(tok.size());
tok.push_back(token);
root.push_back(r);
alive.push_back(1);
place(i);
return i;
}
// Reserves dead slots that disjoint frame ranges then fill concurrently via ``place``.
InstanceId reserve(size_t count) {
const size_t base = tok.size();
if (base + count > static_cast<size_t>(kNoInstance)) throw std::length_error("instance ids");
tok.resize(base + count, 0);
root.resize(base + count, 0);
alive.resize(base + count, 0);
return static_cast<InstanceId>(base);
}
void place(InstanceId i) {
const int32_t token = tok[i];
const int64_t r = root[i];
alive[i] = 1;
for (int32_t o = 0; o < tb->span[token]; ++o) {
uint32_t m = tb->mask(token, o);
while (m) {
const int32_t d = __builtin_ctz(m);
m &= m - 1;
InstanceId& slot = member[static_cast<size_t>(r + o) * tb->dims + d];
if (slot != kNoInstance) throw std::runtime_error("instance overlap");
slot = i;
}
}
}
void kill(InstanceId i) {
alive[i] = 0;
const int32_t token = tok[i];
for (int32_t o = 0; o < tb->span[token]; ++o) {
uint32_t m = tb->mask(token, o);
while (m) {
const int32_t d = __builtin_ctz(m);
m &= m - 1;
member[static_cast<size_t>(root[i] + o) * tb->dims + d] = kNoInstance;
}
}
}
bool first(InstanceId a, InstanceId b) const {
return std::make_pair(root[a], tb->rdim[tok[a]]) < std::make_pair(root[b], tb->rdim[tok[b]]);
}
// Ordered candidate pairs (a, b, offset) involving i inside the lookahead window.
// Each neighbour is emitted once, at its first support cell inside the window.
template <class F>
void partners(InstanceId i, bool allow_zero, F&& emit) const {
const int64_t r = root[i];
const int32_t L = tb->lookahead;
const int64_t lo = std::max<int64_t>(0, r - L + 1), hi = std::min<int64_t>(frames - 1, r + L - 1);
for (int64_t f = lo; f <= hi; ++f) {
for (int32_t d = 0; d < tb->dims; ++d) {
const InstanceId j = cell(f, d);
if (j == kNoInstance || j == i) continue;
const int32_t tj = tok[j];
const int64_t o = f - root[j];
const int64_t k = std::max<int64_t>(0, lo - root[j]);
if (tb->first_from[static_cast<size_t>(tj) * L + k] != o) continue;
if (__builtin_ctz(tb->mask(tj, static_cast<int32_t>(o))) != d) continue;
const bool ij = first(i, j);
const InstanceId a = ij ? i : j, b = ij ? j : i;
const int64_t off = root[b] - root[a];
if (off >= L || (off == 0 && !allow_zero)) continue;
if (std::max<int64_t>(tb->span[tok[a]], off + tb->span[tok[b]]) > L) continue;
emit(a, b, static_cast<int32_t>(off));
}
}
}
// Would merging a (rooted first) with b keep the stream decodable?
bool valid(InstanceId a, InstanceId b, int32_t off) const {
const int32_t ta = tok[a], tb_ = tok[b];
const int32_t sp = std::max(tb->span[ta], off + tb->span[tb_]);
uint32_t emask = tb->emask[ta] | (tb->emask[tb_] << off);
const Key km(root[a], emask ? 1 : 0, emask ? root[a] + __builtin_ctz(emask) : 0,
tb->rdim[ta]);
for (int32_t o = 0; o < sp; ++o) {
uint32_t m = (o < tb->span[ta] ? tb->mask(ta, o) : 0) |
(o >= off && o - off < tb->span[tb_] ? tb->mask(tb_, o - off) : 0);
if (!m) continue;
const int64_t u = root[a] + o;
const InstanceId h = cell(u, tb->e_dim);
if (h != kNoInstance && h != a && h != b && !(key(h) > km)) return false;
if (emask >> o & 1u) {
for (int32_t d = 0; d < tb->dims; ++d) {
const InstanceId j = cell(u, d);
if (j != kNoInstance && j != a && j != b && !(key(j) < km)) return false;
}
}
}
return true;
}
// Right partner of left instance a under rule (left, right, off), or -1.
InstanceId right_of(InstanceId a, int32_t right, int32_t off) const {
const int64_t f = root[a] + off;
if (f >= frames) return kNoInstance;
const InstanceId b = cell(f, tb->rdim[right]);
if (b == kNoInstance || b == a || !alive[b] || tok[b] != right || root[b] != f) return kNoInstance;
return b;
}
// Left partner of right instance b under rule (left, right, off), or -1.
InstanceId left_of(InstanceId b, int32_t left, int32_t off) const {
const int64_t f = root[b] - off;
if (f < 0) return kNoInstance;
const InstanceId a = cell(f, tb->rdim[left]);
if (a == kNoInstance || a == b || !alive[a] || tok[a] != left || root[a] != f) return kNoInstance;
return a;
}
std::vector<int32_t> serialize() const {
std::vector<InstanceId> live;
for (size_t i = 0; i < tok.size(); ++i)
if (alive[i]) live.push_back(static_cast<InstanceId>(i));
std::sort(live.begin(), live.end(), [&](InstanceId x, InstanceId y) { return key(x) < key(y); });
std::vector<int32_t> out;
out.reserve(live.size());
for (InstanceId i : live) out.push_back(tok[i]);
return out;
}
};
struct Stat {
bool retired = false; // a rule is learned at most once; leftovers stay unmerged
int64_t raw = 0;
uint32_t version = 0;
int64_t exact = -1; // matching- and validity-certified count at exact_version
uint32_t exact_version = 0;
int64_t added = 0; // cumulative pair insertions; bounds growth since certification
int64_t added_at_exact = 0;
uint32_t dirty_epoch = 0; // merge batch that last queued this key for re-push
bool live = false; // a heap entry (queued_version, queued_count) represents the key
uint32_t queued_version = 0;
int64_t queued_count = 0;
};
// Open-addressing rule table; references are invalidated only by inserting a new key.
class StatMap {
public:
Stat* find(uint64_t key) {
if (keys_.empty()) return nullptr;
for (size_t h = slot(key);; h = (h + 1) & (keys_.size() - 1)) {
if (keys_[h] == key) return &values_[h];
if (keys_[h] == kEmpty) return nullptr;
}
}
Stat& operator[](uint64_t key) {
if ((used_ + 1) * 2 > keys_.size()) grow();
size_t h = slot(key);
while (keys_[h] != kEmpty && keys_[h] != key) h = (h + 1) & (keys_.size() - 1);
if (keys_[h] == kEmpty) {
keys_[h] = key;
++used_;
}
return values_[h];
}
size_t size() const { return used_; }
template <class F>
void each(F&& visit) const {
for (size_t h = 0; h < keys_.size(); ++h)
if (keys_[h] != kEmpty) visit(keys_[h], values_[h]);
}
private:
static constexpr uint64_t kEmpty = ~0ull;
size_t slot(uint64_t key) const {
return static_cast<size_t>(mix(key) >> 20) & (keys_.size() - 1);
}
void grow() {
std::vector<uint64_t> keys = std::move(keys_);
std::vector<Stat> values = std::move(values_);
const size_t capacity = std::max<size_t>(1 << 16, keys.size() * 2);
keys_.assign(capacity, kEmpty);
values_.assign(capacity, Stat{});
used_ = 0;
for (size_t h = 0; h < keys.size(); ++h)
if (keys[h] != kEmpty) (*this)[keys[h]] = values[h];
}
std::vector<uint64_t> keys_;
std::vector<Stat> values_;
size_t used_ = 0;
};
// Rule table split by key hash so disjoint shards can be updated concurrently.
class ShardedStats {
public:
Stat* find(uint64_t key) { return shards_[shard_of(key)].find(key); }
Stat& operator[](uint64_t key) { return shards_[shard_of(key)][key]; }
StatMap& shard(size_t p) { return shards_[p]; }
size_t size() const {
size_t n = 0;
for (const auto& s : shards_) n += s.size();
return n;
}
template <class F>
void each(F&& visit) const {
for (const auto& s : shards_) s.each(visit);
}
private:
StatMap shards_[kShards];
};
struct Entry {
int64_t count;
uint64_t key;
uint32_t version;
bool operator<(const Entry& o) const {
if (count != o.count) return count < o.count;
return key > o.key;
}
};
class Trainer {
public:
Trainer(std::vector<uint32_t> spatial_support, int32_t dims, int32_t lookahead,
int64_t vocab_size, bool allow_zero, int64_t min_count, int threads,
bool dense_census = false, int64_t max_instances = UINT32_MAX)
: vocab_size_(vocab_size), allow_zero_(allow_zero), dense_census_(dense_census),
min_count_(min_count), max_instances_(max_instances) {
if (max_instances_ < 1 || max_instances_ > UINT32_MAX)
throw std::invalid_argument("max_instances");
tables_.init(spatial_support, dims, lookahead);
#ifdef _OPENMP
threads_ = threads > 0 ? threads : std::min(omp_get_max_threads(), 32);
#else
(void)threads;
threads_ = 1;
#endif
}
// anchors: global frame index per instance (chunks already separated by >= lookahead).
void set_instances(py::array_t<int64_t, py::array::c_style> anchors,
py::array_t<int32_t, py::array::c_style> tokens, int64_t frames) {
auto a = anchors.unchecked<1>();
auto t = tokens.unchecked<1>();
if (a.shape(0) != t.shape(0)) throw std::invalid_argument("anchors/tokens length");
if (a.shape(0) > max_instances_)
throw std::length_error("live instance capacity exceeded before training");
st_.reset(&tables_, frames);
// Pre-sized so references into the outer vector stay valid while children are added.
occurrence_.assign(static_cast<size_t>(vocab_size_) + 16, {});
dead_.assign(static_cast<size_t>(vocab_size_) + 16, 0);
for (py::ssize_t i = 0; i < a.shape(0); ++i) {
const int32_t tk = t(i);
if (tk < 0 || tk >= tables_.base) throw std::invalid_argument("token outside spatial vocab");
occurrence_[tk].push_back(st_.add(tk, a(i)));
}
for (auto& occ : occurrence_)
std::sort(occ.begin(), occ.end(), [&](InstanceId x, InstanceId y) { return st_.root[x] < st_.root[y]; });
initial_instances_ = a.shape(0);
// Supports never change, so frames preceded by >= lookahead empty frames split the
// axis into ranges whose merges never read or write each other's cells.
safe_.clear();
int64_t empty = tables_.lookahead;
for (int64_t f = 0; f < frames; ++f) {
bool occupied = false;
for (int32_t d = 0; d < tables_.dims && !occupied; ++d) occupied = st_.cell(f, d) != kNoInstance;
if (!occupied) {
++empty;
continue;
}
if (empty >= tables_.lookahead) safe_.push_back(f);
empty = 0;
}
}
py::dict train(bool progress) {
const auto t0 = std::chrono::steady_clock::now();
{
py::gil_scoped_release release;
census();
std::vector<uint64_t> keys;
keys.reserve(stats_.size());
stats_.each([&](uint64_t k, const Stat&) { keys.push_back(k); });
for (uint64_t k : keys) queue(k, stats_[k], stats_[k].raw);
initial_types_ = stats_.size();
census_seconds_ = seconds_since(t0);
run(progress);
}
py::dict report;
report["initial_instances"] = initial_instances_;
report["final_instances"] = live_;
report["initial_rule_types"] = initial_types_;
report["merges"] = static_cast<int64_t>(merges_.size());
report["dry_runs"] = dry_runs_;
report["rejected_by_validity"] = rejected_;
report["applied"] = applied_;
report["census_seconds"] = census_seconds_;
report["dry_run_seconds"] = dry_seconds_;
report["apply_seconds"] = apply_seconds_;
report["collect_seconds"] = collect_seconds_;
report["parallel_seconds"] = parallel_seconds_;
report["serial_seconds"] = serial_seconds_;
report["merge_seconds"] = merge_seconds_;
report["changes"] = changes_;
report["compactions"] = compactions_;
report["total_seconds"] = seconds_since(t0);
return report;
}
std::vector<std::tuple<int32_t, int32_t, int32_t>> merges() const { return merges_; }
// Count every forward pair once: blockwise parallel key generation, parallel sort, RLE.
void census() {
const int64_t n = static_cast<int64_t>(st_.tok.size());
const int64_t base = tables_.base;
const int64_t universe = base * base * tables_.lookahead;
if (dense_census_ && universe <= 4000000) {
std::vector<std::vector<int64_t>> local(
threads_, std::vector<int64_t>(static_cast<size_t>(universe), 0));
#ifdef _OPENMP
#pragma omp parallel num_threads(threads_)
#endif
{
#ifdef _OPENMP
auto& counts = local[omp_get_thread_num()];
#else
auto& counts = local[0];
#endif
#ifdef _OPENMP
#pragma omp for schedule(static)
#endif
for (int64_t i = 0; i < n; ++i) {
const InstanceId ii = static_cast<InstanceId>(i);
st_.partners(ii, allow_zero_, [&](InstanceId a, InstanceId b, int32_t off) {
if (a == ii) {
const int64_t key = (static_cast<int64_t>(st_.tok[a]) * base + st_.tok[b]) *
tables_.lookahead + off;
++counts[static_cast<size_t>(key)];
}
});
}
}
for (int64_t key = 0; key < universe; ++key) {
int64_t count = 0;
for (const auto& counts : local) count += counts[static_cast<size_t>(key)];
if (!count) continue;
const int32_t offset = static_cast<int32_t>(key % tables_.lookahead);
const int64_t pair = key / tables_.lookahead;
const int32_t left = static_cast<int32_t>(pair / base);
const int32_t right = static_cast<int32_t>(pair % base);
stats_[rule_key(left, right, offset)].raw += count;
}
return;
}
const int64_t block = int64_t{1} << 22;
std::vector<std::vector<uint64_t>> local(threads_);
std::vector<uint64_t> keys;
for (int64_t start = 0; start < n; start += block) {
const int64_t stop = std::min(n, start + block);
#ifdef _OPENMP
#pragma omp parallel num_threads(threads_)
#endif
{
#ifdef _OPENMP
auto& out = local[omp_get_thread_num()];
#else
auto& out = local[0];
#endif
out.clear();
#ifdef _OPENMP
#pragma omp for schedule(static)
#endif
for (int64_t i = start; i < stop; ++i) {
const InstanceId ii = static_cast<InstanceId>(i);
st_.partners(ii, allow_zero_, [&](InstanceId a, InstanceId b, int32_t off) {
if (a == ii) out.push_back(rule_key(st_.tok[a], st_.tok[b], off));
});
}
}
keys.clear();
for (auto& v : local) keys.insert(keys.end(), v.begin(), v.end());
AC2_PARALLEL_SORT(keys.begin(), keys.end());
for (size_t k = 0; k < keys.size();) {
size_t e = k;
while (e < keys.size() && keys[e] == keys[k]) ++e;
stats_[keys[k]].raw += static_cast<int64_t>(e - k);
k = e;
}
}
}
private:
// Reclaim dead slots without changing token order, coverage, or candidate statistics.
void compact_instances() {
std::vector<InstanceId> remap(st_.tok.size(), kNoInstance);
std::vector<int32_t> tokens;
std::vector<int64_t> roots;
std::vector<uint8_t> alive;
tokens.reserve(static_cast<size_t>(live_));
roots.reserve(static_cast<size_t>(live_));
alive.reserve(static_cast<size_t>(live_));
for (size_t old = 0; old < st_.tok.size(); ++old) {
if (!st_.alive[old]) continue;
remap[old] = static_cast<InstanceId>(tokens.size());
tokens.push_back(st_.tok[old]);
roots.push_back(st_.root[old]);
alive.push_back(1);
}
if (tokens.size() != static_cast<size_t>(live_))
throw std::runtime_error("instance compaction live count mismatch");
for (auto& occ : occurrence_) {
size_t next = 0;
for (InstanceId old : occ) {
const InstanceId mapped = remap[static_cast<size_t>(old)];
if (mapped != kNoInstance) occ[next++] = mapped;
}
occ.resize(next);
}
for (auto& slot : st_.member)
if (slot != kNoInstance) slot = remap[static_cast<size_t>(slot)];
st_.tok.swap(tokens);
st_.root.swap(roots);
st_.alive.swap(alive);
std::fill(dead_.begin(), dead_.end(), 0);
++compactions_;
}
void change(uint64_t k, int64_t delta, StatMap& table, std::vector<uint64_t>& dirty) {
Stat& s = table[k];
s.raw += delta;
if (delta > 0) s.added += delta;
++s.version;
if (s.dirty_epoch != epoch_) {
s.dirty_epoch = epoch_;
dirty.push_back(k);
}
}
std::vector<InstanceId>& lefts(int32_t left) {
auto& occ = occurrence_[left];
// Lists are root-sorted by construction: atoms are sorted once, and every child list
// is appended in the root order of its single merge batch.
if (dead_[left] * 4 > static_cast<int64_t>(occ.size())) {
occ.erase(std::remove_if(occ.begin(), occ.end(), [&](InstanceId i) { return !st_.alive[i]; }),
occ.end());
dead_[left] = 0;
}
return occ;
}
static double seconds_since(std::chrono::steady_clock::time_point t) {
return std::chrono::duration<double>(std::chrono::steady_clock::now() - t).count();
}
// Candidate (a, b) pairs are visited in increasing root order of a by scanning the
// shorter occurrence list; the partner is one cell lookup away in either direction.
const std::vector<InstanceId>& scanned(int32_t left, int32_t right, bool& forward) const {
const auto& lo = occurrence_[left];
const auto& ro = occurrence_[right];
forward = left == right || lo.size() <= ro.size();
return forward ? lo : ro;
}
template <class F>
void pairs_in(int32_t left, int32_t right, int32_t off, size_t begin, size_t end,
F&& visit) const {
bool forward;
const auto& scan = scanned(left, right, forward);
for (size_t i = begin; i < end; ++i) {
const InstanceId x = scan[i];
if (!st_.alive[x]) continue;
if (forward) {
const InstanceId b = st_.right_of(x, right, off);
if (b != kNoInstance) visit(x, b);
} else {
const InstanceId a = st_.left_of(x, left, off);
if (a != kNoInstance) visit(a, x);
}
}
}
// Cuts of a root-sorted list into blocks of about ``grain`` entries at safe frames, so
// per-block replay (self-rule chains included) equals one sequential pass.
std::vector<size_t> blocks(const std::vector<InstanceId>& scan, size_t grain) const {
std::vector<size_t> cuts{0};
for (size_t next = grain; next < scan.size();) {
auto frame = std::upper_bound(safe_.begin(), safe_.end(), st_.root[scan[next]]);
if (frame == safe_.end()) break;
const size_t at = static_cast<size_t>(
std::lower_bound(scan.begin() + next, scan.end(), *frame,
[&](InstanceId i, int64_t f) { return st_.root[i] < f; }) -
scan.begin());
if (at >= scan.size()) break;
cuts.push_back(at);
next = at + grain;
}
cuts.push_back(scan.size());
return cuts;
}
// Candidate pairs in scan order, gathered by blocks in parallel.
std::vector<std::pair<InstanceId, InstanceId>> collect(int32_t left, int32_t right, int32_t off) const {
bool forward;
const std::vector<size_t> cuts = blocks(scanned(left, right, forward), 1 << 16);
std::vector<std::vector<std::pair<InstanceId, InstanceId>>> parts(cuts.size() - 1);
#ifdef _OPENMP
#pragma omp parallel for schedule(dynamic, 1) num_threads(threads_) if (parts.size() > 1)
#endif
for (size_t k = 0; k < parts.size(); ++k)
pairs_in(left, right, off, cuts[k], cuts[k + 1],
[&](InstanceId a, InstanceId b) { parts[k].emplace_back(a, b); });
std::vector<std::pair<InstanceId, InstanceId>> pairs;
for (auto& part : parts) pairs.insert(pairs.end(), part.begin(), part.end());
return pairs;
}
// Exact count of one block under the real merge's matching and validity tests.
int64_t count_exact(int32_t left, int32_t right, int32_t off, size_t begin, size_t end) const {
int64_t n = 0;
// Self rules chain: a may be the right half of a recent match (roots increase).
std::vector<InstanceId> recent;
size_t head = 0;
pairs_in(left, right, off, begin, end, [&](InstanceId a, InstanceId b) {
if (left == right) {
while (head < recent.size() && st_.root[recent[head]] < st_.root[a]) ++head;
if (std::find(recent.begin() + head, recent.end(), a) != recent.end()) return;
}
if (!st_.valid(a, b, off)) return;
if (left == right) recent.push_back(b);
++n;
});
return n;
}
// Certifies a batch of keys; every key is split into blocks so one large list does
// not serialize the batch. Read-only on shared state.
void certify(std::vector<Entry>& batch) {
struct Timer {
double& sink;
std::chrono::steady_clock::time_point t;
~Timer() { sink += seconds_since(t); }
} timer{dry_seconds_, std::chrono::steady_clock::now()};
dry_runs_ += static_cast<int64_t>(batch.size());
struct Task {
size_t key, begin, end;
};
std::vector<Task> tasks;
for (size_t k = 0; k < batch.size(); ++k) {
int32_t l, r, o;
unpack(batch[k].key, l, r, o);
lefts(l); // compaction mutates lists: do it before the parallel region
lefts(r);
bool forward;
const std::vector<size_t> cuts = blocks(scanned(l, r, forward), 1 << 15);
for (size_t c = 0; c + 1 < cuts.size(); ++c) tasks.push_back({k, cuts[c], cuts[c + 1]});
}
std::vector<int64_t> partial(tasks.size());
#ifdef _OPENMP
#pragma omp parallel for schedule(dynamic, 1) num_threads(threads_)
#endif
for (size_t t = 0; t < tasks.size(); ++t) {
int32_t l, r, o;
unpack(batch[tasks[t].key].key, l, r, o);
partial[t] = count_exact(l, r, o, tasks[t].begin, tasks[t].end);
}
std::vector<int64_t> exact(batch.size(), 0);
for (size_t t = 0; t < tasks.size(); ++t) exact[tasks[t].key] += partial[t];
for (size_t k = 0; k < batch.size(); ++k) {
Stat& s = stats_[batch[k].key];
s.exact = exact[k];
s.exact_version = s.version;
s.added_at_exact = s.added;
if (s.exact >= min_count_) queue(batch[k].key, s, s.exact);
}
}
// Pairs are applied in root order; ranges split at safe frames are independent, so
// each is replayed sequentially by one thread with exactly the serial outcome.
void apply(int32_t left, int32_t right, int32_t off, int32_t child) {
struct Timer {
double& sink;
std::chrono::steady_clock::time_point t;
~Timer() { sink += seconds_since(t); }
} timer{apply_seconds_, std::chrono::steady_clock::now()};
lefts(left);
lefts(right);
// Candidates are collected first: applying merges kills instances in both lists.
const auto collect_start = std::chrono::steady_clock::now();
std::vector<std::pair<InstanceId, InstanceId>> pairs = collect(left, right, off);
if (st_.tok.size() + pairs.size() > static_cast<size_t>(max_instances_)) {
pairs.clear();
pairs.shrink_to_fit();
compact_instances();
pairs = collect(left, right, off);
if (st_.tok.size() + pairs.size() > static_cast<size_t>(max_instances_))
throw std::length_error("live instance capacity exceeded after compaction");
}
collect_seconds_ += seconds_since(collect_start);
const std::vector<size_t> cuts = split(pairs);
const size_t groups = cuts.size() - 1;
const InstanceId base = st_.reserve(pairs.size());
if (pool_.size() < groups) pool_.resize(groups, std::vector<DeltaMap>(kShards));
dirty_shards_.resize(kShards);
std::vector<std::vector<InstanceId>> born(groups);
std::vector<int64_t> applied(groups, 0), rejected(groups, 0);
const auto loop_start = std::chrono::steady_clock::now();
#ifdef _OPENMP
#pragma omp parallel for schedule(dynamic, 1) num_threads(std::min<size_t>(threads_, groups))
#endif
for (size_t g = 0; g < groups; ++g) {
auto& delta = pool_[g];
auto add = [&](uint64_t key, int64_t d) { delta[shard_of(key)].add(key, d); };
for (size_t k = cuts[g]; k < cuts[g + 1]; ++k) {
const InstanceId a = pairs[k].first, b = pairs[k].second;
if (!st_.alive[a] || !st_.alive[b]) continue;
if (!st_.valid(a, b, off)) {
++rejected[g];
continue;
}
st_.partners(a, allow_zero_, [&](InstanceId x, InstanceId y, int32_t o) {
add(rule_key(st_.tok[x], st_.tok[y], o), -1);
});
st_.partners(b, allow_zero_, [&](InstanceId x, InstanceId y, int32_t o) {
if (x != a && y != a) add(rule_key(st_.tok[x], st_.tok[y], o), -1);
});
const int64_t r = st_.root[a];
st_.kill(a);
st_.kill(b);
const InstanceId m = base + static_cast<InstanceId>(k);
st_.tok[m] = child;
st_.root[m] = r;
st_.place(m);
born[g].push_back(m);
st_.partners(m, allow_zero_, [&](InstanceId x, InstanceId y, int32_t o) {
add(rule_key(st_.tok[x], st_.tok[y], o), +1);
});
++applied[g];
}
}
(groups > 1 ? parallel_seconds_ : serial_seconds_) += seconds_since(loop_start);
const auto merge_start = std::chrono::steady_clock::now();
std::vector<int64_t> changes(kShards, 0);
#ifdef _OPENMP
#pragma omp parallel for schedule(dynamic, 1) num_threads(threads_) if (groups > 1)
#endif
for (size_t p = 0; p < kShards; ++p) {
StatMap& table = stats_.shard(p);
for (size_t g = 0; g < groups; ++g) {
pool_[g][p].each([&](uint64_t k, int64_t d) {
change(k, d, table, dirty_shards_[p]);
++changes[p];
});
pool_[g][p].clear();
}
}
for (size_t p = 0; p < kShards; ++p) {
dirty_.insert(dirty_.end(), dirty_shards_[p].begin(), dirty_shards_[p].end());
dirty_shards_[p].clear();
changes_ += changes[p];
}
for (size_t g = 0; g < groups; ++g) {
auto& occ = occurrence_[child];
occ.insert(occ.end(), born[g].begin(), born[g].end());
dead_[left] += applied[g];
dead_[right] += applied[g];
applied_ += applied[g];
rejected_ += rejected[g];
live_ -= applied[g];
}
merge_seconds_ += seconds_since(merge_start);
}
// Cut root-ordered pairs into contiguous groups at safe frames.
std::vector<size_t> split(const std::vector<std::pair<InstanceId, InstanceId>>& pairs) const {
std::vector<size_t> cuts{0};
const size_t n = pairs.size();
const size_t groups = std::max<size_t>(1, std::min<size_t>(threads_, n / kApplyGrain));
for (size_t g = 1; g < groups; ++g) {
const size_t target = std::max(cuts.back(), n * g / groups);
if (target >= n) break;
auto frame = std::upper_bound(safe_.begin(), safe_.end(), st_.root[pairs[target].first]);
if (frame == safe_.end()) break;
const size_t cut = static_cast<size_t>(
std::lower_bound(pairs.begin() + target, pairs.end(), *frame,
[&](const std::pair<InstanceId, InstanceId>& p, int64_t f) {
return st_.root[p.first] < f;
}) -
pairs.begin());
if (cut > cuts.back() && cut < n) cuts.push_back(cut);
}
cuts.push_back(n);
return cuts;
}
static bool certified(const Stat& s) { return s.exact >= 0 && s.exact_version == s.version; }
static int64_t bound(const Stat& s) {
if (certified(s)) return s.exact;
return s.exact >= 0 ? std::min(s.raw, s.exact + s.added - s.added_at_exact) : s.raw;
}
void queue(uint64_t key, Stat& s, int64_t count) {
heap_.push({count, key, s.version});
s.live = true;
s.queued_version = s.version;
s.queued_count = count;
}
bool take(Entry& e, Stat*& s) {
e = heap_.top();
heap_.pop();
s = stats_.find(e.key);
if (!s || s->retired || !s->live || e.version != s->queued_version) return false;
s->live = false;
if (e.version == s->version) return true;
const int64_t b = bound(*s);
if (b > 0) queue(e.key, *s, b);
return false;
}
bool current_certified(const Entry& e) {
const Stat* s = stats_.find(e.key);
return s && !s->retired && s->live && e.version == s->queued_version &&
e.version == s->version && certified(*s);
}
void run(bool progress) {
live_ = initial_instances_;
while (tables_.size() < vocab_size_ && !heap_.empty()) {
Entry top;
Stat* found;
if (!take(top, found)) continue;
Stat& s = *found;
if (!certified(s)) {
std::vector<Entry> batch{top};
while (!heap_.empty() && batch.size() < kCertifyBatch) {
if (current_certified(heap_.top())) break;
Entry e;
Stat* next;
if (take(e, next)) batch.push_back(e);
}
certify(batch);
continue;
}
if (top.count != s.exact) continue;
if (s.exact < min_count_) continue;
int32_t l, r, o;
unpack(top.key, l, r, o);
std::vector<uint32_t> m;
if (!tables_.compose(l, r, o, m)) {
s.exact = 0;
continue;
}
const int32_t child = tables_.add_merge(l, r, o);
merges_.emplace_back(l, r, o);
s.retired = true;
dirty_.clear();
++epoch_;
apply(l, r, o, child);
// Lazy decrease-key: only a grown bound needs a new entry now.
for (uint64_t k : dirty_) {
Stat& d = stats_[k];
if (d.retired) continue;
const int64_t b = bound(d);
if (b > 0 && (!d.live || b > d.queued_count)) queue(k, d, b);
}
dirty_.clear();
if (progress && merges_.size() % 256 == 0) {
py::gil_scoped_acquire acquire;
py::print("sparse temporal merges", merges_.size(), "live", live_, py::arg("flush") = true);
}
}
}
Tables tables_;
Instances st_;
int64_t vocab_size_;
bool allow_zero_;
bool dense_census_;
int64_t min_count_;
int64_t max_instances_;
std::vector<std::vector<InstanceId>> occurrence_;
ShardedStats stats_;
std::vector<std::vector<DeltaMap>> pool_; // [apply group][shard]
std::vector<std::vector<uint64_t>> dirty_shards_; // [shard]
std::priority_queue<Entry> heap_;
std::vector<uint64_t> dirty_;
std::vector<std::tuple<int32_t, int32_t, int32_t>> merges_;
std::vector<int64_t> dead_;
std::vector<int64_t> safe_;
uint32_t epoch_ = 0;
int threads_ = 1;
int64_t initial_instances_ = 0, live_ = 0, dry_runs_ = 0, rejected_ = 0, applied_ = 0;
int64_t compactions_ = 0;
size_t initial_types_ = 0;
int64_t changes_ = 0;
double census_seconds_ = 0, dry_seconds_ = 0, apply_seconds_ = 0, collect_seconds_ = 0,
parallel_seconds_ = 0, serial_seconds_ = 0, merge_seconds_ = 0;
};
struct Spatial {
std::vector<int64_t> offsets;
std::vector<uint32_t> support;
std::vector<std::vector<int32_t>> expansion; // token -> value per sparse slot (-1 absent)
std::unordered_map<uint64_t, std::pair<int32_t, int32_t>> rules; // (lo, hi) -> (child, rank)
int32_t dims = 0;
static uint64_t pair_key(int32_t a, int32_t b) {
if (a > b) std::swap(a, b);
return (static_cast<uint64_t>(a) << 32) | static_cast<uint32_t>(b);
}
void init(const std::vector<int64_t>& bins, const std::vector<std::pair<int32_t, int32_t>>& merges) {
dims = static_cast<int32_t>(bins.size());
offsets.assign(1, 0);
for (int64_t b : bins) {
if (b < 1) throw std::invalid_argument("every sparse slot needs at least one value");
offsets.push_back(offsets.back() + b);
}
for (int32_t d = 0; d < dims; ++d)
for (int64_t v = 0; v < bins[d]; ++v) {
support.push_back(1u << d);
std::vector<int32_t> e(dims, -1);
e[d] = static_cast<int32_t>(v);
expansion.push_back(std::move(e));
}
const int32_t atoms = static_cast<int32_t>(support.size());
for (size_t r = 0; r < merges.size(); ++r) {
const auto [a, b] = merges[r];
const int32_t child = atoms + static_cast<int32_t>(r);
if (a < 0 || b < 0 || a >= child || b >= child || (support[a] & support[b]))
throw std::invalid_argument("invalid spatial merge");
support.push_back(support[a] | support[b]);
std::vector<int32_t> e(dims);
for (int32_t d = 0; d < dims; ++d) e[d] = std::max(expansion[a][d], expansion[b][d]);
expansion.push_back(std::move(e));
rules.emplace(pair_key(a, b), std::make_pair(child, static_cast<int32_t>(r)));
}
}
std::vector<int32_t> encode(const int32_t* row) const {
std::set<int32_t> tokens;
for (int32_t d = 0; d < dims; ++d)
if (row[d] >= 0) tokens.insert(static_cast<int32_t>(offsets[d] + row[d]));
using Item = std::tuple<int32_t, int32_t, int32_t, int32_t>;
std::priority_queue<Item, std::vector<Item>, std::greater<Item>> heap;
auto offer = [&](int32_t a, int32_t b) {
auto it = rules.find(pair_key(a, b));
if (it != rules.end())
heap.emplace(it->second.second, it->second.first, std::min(a, b), std::max(a, b));
};
for (auto i = tokens.begin(); i != tokens.end(); ++i)
for (auto j = std::next(i); j != tokens.end(); ++j) offer(*i, *j);
while (!heap.empty()) {
const auto [rank, child, left, right] = heap.top();
heap.pop();
if (!tokens.count(left) || !tokens.count(right)) continue;
tokens.erase(left);
tokens.erase(right);
for (int32_t other : tokens) offer(child, other);
tokens.insert(child);
}
return {tokens.begin(), tokens.end()};
}
};
// Runtime codec over dense primitive cells ``(T, D)`` with ``-1`` for inactive slots.
class SparseCodec {
public:
SparseCodec(std::vector<int64_t> bins, std::vector<int64_t> zero,
std::vector<std::pair<int32_t, int32_t>> spatial_merges, int32_t lookahead,
std::vector<std::tuple<int32_t, int32_t, int32_t>> temporal_merges)
: bins_(std::move(bins)), zero_(std::move(zero)) {
if (bins_.empty() || bins_.size() != zero_.size())
throw std::invalid_argument("bins and zero must have one entry per slot");
dims_ = static_cast<int32_t>(bins_.size());
std::vector<int64_t> sparse;
for (int32_t d = 0; d < dims_; ++d) {
// -2 retains every dense value; -1 change-codes an absolute slot.
if (zero_[d] >= bins_[d] || zero_[d] < -2) throw std::invalid_argument("zero index");
sparse.push_back(bins_[d] - (zero_[d] >= 0));
}
sparse.push_back(1); // E
spatial_.init(sparse, spatial_merges);
tables_.init(spatial_.support, dims_ + 1, lookahead);
for (size_t r = 0; r < temporal_merges.size(); ++r) {
auto [a, b, o] = temporal_merges[r];
if (a < 0 || b < 0 || a >= tables_.size() || b >= tables_.size())
throw std::invalid_argument("merge order");
tables_.add_merge(a, b, o);
if (!rank_.emplace(rule_key(a, b, o), static_cast<int32_t>(r)).second)
throw std::invalid_argument("duplicate temporal rule");
}
}
int32_t vocab_size() const { return tables_.size(); }
int32_t spatial_vocab_size() const { return static_cast<int32_t>(spatial_.support.size()); }
std::vector<std::vector<int32_t>> encode_batch(
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& cells,
int threads) const {
if (cells.ndim() != 3 || cells.shape(2) != dims_ || cells.shape(0) < 1 || cells.shape(1) < 1)
throw std::invalid_argument("cells must have shape (B, T, D) with B, T >= 1");
const int64_t batch = cells.shape(0), steps = cells.shape(1);
const int32_t* data = cells.data();
std::vector<std::vector<int32_t>> out(batch);
{
py::gil_scoped_release release;
ac2::parallel_for(batch, threads, [&](int64_t b) {
out[b] = encode_one(data + b * steps * dims_, steps);
});
}
return out;
}
py::array_t<int32_t> decode_batch(const std::vector<std::vector<int32_t>>& rows,
int64_t horizon, const std::vector<bool>& active) const {
if (horizon < 1) throw std::invalid_argument("horizon must be positive");
if (static_cast<int32_t>(active.size()) != dims_)
throw std::invalid_argument("active must have one flag per slot");
py::array_t<int32_t> out({static_cast<int64_t>(rows.size()), horizon, int64_t{dims_}});
int32_t* data = out.mutable_data();
for (size_t k = 0; k < rows.size(); ++k) decode_one(rows[k], horizon, active, data + k * horizon * dims_);
return out;
}
// Spatial tokens of sparse frames ``(N, D + 1)``: per-frame counts and flat tokens.
py::tuple spatial_encode(
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& frames) const {
if (frames.ndim() != 2 || frames.shape(1) != dims_ + 1)
throw std::invalid_argument("frames must have shape (N, D + 1)");
const int64_t n = frames.shape(0);
py::array_t<int64_t> counts(n);
std::vector<int32_t> flat;
for (int64_t i = 0; i < n; ++i) {
const int32_t* row = frames.data() + i * (dims_ + 1);
for (int32_t d = 0; d <= dims_; ++d)
if (row[d] >= static_cast<int32_t>(spatial_.offsets[d + 1] - spatial_.offsets[d]))
throw std::invalid_argument("sparse value out of range");
const auto tokens = spatial_.encode(row);
counts.mutable_data()[i] = static_cast<int64_t>(tokens.size());
flat.insert(flat.end(), tokens.begin(), tokens.end());
}
return py::make_tuple(counts, py::array_t<int32_t>(flat.size(), flat.data()));
}
std::vector<std::vector<std::pair<int32_t, int32_t>>> patterns() const { return tables_.pattern; }
private:
// Dense chunk -> sparse frames (T, D + 1); validates slot activity and ranges.
std::vector<int32_t> sparsify(const int32_t* cells, int64_t steps) const {
const int32_t width = dims_ + 1;
std::vector<int32_t> sparse(static_cast<size_t>(steps) * width, -1);
int64_t last = -1;
for (int32_t d = 0; d < dims_; ++d) {
const bool active = cells[d] >= 0;
for (int64_t t = 0; t < steps; ++t) {
const int32_t v = cells[t * dims_ + d];
if ((v >= 0) != active)
throw std::invalid_argument("a slot must be active for the whole chunk or not at all");
if (v >= bins_[d]) throw std::invalid_argument("cell outside the primitive range");
if (!active) continue;
int32_t s = -1;
if (zero_[d] == -2) {
s = v;
} else if (zero_[d] < 0) {
if (t == 0 || v != cells[(t - 1) * dims_ + d]) s = v;
} else if (v != zero_[d]) {
s = v - (v > zero_[d]);
}
sparse[t * width + d] = s;
if (s >= 0) last = std::max(last, t);
}
}
for (int64_t t = 0; t < last; ++t) sparse[t * width + dims_] = 0;
return sparse;
}
std::vector<int32_t> encode_one(const int32_t* cells, int64_t steps) const {
const std::vector<int32_t> sparse = sparsify(cells, steps);
Instances s;
s.reset(&tables_, steps + tables_.lookahead);
for (int64_t t = 0; t < steps; ++t)
for (int32_t token : spatial_.encode(sparse.data() + t * (dims_ + 1))) s.add(token, t);
using Cand = std::tuple<int32_t, int64_t, int32_t, int64_t, InstanceId, InstanceId, int32_t>;
std::priority_queue<Cand, std::vector<Cand>, std::greater<Cand>> heap;
auto offer = [&](InstanceId i) {
s.partners(i, true, [&](InstanceId a, InstanceId b, int32_t off) {
auto it = rank_.find(rule_key(s.tok[a], s.tok[b], off));
if (it == rank_.end()) return;
heap.emplace(it->second, s.root[a], tables_.rdim[s.tok[a]], s.root[b], a, b, off);
});
};
for (size_t i = 0; i < s.tok.size(); ++i) offer(static_cast<InstanceId>(i));
while (!heap.empty()) {
auto [r, ra, da, rb, a, b, off] = heap.top();
heap.pop();
if (!s.alive[a] || !s.alive[b] || !s.valid(a, b, off)) continue;
const int64_t root = s.root[a];
s.kill(a);
s.kill(b);
offer(s.add(tables_.base + r, root));
}
return s.serialize();
}
void decode_one(const std::vector<int32_t>& row, int64_t horizon,
const std::vector<bool>& active, int32_t* cells) const {
const int32_t L = tables_.lookahead, width = dims_ + 1;
uint32_t allowed = 1u << dims_;
for (int32_t d = 0; d < dims_; ++d)
if (active[d]) allowed |= 1u << d;
std::vector<uint32_t> cover(horizon + L, 0);
std::vector<int32_t> sparse(static_cast<size_t>(horizon) * width, -1);
int64_t front = 0;
std::vector<uint32_t> own(L, 0);
std::vector<int32_t> frames;
for (int32_t token : row) {
if (token < 0 || token >= tables_.size()) throw std::invalid_argument("token outside vocab");
if (front >= horizon) throw std::invalid_argument("token after the last open frame");
frames.clear();
for (auto [o, t] : tables_.pattern[token]) {
const int64_t u = front + o;
const uint32_t m = spatial_.support[t];
if (u >= horizon) throw std::invalid_argument("token exceeds horizon");
if (m & ~allowed) throw std::invalid_argument("token touches an inactive slot");
if (m & own[o]) throw std::invalid_argument("overlapping token supports");
if ((m >> dims_ & 1u) && u == horizon - 1)
throw std::invalid_argument("end-of-frame atom on the last frame");
if (!own[o]) frames.push_back(o);
own[o] |= m;
for (int32_t d = 0; d < width; ++d)
if (spatial_.expansion[t][d] >= 0) sparse[u * width + d] = spatial_.expansion[t][d];
}
for (int32_t o : frames) {
const int64_t u = front + o;
if (cover[u] & own[o]) throw std::invalid_argument("overlapping token supports");
if (cover[u] >> dims_ & 1u) throw std::invalid_argument("token touches a closed frame");
cover[u] |= own[o];
own[o] = 0;
}
while (front < horizon && (cover[front] >> dims_ & 1u)) ++front;
}
for (int32_t d = 0; d < dims_; ++d) {
int32_t carried = -1;
for (int64_t t = 0; t < horizon; ++t) {
const int32_t v = sparse[t * width + d];
int32_t& out = cells[t * dims_ + d];
if (!active[d]) {
out = -1;
} else if (zero_[d] == -2) {
if (v < 0) throw std::invalid_argument("dense active slot is missing");
out = v;
} else if (zero_[d] < 0) {
if (v >= 0) carried = v;
if (carried < 0) throw std::invalid_argument("an absolute slot must be stated on frame 0");
out = carried;
} else {
out = v < 0 ? static_cast<int32_t>(zero_[d]) : v + (v >= zero_[d]);
}
}
}
}
std::vector<int64_t> bins_, zero_;
int32_t dims_ = 0;
Spatial spatial_;
Tables tables_;
std::unordered_map<uint64_t, int32_t> rank_;
};
} // namespace
void register_sparse(py::module_& m) {
py::class_<Trainer>(m, "SparseTrainer")
.def(py::init<std::vector<uint32_t>, int32_t, int32_t, int64_t, bool, int64_t, int, bool, int64_t>(),
py::arg("spatial_support"), py::arg("dims"), py::arg("lookahead"),
py::arg("vocab_size"), py::arg("allow_zero_offset"), py::arg("min_count") = 2,
py::arg("threads") = 0, py::arg("dense_census") = false,
py::arg("max_instances") = UINT32_MAX)
.def("set_instances", &Trainer::set_instances)
.def("train", &Trainer::train, py::arg("progress") = false)
.def("merges", &Trainer::merges);
py::class_<SparseCodec>(m, "SparseCodec")
.def(py::init<std::vector<int64_t>, std::vector<int64_t>,
std::vector<std::pair<int32_t, int32_t>>, int32_t,
std::vector<std::tuple<int32_t, int32_t, int32_t>>>(),
py::arg("bins"), py::arg("zero"), py::arg("spatial_merges"), py::arg("lookahead"),
py::arg("temporal_merges"))
.def_property_readonly("vocab_size", &SparseCodec::vocab_size)
.def_property_readonly("spatial_vocab_size", &SparseCodec::spatial_vocab_size)
.def("encode_batch", &SparseCodec::encode_batch, py::arg("cells"), py::arg("threads") = 0)
.def("decode_batch", &SparseCodec::decode_batch, py::arg("rows"), py::arg("horizon"),
py::arg("active"))
.def("spatial_encode", &SparseCodec::spatial_encode, py::arg("frames"))
.def("patterns", &SparseCodec::patterns);
}
namespace py = pybind11;
namespace {
constexpr int kMaxLookahead = 32;
constexpr size_t kCertifyBatch = 32;
constexpr size_t kApplyGrain = 256; // minimum candidate pairs per parallel apply group
constexpr size_t kShards = 64; // rule-table shards updated concurrently after a merge
inline uint64_t mix(uint64_t key) { return key * 0x9E3779B97F4A7C15ull; }
inline size_t shard_of(uint64_t key) { return static_cast<size_t>(mix(key) >> 58); }
// Small open-addressing accumulator for rule-count deltas; clearing costs O(used).
class DeltaMap {
public:
void add(uint64_t key, int64_t delta) {
if ((filled_.size() + 1) * 2 > keys_.size()) grow();
size_t h = slot(key);
while (keys_[h] != kEmpty && keys_[h] != key) h = (h + 1) & (keys_.size() - 1);
if (keys_[h] == kEmpty) {
keys_[h] = key;
filled_.push_back(static_cast<uint32_t>(h));
}
values_[h] += delta;
}
template <class F>
void each(F&& visit) const {
for (uint32_t h : filled_) visit(keys_[h], values_[h]);
}
void clear() {
for (uint32_t h : filled_) {
keys_[h] = kEmpty;
values_[h] = 0;
}
filled_.clear();
}
private:
static constexpr uint64_t kEmpty = ~0ull;
size_t slot(uint64_t key) const {
return static_cast<size_t>(mix(key) >> 20) & (keys_.size() - 1);
}
void grow() {
std::vector<uint64_t> keys = std::move(keys_);
std::vector<int64_t> values = std::move(values_);
std::vector<uint32_t> filled = std::move(filled_);
const size_t capacity = std::max<size_t>(64, keys.size() * 2);
keys_.assign(capacity, kEmpty);
values_.assign(capacity, 0);
filled_.clear();
for (uint32_t h : filled) add(keys[h], values[h]);
}
std::vector<uint64_t> keys_;
std::vector<int64_t> values_;
std::vector<uint32_t> filled_;
};
inline uint64_t rule_key(int32_t left, int32_t right, int32_t offset) {
return (static_cast<uint64_t>(left) << 40) | (static_cast<uint64_t>(right) << 16) |
static_cast<uint64_t>(offset);
}
inline void unpack(uint64_t key, int32_t& left, int32_t& right, int32_t& offset) {
left = static_cast<int32_t>(key >> 40);
right = static_cast<int32_t>((key >> 16) & 0xffffffull);
offset = static_cast<int32_t>(key & 0xffffull);
}
// Immutable-per-token geometry shared by the trainer, encoder and decoder.