File size: 6,008 Bytes
1bfc3f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | // Parallel replay of a learned Spatial Set-BPE vocabulary over action cells.
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <algorithm>
#include <cstdint>
#include <functional>
#include <queue>
#include <tuple>
#include <vector>
#ifdef _OPENMP
#include <omp.h>
#endif
namespace py = pybind11;
namespace {
struct Replay {
int32_t dimension;
int32_t vocab_size;
std::vector<int32_t> offsets;
std::vector<int32_t> rule;
std::vector<int32_t> anchor;
int32_t child(int32_t left, int32_t right) const {
if (left > right) std::swap(left, right);
return rule[static_cast<size_t>(left) * vocab_size + right];
}
std::vector<int32_t> encode(const int32_t* code) const {
std::vector<int32_t> active;
active.reserve(dimension);
for (int32_t axis = 0; axis < dimension; ++axis) {
if (code[axis] < 0) continue;
active.push_back(offsets[axis] + code[axis]);
}
using Candidate = std::tuple<int32_t, int32_t, int32_t>;
std::priority_queue<Candidate, std::vector<Candidate>, std::greater<Candidate>> heap;
auto offer = [&](int32_t left, int32_t right) {
const int32_t value = child(left, right);
if (value >= 0) heap.emplace(value, std::min(left, right), std::max(left, right));
};
for (int32_t a = 0; a < static_cast<int32_t>(active.size()); ++a) {
for (int32_t b = a + 1; b < static_cast<int32_t>(active.size()); ++b) {
offer(active[a], active[b]);
}
}
while (!heap.empty()) {
const auto [child_token, left, right] = heap.top();
heap.pop();
auto left_it = std::find(active.begin(), active.end(), left);
auto right_it = std::find(active.begin(), active.end(), right);
if (left_it == active.end() || right_it == active.end()) continue;
*left_it = child_token;
active.erase(right_it);
for (const int32_t other : active) {
if (other != child_token) offer(child_token, other);
}
}
std::sort(active.begin(), active.end(), [&](int32_t left, int32_t right) {
if (anchor[left] != anchor[right]) return anchor[left] < anchor[right];
return left < right;
});
return active;
}
};
} // namespace
py::tuple spatial_encode_instances(
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& codes,
const py::array_t<int64_t, py::array::c_style | py::array::forcecast>& row_anchors,
const std::vector<int32_t>& bins,
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& merges,
int threads) {
if (codes.ndim() != 2 || row_anchors.ndim() != 1 ||
row_anchors.shape(0) != codes.shape(0) ||
codes.shape(1) != static_cast<int64_t>(bins.size()) || merges.ndim() != 2 ||
merges.shape(1) != 2) {
throw std::invalid_argument("invalid spatial replay inputs");
}
Replay replay;
replay.dimension = static_cast<int32_t>(bins.size());
replay.offsets.assign(replay.dimension + 1, 0);
for (int32_t axis = 0; axis < replay.dimension; ++axis) {
replay.offsets[axis + 1] = replay.offsets[axis] + bins[axis];
}
const int32_t atoms = replay.offsets.back();
replay.vocab_size = atoms + static_cast<int32_t>(merges.shape(0));
replay.rule.assign(static_cast<size_t>(replay.vocab_size) * replay.vocab_size, -1);
replay.anchor.assign(replay.vocab_size, -1);
for (int32_t axis = 0; axis < replay.dimension; ++axis) {
for (int32_t token = replay.offsets[axis]; token < replay.offsets[axis + 1]; ++token) {
replay.anchor[token] = axis;
}
}
auto rules = merges.unchecked<2>();
for (int32_t index = 0; index < merges.shape(0); ++index) {
const int32_t left = rules(index, 0), right = rules(index, 1), child = atoms + index;
if (left < 0 || right < 0 || left >= child || right >= child) {
throw std::invalid_argument("invalid spatial merge rule");
}
replay.rule[static_cast<size_t>(std::min(left, right)) * replay.vocab_size +
std::max(left, right)] = child;
replay.anchor[child] = std::min(replay.anchor[left], replay.anchor[right]);
}
const int64_t rows = codes.shape(0);
std::vector<int32_t> counts(rows);
const int32_t* source = codes.data();
for (int64_t row = 0; row < rows; ++row) {
int32_t active = 0;
for (int32_t axis = 0; axis < replay.dimension; ++axis) {
const int32_t value = source[row * replay.dimension + axis];
if (value < -1 || value >= bins[axis]) {
throw std::invalid_argument("a cell index is outside its coordinate's bin count");
}
active += value >= 0;
}
if (active == 0) throw std::invalid_argument("every row needs an active coordinate");
}
#ifndef _OPENMP
(void)threads;
#endif
{
py::gil_scoped_release release;
#ifdef _OPENMP
#pragma omp parallel for schedule(static) num_threads(threads > 0 ? threads : omp_get_max_threads())
#endif
for (int64_t row = 0; row < rows; ++row) {
counts[row] = static_cast<int32_t>(
replay.encode(source + row * replay.dimension).size());
}
}
std::vector<int64_t> offsets(rows + 1, 0);
for (int64_t row = 0; row < rows; ++row) offsets[row + 1] = offsets[row] + counts[row];
py::array_t<int32_t> tokens(offsets.back());
py::array_t<int64_t> anchors(offsets.back());
int32_t* output_tokens = tokens.mutable_data();
int64_t* output_anchors = anchors.mutable_data();
const int64_t* input_anchors = row_anchors.data();
{
py::gil_scoped_release release;
#ifdef _OPENMP
#pragma omp parallel for schedule(static) num_threads(threads > 0 ? threads : omp_get_max_threads())
#endif
for (int64_t row = 0; row < rows; ++row) {
const auto encoded = replay.encode(source + row * replay.dimension);
const int64_t begin = offsets[row];
for (int32_t index = 0; index < static_cast<int32_t>(encoded.size()); ++index) {
output_tokens[begin + index] = encoded[index];
output_anchors[begin + index] = input_anchors[row];
}
}
}
return py::make_tuple(tokens, anchors);
}
|