From Python Call to GPU Kernel: The C++ Underneath ML Libraries

import os, sys, gc, re, json, time, glob, random, shutil, platform, datetime, sysconfig
import subprocess, traceback, contextlib, threading
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
import numpy as np, pandas as pd
import matplotlib.pyplot as plt

subprocess.run([sys.executable, "-m", "pip", "install", "-q", "pybind11"], check=True)
import pybind11

INK, LINE = "#2a2a2a", "#d9d3c7"
BLUE, EMBER, FOREST, GRAY, PLUM = "#3b6ea5", "#c4521e", "#2e7d5b", "#9a9384", "#7a4f8c"
plt.rcParams.update({"font.size": 9, "axes.edgecolor": LINE, "axes.spines.top": False,
                     "axes.spines.right": False, "figure.dpi": 120})

RESULTS = {"meta": {}, "errors": {}}


@contextlib.contextmanager
def experiment(name):
    """One experiment per cell. A failure is printed and recorded, and the run moves on."""
    t0 = time.time()
    try:
        yield
    except Exception:
        RESULTS["errors"][name] = traceback.format_exc()
        print(f"!! {name} failed:\n{traceback.format_exc()}")
    finally:
        print(f"[{name}: {time.time() - t0:.1f} s]")


def sh(cmd, check=True):
    """Run a shell command; return the finished process (stdout and stderr as text)."""
    p = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if check and p.returncode != 0:
        raise RuntimeError(f"command failed ({p.returncode}): {cmd}\n{p.stderr[-4000:]}")
    return p


def build(src, out, flags="-O3 -march=native", std="c++20", extra=""):
    """Compile one benchmark program with g++; return the compile time in seconds."""
    t0 = time.perf_counter()
    sh(f"g++ -std={std} {flags} -pthread -Itinyknn/include {extra} {src} -o {out}")
    return time.perf_counter() - t0


def run_json(cmd):
    """Run a benchmark program and parse the JSON line it prints."""
    lines = [l for l in sh(cmd).stdout.splitlines() if l.startswith("{")]
    return json.loads(lines[-1])


def best_time(fn, reps=5, warmup=1):
    """Fastest of `reps` runs of fn(), in seconds."""
    for _ in range(warmup):
        fn()
    best = float("inf")
    for _ in range(reps):
        t0 = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - t0)
    return best


def jsonable(o):
    if isinstance(o, (np.floating, np.integer)):
        return o.item()
    if isinstance(o, np.ndarray):
        return o.tolist()
    return str(o)


for d in ["tinyknn/include/tinyknn", "tinyknn/src", "tinyknn/python", "bench", "build"]:
    Path(d).mkdir(parents=True, exist_ok=True)

CUDA_HOME = next((p for p in [os.environ.get("CUDA_HOME"), "/usr/local/cuda"]
                  if p and Path(p, "bin", "nvcc").exists()), None)
NVCC = f"{CUDA_HOME}/bin/nvcc" if CUDA_HOME else None
CPU = sh("lscpu | grep 'Model name' | head -1", check=False).stdout.split(":")[-1].strip()
GPU = sh("nvidia-smi --query-gpu=name --format=csv,noheader | head -1", check=False).stdout.strip()
RESULTS["meta"].update(
    cpu=CPU, cores=os.cpu_count(), gpu=GPU or None, python=platform.python_version(),
    numpy=np.__version__, gcc=sh("g++ --version | head -1").stdout.strip(),
    nvcc=sh(f"{NVCC} --version | tail -1", check=False).stdout.strip() if NVCC else None,
)
print(json.dumps(RESULTS["meta"], indent=1))

%%writefile tinyknn/include/tinyknn/buffer.hpp
// tinyknn/buffer.hpp — the one type in the library that owns host memory.
#pragma once
#include <cstddef>
#include <cstring>
#include <new>
#include <utility>

namespace tinyknn {

// n elements of T, starting on a 64-byte boundary (one cache line).
// The constructor allocates, the destructor frees: whoever holds the buffer owns the bytes,
// and the bytes live exactly as long as the buffer does (RAII).
// Copying is disabled on purpose. Moving hands the bytes over; clone() is the only way to
// pay for a second copy, and it has to be written out.
template <class T>
class host_buffer {
 public:
  host_buffer() = default;
  explicit host_buffer(std::size_t n)
      : data_(static_cast<T*>(::operator new(n * sizeof(T), std::align_val_t{64}))), size_(n) {}
  ~host_buffer() { release(); }

  host_buffer(const host_buffer&) = delete;
  host_buffer& operator=(const host_buffer&) = delete;

  host_buffer(host_buffer&& other) noexcept
      : data_(std::exchange(other.data_, nullptr)), size_(std::exchange(other.size_, 0)) {}
  host_buffer& operator=(host_buffer&& other) noexcept {
    if (this != &other) {
      release();
      data_ = std::exchange(other.data_, nullptr);
      size_ = std::exchange(other.size_, 0);
    }
    return *this;
  }

  host_buffer clone() const {
    host_buffer out(size_);
    std::memcpy(out.data_, data_, size_ * sizeof(T));
    return out;
  }

  T* data() noexcept { return data_; }
  const T* data() const noexcept { return data_; }
  std::size_t size() const noexcept { return size_; }

 private:
  void release() noexcept {
    if (data_) ::operator delete(data_, std::align_val_t{64});
  }
  T* data_ = nullptr;
  std::size_t size_ = 0;
};

}  // namespace tinyknn

%%writefile tinyknn/include/tinyknn/view.hpp
// tinyknn/view.hpp — non-owning views. A view answers "how do I read these bytes",
// never "who frees them". It is three numbers and a pointer, cheap to pass by value.
#pragma once
#include <cstddef>
#include <type_traits>

#if defined(__CUDACC__)
#define TINYKNN_HD __host__ __device__
#else
#define TINYKNN_HD
#endif

namespace tinyknn {

// Where the bytes are. A view carries it so that a CPU function cannot be handed GPU memory
// by mistake: the check is one comparison, made before any work starts.
enum class memory { host, device };

// rows × cols elements of T. Row i starts `stride` elements after row i-1, so a view can
// describe a slice of a bigger matrix without copying it. T = const float gives a read-only view.
template <class T>
struct matrix_view {
  T* data = nullptr;
  std::size_t rows = 0, cols = 0, stride = 0;
  memory where = memory::host;

  TINYKNN_HD T* row(std::size_t i) const { return data + i * stride; }
  TINYKNN_HD T& operator()(std::size_t i, std::size_t j) const { return data[i * stride + j]; }
  bool contiguous() const { return stride == cols; }

  // Mutable → read-only converts implicitly; read-only → mutable does not exist.
  template <class U = T, std::enable_if_t<!std::is_const_v<U>, int> = 0>
  operator matrix_view<const U>() const { return {data, rows, cols, stride, where}; }

  // Rows [begin, end) of the same memory — a smaller view, no copy.
  matrix_view slice_rows(std::size_t begin, std::size_t end) const {
    return {data + begin * stride, end - begin, cols, stride, where};
  }
};

template <class T>
matrix_view<T> make_view(T* data, std::size_t rows, std::size_t cols,
                         memory where = memory::host) {
  return {data, rows, cols, cols, where};
}

}  // namespace tinyknn

%%writefile tinyknn/include/tinyknn/distance.hpp
// tinyknn/distance.hpp — metrics as types. Which metric to use is decided when the code is
// compiled, so the inner loop contains the arithmetic of one metric and nothing else.
#pragma once
#include <cstddef>
#include "view.hpp"
#if defined(__cpp_concepts)
#include <concepts>
#endif

namespace tinyknn {

// Every metric is "smaller is closer", so one top-k works for all of them.
struct L2 {
  static constexpr const char* name = "l2";
  static constexpr bool needs_norms = true;
  template <class T>
  TINYKNN_HD static float term(T a, T b) {
    float d = float(a) - float(b);
    return d * d;
  }
  // The same distance rewritten around a dot product: |q - x|² = |q|² + |x|² - 2·q·x.
  // This is what lets a matrix multiplication do most of the work.
  TINYKNN_HD static float from_dot(float q_norm, float x_norm, float dot) {
    return q_norm + x_norm - 2.0f * dot;
  }
};

struct InnerProduct {
  static constexpr const char* name = "ip";
  static constexpr bool needs_norms = false;
  template <class T>
  TINYKNN_HD static float term(T a, T b) { return -float(a) * float(b); }
  TINYKNN_HD static float from_dot(float, float, float dot) { return -dot; }
};

#if defined(__cpp_concepts)
// What the search code needs from a metric, stated once. A type that doesn't fit fails here,
// with this name in the error, instead of deep inside the loop that tried to use it.
template <class M>
concept Metric = requires(float a, float b) {
  { M::template term<float>(a, b) } -> std::convertible_to<float>;
  { M::from_dot(a, a, b) } -> std::convertible_to<float>;
  { M::needs_norms } -> std::convertible_to<bool>;
};
#define TINYKNN_METRIC Metric
#else
#define TINYKNN_METRIC class
#endif

// Distance between two d-dimensional points, summed in LANES independent partial sums.
// Floating-point addition is not associative, so a compiler may not reorder one running sum
// into vector lanes on its own. Writing the lanes out states the order explicitly — the result
// is the same on every run, and the loop is free to become SIMD.
template <class M, int LANES = 8, class T>
TINYKNN_HD inline float distance(const T* a, const T* b, std::size_t d) {
  float acc[LANES] = {};
  std::size_t i = 0;
  for (; i + LANES <= d; i += LANES)
    for (int l = 0; l < LANES; ++l) acc[l] += M::term(a[i + l], b[i + l]);
  for (; i < d; ++i) acc[0] += M::term(a[i], b[i]);
  float s = 0.0f;
  for (int l = 0; l < LANES; ++l) s += acc[l];
  return s;
}

// Four queries against one database row. Every element of x is loaded once and used four
// times, and the 4 × LANES partial sums stay in registers. DIM > 0 fixes the dimension at
// compile time: the loop bound becomes a constant the compiler can unroll around.
// Same summation order as distance<M, LANES>, so both give bit-identical results.
template <class M, std::size_t DIM = 0, int LANES = 8, class T>
inline void distance_x4(const T* q, std::size_t q_stride, const T* x, std::size_t d_runtime,
                        float* out) {
  const std::size_t d = DIM ? DIM : d_runtime;
  float acc[4][LANES] = {};
  std::size_t i = 0;
  for (; i + LANES <= d; i += LANES)
    for (int b = 0; b < 4; ++b)
      for (int l = 0; l < LANES; ++l) acc[b][l] += M::term(q[b * q_stride + i + l], x[i + l]);
  if constexpr (DIM == 0 || DIM % LANES != 0)
    for (; i < d; ++i)
      for (int b = 0; b < 4; ++b) acc[b][0] += M::term(q[b * q_stride + i], x[i]);
  for (int b = 0; b < 4; ++b) {
    float s = 0.0f;
    for (int l = 0; l < LANES; ++l) s += acc[b][l];
    out[b] = s;
  }
}

}  // namespace tinyknn

%%writefile tinyknn/include/tinyknn/topk.hpp
// tinyknn/topk.hpp — keep the k smallest of a stream of (distance, id) pairs.
#pragma once
#include <algorithm>
#include <cstdint>
#include <limits>
#include <utility>
#include <vector>

namespace tinyknn {

// A max-heap of the k best so far: the root is the worst of the kept candidates, so a new
// candidate is compared with one number and usually rejected. With N ≫ k almost every
// candidate is, and the heap is touched O(k log N) times in total.
class topk {
 public:
  explicit topk(std::size_t k) : k_(k) { heap_.reserve(k); }

  float worst() const {
    return heap_.size() < k_ ? std::numeric_limits<float>::infinity() : heap_.front().first;
  }
  void push(float dist, std::int64_t id) {
    if (heap_.size() < k_) {
      heap_.emplace_back(dist, id);
      std::push_heap(heap_.begin(), heap_.end());
    } else if (dist < heap_.front().first) {
      std::pop_heap(heap_.begin(), heap_.end());
      heap_.back() = {dist, id};
      std::push_heap(heap_.begin(), heap_.end());
    }
  }
  // Write the kept pairs, closest first. Unfilled slots (N < k) get id -1.
  void write(std::int64_t* ids, float* dists) {
    std::sort_heap(heap_.begin(), heap_.end());
    for (std::size_t i = 0; i < k_; ++i) {
      bool have = i < heap_.size();
      ids[i] = have ? heap_[i].second : -1;
      dists[i] = have ? heap_[i].first : std::numeric_limits<float>::infinity();
    }
    heap_.clear();
  }

 private:
  std::size_t k_;
  std::vector<std::pair<float, std::int64_t>> heap_;
};

}  // namespace tinyknn

%%writefile tinyknn/include/tinyknn/search_cpu.hpp
// tinyknn/search_cpu.hpp — brute-force k-nearest-neighbour search on the CPU.
// A template: it is compiled once for every (metric, element type, dimension) it is used with.
#pragma once
#include <algorithm>
#include <cstdint>
#include <stdexcept>
#include <thread>
#include <vector>
#include "distance.hpp"
#include "topk.hpp"
#include "view.hpp"

namespace tinyknn {
namespace detail {

// Queries [q0, q1) against the whole database. The database is walked in blocks of
// `db_block` rows; every block is used by all of these queries while it is still in cache,
// so each row comes from main memory once per thread instead of once per four queries.
// db_block = 0 turns that off: each group of four queries walks the whole database itself.
template <class M, std::size_t DIM, class T>
void scan(matrix_view<const T> queries, matrix_view<const T> db, std::size_t q0, std::size_t q1,
          std::size_t db_block, std::vector<topk>& best) {
  const std::size_t dim = db.cols, block = db_block ? db_block : db.rows;
  for (std::size_t r0 = 0; r0 < db.rows; r0 += block) {
    const std::size_t r1 = std::min(db.rows, r0 + block);
    std::size_t q = q0;
    for (; q + 4 <= q1; q += 4)
      for (std::size_t n = r0; n < r1; ++n) {
        float d[4];
        distance_x4<M, DIM>(queries.row(q), queries.stride, db.row(n), dim, d);
        for (int b = 0; b < 4; ++b) best[q - q0 + b].push(d[b], std::int64_t(n));
      }
    for (; q < q1; ++q)
      for (std::size_t n = r0; n < r1; ++n)
        best[q - q0].push(distance<M>(queries.row(q), db.row(n), dim), std::int64_t(n));
  }
}

// The dimension is a run-time number; the kernels are faster when it is a compile-time one.
// This switch turns the common sizes into template arguments — one instantiation each —
// and sends everything else to the general version.
template <class M, class T, class... Args>
void scan_any_dim(std::size_t dim, bool fixed_dim, Args&&... args) {
  if (fixed_dim) switch (dim) {
      case 64: return scan<M, 64, T>(args...);
      case 128: return scan<M, 128, T>(args...);
      case 256: return scan<M, 256, T>(args...);
      case 384: return scan<M, 384, T>(args...);
      case 768: return scan<M, 768, T>(args...);
    }
  scan<M, 0, T>(args...);
}

}  // namespace detail

// queries: Q × D, db: N × D (read-only).  ids, dists: Q × k (written).
// Nothing is allocated for the data and nothing is returned: the caller decides where every
// byte lives, and the signature says which bytes the function may change.
template <TINYKNN_METRIC M, class T>
void search_cpu(matrix_view<const T> queries, matrix_view<const T> db,
                matrix_view<std::int64_t> ids, matrix_view<float> dists,
                int threads = 1, std::size_t db_block = 512, bool fixed_dim = true) {
  if (queries.where != memory::host || db.where != memory::host ||
      ids.where != memory::host || dists.where != memory::host)
    throw std::invalid_argument("search_cpu: every view must point to host memory");
  if (queries.cols != db.cols) throw std::invalid_argument("search_cpu: dimension mismatch");
  if (ids.rows != queries.rows || dists.rows != queries.rows || ids.cols != dists.cols)
    throw std::invalid_argument("search_cpu: output shape must be Q × k");
  const std::size_t k = ids.cols, Q = queries.rows;
  if (k == 0 || Q == 0) return;

  // Each worker owns a contiguous range of queries and one top-k per query. Workers share the
  // database read-only and write disjoint rows of the output: nothing to lock.
  auto work = [&](std::size_t q0, std::size_t q1) {
    std::vector<topk> best(q1 - q0, topk(k));
    detail::scan_any_dim<M, T>(db.cols, fixed_dim, queries, db, q0, q1, db_block, best);
    for (std::size_t q = q0; q < q1; ++q) best[q - q0].write(ids.row(q), dists.row(q));
  };

  threads = std::max(1, std::min<int>(threads, int(Q)));
  if (threads == 1) return work(0, Q);
  std::vector<std::thread> pool;
  const std::size_t per = (Q + threads - 1) / threads;
  for (int t = 0; t < threads; ++t) {
    const std::size_t q0 = t * per, q1 = std::min(Q, q0 + per);
    if (q0 < q1) pool.emplace_back(work, q0, q1);
  }
  for (auto& th : pool) th.join();
}

}  // namespace tinyknn

%%writefile tinyknn/include/tinyknn/api.hpp
// tinyknn/api.hpp — what libtinyknn.so exports. Templates cannot cross a shared-library
// boundary (they are recipes, not machine code), so the public surface is a handful of
// ordinary functions and one class, each compiled once inside the library.
#pragma once
#include <cstdint>
#include <memory>
#include "view.hpp"

struct CUstream_st;  // cudaStream_t without including CUDA headers in every user of the library

namespace tinyknn {

enum class metric { l2, inner_product };

// CPU search; every view points to host memory.
void search(metric m, matrix_view<const float> queries, matrix_view<const float> db,
            matrix_view<std::int64_t> ids, matrix_view<float> dists,
            int threads = 1, std::size_t db_block = 512, bool fixed_dim = true);

bool has_cuda();

// How the GPU computes the Q × N distance matrix before the top-k pass.
enum class gpu_algo {
  naive,  // one thread per (query, row) pair, walking the row itself
  gemm,   // norms + one matrix multiplication (cuBLAS), then norms added back in the top-k pass
};

// How the GPU picks the k smallest of each row of distances.
enum class topk_algo {
  runtime_k,       // one block per query; k is a run-time number, so each thread's list of
                   // candidates is indexed at run time and lives in slow per-thread memory
  compile_time_k,  // one block per query; k is a template argument (1, 4, 8, 10 or 16), so the
                   // list is fully unrolled into registers
  split_rows,      // as compile_time_k, but a row is split across many blocks, and a second
                   // pass merges their candidates — so a handful of queries still fills the GPU
};

// Low level: every pointer is device memory, nothing is allocated, and the call returns as soon
// as the work is queued on `stream`, typically long before the GPU has done it.
// scratch must hold queries.rows × (db.rows + 2048) floats; db_norms / query_norms may be null
// for metrics that don't need them (query_norms is scratch space this call fills).
void search_async(metric m, gpu_algo algo, topk_algo select,
                  matrix_view<const float> queries, matrix_view<const float> db,
                  const float* db_norms, float* query_norms, float* scratch,
                  matrix_view<std::int64_t> ids, matrix_view<float> dists,
                  CUstream_st* stream);

// Stage times of one gpu_index::search call, in milliseconds.
struct trace {
  double stage_in = 0;   // host: caller's queries → pinned staging buffer
  double queue = 0;      // host: issuing the copies and kernels (they run later)
  double h2d = 0;        // staging → GPU
  double distances = 0;  // Q × N distances
  double topk = 0;       // k smallest per query
  double d2h = 0;        // results → pinned staging
  double stage_out = 0;  // pinned staging → caller's arrays
  double total = 0;      // whole call, host clock
};

// High level: owns a copy of the database on the GPU, plus everything a search needs.
// The implementation (CUDA types, cuBLAS handle, buffers) is hidden behind a pointer, so this
// header — and the binary layout of the class — doesn't change when the internals do.
class gpu_index {
 public:
  gpu_index(matrix_view<const float> db, metric m);  // copies db to the GPU once
  ~gpu_index();
  gpu_index(gpu_index&&) noexcept;
  gpu_index& operator=(gpu_index&&) noexcept;

  std::size_t size() const;
  std::size_t dim() const;

  // Synchronous: returns when the results are in ids/dists (host memory).
  // With t != nullptr it stops after every stage to time it, which makes the call slower.
  void search(matrix_view<const float> queries, matrix_view<std::int64_t> ids,
              matrix_view<float> dists, gpu_algo algo,
              topk_algo select = topk_algo::split_rows, trace* t = nullptr);

  // Asynchronous, for at most chunk() queries: queue the whole search and return.
  // The results land in pinned host memory owned by the index (result_ids / result_dists),
  // and are only valid after wait().
  void enqueue(matrix_view<const float> queries, std::size_t k, gpu_algo algo,
               topk_algo select = topk_algo::split_rows);
  void wait();
  const std::int64_t* result_ids() const;
  const float* result_dists() const;
  std::size_t chunk() const;

 private:
  struct impl;
  std::unique_ptr<impl> p_;
};

}  // namespace tinyknn

%%writefile tinyknn/src/cpu.cpp
// src/cpu.cpp — the CPU half of libtinyknn.so. The metric is a run-time value in the public
// API and a compile-time type inside: the switch below is the one place where one becomes
// the other, once per call, outside every loop.
#include "tinyknn/api.hpp"
#include "tinyknn/search_cpu.hpp"

namespace tinyknn {

void search(metric m, matrix_view<const float> queries, matrix_view<const float> db,
            matrix_view<std::int64_t> ids, matrix_view<float> dists,
            int threads, std::size_t db_block, bool fixed_dim) {
  switch (m) {
    case metric::l2:
      return search_cpu<L2>(queries, db, ids, dists, threads, db_block, fixed_dim);
    case metric::inner_product:
      return search_cpu<InnerProduct>(queries, db, ids, dists, threads, db_block, fixed_dim);
  }
}

#ifndef TINYKNN_WITH_CUDA
bool has_cuda() { return false; }
#endif

}  // namespace tinyknn

%%writefile tinyknn/src/cuda.cu
// src/cuda.cu — the GPU half of libtinyknn.so. Compiled by nvcc; nothing outside this file
// sees a CUDA type except the opaque stream pointer in api.hpp.
#include <cublas_v2.h>
#include <cuda_runtime.h>

#include <chrono>
#include <cmath>
#include <cstring>
#include <stdexcept>
#include <algorithm>
#include <string>
#include <type_traits>
#include <utility>

#include "tinyknn/api.hpp"
#include "tinyknn/distance.hpp"

namespace tinyknn {
namespace {

void check(cudaError_t e, const char* what) {
  if (e != cudaSuccess) throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(e));
}
void check(cublasStatus_t s, const char* what) {
  if (s != CUBLAS_STATUS_SUCCESS)
    throw std::runtime_error(std::string(what) + ": cuBLAS status " + std::to_string(int(s)));
}
#define CK(x) check((x), #x)

// GPU memory with the same contract as host_buffer: allocated in the constructor, freed in
// the destructor, moved but never copied.
template <class T>
class device_buffer {
 public:
  device_buffer() = default;
  explicit device_buffer(std::size_t n) : n_(n) {
    if (n) CK(cudaMalloc(reinterpret_cast<void**>(&p_), n * sizeof(T)));
  }
  ~device_buffer() {
    if (p_) cudaFree(p_);
  }
  device_buffer(const device_buffer&) = delete;
  device_buffer& operator=(const device_buffer&) = delete;
  device_buffer(device_buffer&& o) noexcept
      : p_(std::exchange(o.p_, nullptr)), n_(std::exchange(o.n_, 0)) {}
  device_buffer& operator=(device_buffer&& o) noexcept {
    if (this != &o) {
      if (p_) cudaFree(p_);
      p_ = std::exchange(o.p_, nullptr);
      n_ = std::exchange(o.n_, 0);
    }
    return *this;
  }
  T* data() const { return p_; }
  std::size_t size() const { return n_; }

 private:
  T* p_ = nullptr;
  std::size_t n_ = 0;
};

// Page-locked ("pinned") host memory. The GPU's copy engine can read and write it directly,
// which is what makes an asynchronous copy possible at all.
template <class T>
class pinned_buffer {
 public:
  pinned_buffer() = default;
  explicit pinned_buffer(std::size_t n) : n_(n) {
    if (n) CK(cudaMallocHost(reinterpret_cast<void**>(&p_), n * sizeof(T)));
  }
  ~pinned_buffer() {
    if (p_) cudaFreeHost(p_);
  }
  pinned_buffer(const pinned_buffer&) = delete;
  pinned_buffer& operator=(const pinned_buffer&) = delete;
  pinned_buffer(pinned_buffer&& o) noexcept
      : p_(std::exchange(o.p_, nullptr)), n_(std::exchange(o.n_, 0)) {}
  pinned_buffer& operator=(pinned_buffer&& o) noexcept {
    if (this != &o) {
      if (p_) cudaFreeHost(p_);
      p_ = std::exchange(o.p_, nullptr);
      n_ = std::exchange(o.n_, 0);
    }
    return *this;
  }
  T* data() const { return p_; }

 private:
  T* p_ = nullptr;
  std::size_t n_ = 0;
};

constexpr int KMAX = 16;          // largest k the top-k kernel supports
constexpr int TOPK_THREADS = 256;  // one block of this many threads per query

__global__ void row_norms(const float* x, std::size_t rows, std::size_t dim, float* out) {
  std::size_t r = blockIdx.x * std::size_t(blockDim.x) + threadIdx.x;
  if (r >= rows) return;
  const float* p = x + r * dim;
  float s = 0.0f;
  for (std::size_t i = 0; i < dim; ++i) s += p[i] * p[i];
  out[r] = s;
}

// One thread per (query, database row): the same distance<M> the CPU uses, run 25.6 million
// times in parallel. Neighbouring threads read rows `dim` floats apart, so their loads don't
// combine into wide memory transactions — the price of the simplest possible mapping.
template <class M>
__global__ void distances_naive(const float* q, const float* db, std::size_t n, std::size_t dim,
                                float* out) {
  std::size_t row = blockIdx.x * std::size_t(blockDim.x) + threadIdx.x;
  std::size_t qi = blockIdx.y;
  if (row >= n) return;
  out[qi * n + row] = distance<M>(q + qi * dim, db + row * dim, dim);
}

// One block per query. Each thread keeps its own sorted k best while striding over the row,
// then the 256 lists are merged pairwise in shared memory: 8 rounds of O(k) work.
template <class M, bool FROM_DOT>
__global__ void topk_rows(const float* scores, std::size_t n, int k, const float* q_norms,
                          const float* x_norms, std::int64_t* ids, float* dists) {
  __shared__ float sd[TOPK_THREADS * KMAX];
  __shared__ int si[TOPK_THREADS * KMAX];
  const int t = threadIdx.x;
  const std::size_t q = blockIdx.x;
  const float* row = scores + q * n;

  float bd[KMAX];
  int bi[KMAX];
  for (int j = 0; j < KMAX; ++j) {
    bd[j] = INFINITY;
    bi[j] = -1;
  }
  const float qn = (FROM_DOT && q_norms) ? q_norms[q] : 0.0f;
  for (std::size_t i = t; i < n; i += blockDim.x) {
    float v = row[i];
    if (FROM_DOT) v = M::from_dot(qn, x_norms ? x_norms[i] : 0.0f, v);
    if (v < bd[k - 1]) {
      int j = k - 1;
      while (j > 0 && bd[j - 1] > v) {
        bd[j] = bd[j - 1];
        bi[j] = bi[j - 1];
        --j;
      }
      bd[j] = v;
      bi[j] = int(i);
    }
  }
  for (int j = 0; j < k; ++j) {
    sd[t * KMAX + j] = bd[j];
    si[t * KMAX + j] = bi[j];
  }
  for (int s = 1; s < blockDim.x; s *= 2) {
    __syncthreads();
    if (t % (2 * s) == 0) {
      const float* ad = sd + t * KMAX;
      const int* ai = si + t * KMAX;
      const float* od = sd + (t + s) * KMAX;
      const int* oi = si + (t + s) * KMAX;
      float md[KMAX];
      int mi[KMAX];
      int a = 0, b = 0;
      for (int j = 0; j < k; ++j) {
        if (ad[a] <= od[b]) {
          md[j] = ad[a];
          mi[j] = ai[a++];
        } else {
          md[j] = od[b];
          mi[j] = oi[b++];
        }
      }
      for (int j = 0; j < k; ++j) {
        sd[t * KMAX + j] = md[j];
        si[t * KMAX + j] = mi[j];
      }
    }
  }
  __syncthreads();
  if (t < k) {
    ids[q * k + t] = si[t];
    dists[q * k + t] = sd[t];
  }
}

// Insert (v, i) into a sorted list of K held by one thread. With K a compile-time constant the
// loop unrolls completely, every index is a constant, and the list stays in registers.
template <int K>
__device__ __forceinline__ void insert(float (&d)[K], int (&id)[K], float v, int i) {
  if (!(v < d[K - 1])) return;
#pragma unroll
  for (int j = 0; j < K; ++j)
    if (v < d[j]) {
      const float tv = d[j];
      const int ti = id[j];
      d[j] = v;
      id[j] = i;
      v = tv;
      i = ti;
    }
}

// Merge the lists of all threads in a block, pairwise through shared memory.
// Afterwards thread 0 holds the block's K best, in its registers.
template <int K>
__device__ void block_select(float (&d)[K], int (&id)[K], float* sd, int* si) {
  const int t = threadIdx.x;
#pragma unroll
  for (int j = 0; j < K; ++j) {
    sd[t * K + j] = d[j];
    si[t * K + j] = id[j];
  }
  for (int s = 1; s < blockDim.x; s *= 2) {
    __syncthreads();
    if (t % (2 * s) == 0) {
#pragma unroll
      for (int j = 0; j < K; ++j) insert<K>(d, id, sd[(t + s) * K + j], si[(t + s) * K + j]);
#pragma unroll
      for (int j = 0; j < K; ++j) {
        sd[t * K + j] = d[j];
        si[t * K + j] = id[j];
      }
    }
  }
}

// Top-K of a slice of each row. Grid: (slices, queries). With one slice the block writes the
// final answer; with several, each block writes K candidates for topk_merge.
template <class M, bool FROM_DOT, int K>
__global__ void topk_fixed(const float* scores, std::size_t n, std::size_t slice,
                           const float* q_norms, const float* x_norms, std::int64_t* ids,
                           float* dists, float* cand_d, int* cand_i) {
  __shared__ float sd[TOPK_THREADS * K];
  __shared__ int si[TOPK_THREADS * K];
  const std::size_t q = blockIdx.y, s = blockIdx.x, S = gridDim.x;
  const std::size_t begin = s * slice, end = begin + slice < n ? begin + slice : n;
  const float* row = scores + q * n;
  float d[K];
  int id[K];
#pragma unroll
  for (int j = 0; j < K; ++j) {
    d[j] = INFINITY;
    id[j] = -1;
  }
  const float qn = (FROM_DOT && q_norms) ? q_norms[q] : 0.0f;
  for (std::size_t i = begin + threadIdx.x; i < end; i += blockDim.x) {
    float v = row[i];
    if (FROM_DOT) v = M::from_dot(qn, x_norms ? x_norms[i] : 0.0f, v);
    insert<K>(d, id, v, int(i));
  }
  block_select<K>(d, id, sd, si);
  if (threadIdx.x == 0) {
#pragma unroll
    for (int j = 0; j < K; ++j) {
      if (S == 1) {
        ids[q * K + j] = id[j];
        dists[q * K + j] = d[j];
      } else {
        cand_d[(q * S + s) * K + j] = d[j];
        cand_i[(q * S + s) * K + j] = id[j];
      }
    }
  }
}

// Second pass of split_rows: the K best of each query's S × K candidates.
template <int K>
__global__ void topk_merge(const float* cand_d, const int* cand_i, int S, std::int64_t* ids,
                           float* dists) {
  __shared__ float sd[TOPK_THREADS * K];
  __shared__ int si[TOPK_THREADS * K];
  const std::size_t q = blockIdx.x;
  float d[K];
  int id[K];
#pragma unroll
  for (int j = 0; j < K; ++j) {
    d[j] = INFINITY;
    id[j] = -1;
  }
  for (int c = threadIdx.x; c < S * K; c += blockDim.x)
    insert<K>(d, id, cand_d[q * S * K + c], cand_i[q * S * K + c]);
  block_select<K>(d, id, sd, si);
  if (threadIdx.x == 0)
#pragma unroll
    for (int j = 0; j < K; ++j) {
      ids[q * K + j] = id[j];
      dists[q * K + j] = d[j];
    }
}

int sm_count() {
  static const int n = [] {
    int dev = 0, v = 0;
    CK(cudaGetDevice(&dev));
    CK(cudaDeviceGetAttribute(&v, cudaDevAttrMultiProcessorCount, dev));
    return v;
  }();
  return n;
}

cublasHandle_t blas() {
  thread_local cublasHandle_t h = [] {
    cublasHandle_t x;
    CK(cublasCreate(&x));
    return x;
  }();
  return h;
}

void validate_device(matrix_view<const float> queries, matrix_view<const float> db,
                     matrix_view<std::int64_t> ids, matrix_view<float> dists) {
  if (queries.where != memory::device || db.where != memory::device ||
      ids.where != memory::device || dists.where != memory::device)
    throw std::invalid_argument("search_async: every view must point to device memory");
  if (!queries.contiguous() || !db.contiguous() || !ids.contiguous() || !dists.contiguous())
    throw std::invalid_argument("search_async: device views must be contiguous");
  if (queries.cols != db.cols) throw std::invalid_argument("search_async: dimension mismatch");
  if (ids.rows != queries.rows || dists.rows != queries.rows || ids.cols != dists.cols)
    throw std::invalid_argument("search_async: output shape must be Q × k");
  if (ids.cols < 1 || ids.cols > KMAX)
    throw std::invalid_argument("search_async: k must be between 1 and " + std::to_string(KMAX));
}

template <class M>
void launch_distances(gpu_algo algo, matrix_view<const float> q, matrix_view<const float> db,
                      float* query_norms, float* scratch, cudaStream_t s) {
  const std::size_t Q = q.rows, N = db.rows, D = db.cols;
  if (algo == gpu_algo::naive) {
    dim3 grid(unsigned((N + 255) / 256), unsigned(Q));
    distances_naive<M><<<grid, 256, 0, s>>>(q.data, db.data, N, D, scratch);
  } else {
    if (M::needs_norms) row_norms<<<unsigned((Q + 255) / 256), 256, 0, s>>>(q.data, Q, D, query_norms);
    // Row-major Q×D and N×D are column-major D×Q and D×N to cuBLAS. Asking for dbᵀ·q gives an
    // N×Q column-major result, which is exactly the Q×N row-major matrix of dot products.
    const float one = 1.0f, zero = 0.0f;
    CK(cublasSetStream(blas(), s));
    CK(cublasSgemm(blas(), CUBLAS_OP_T, CUBLAS_OP_N, int(N), int(Q), int(D), &one, db.data,
                   int(D), q.data, int(D), &zero, scratch, int(N)));
  }
  CK(cudaGetLastError());
}

template <class M, bool FROM_DOT>
void launch_select(topk_algo select, std::size_t Q, std::size_t N, int k, const float* scores,
                   const float* q_norms, const float* x_norms, std::int64_t* ids, float* dists,
                   float* cand, cudaStream_t s) {
  // k is a run-time number; the fast kernels need it as a type. Same move as the dimension
  // switch on the CPU: a few common values are compiled, anything else takes the general path.
  auto fixed = [&](auto k_const) {
    constexpr int K = decltype(k_const)::value;
    unsigned S = 1;
    if (select == topk_algo::split_rows) {
      // Enough blocks for about four per SM, but no slice shorter than 1,024 values.
      const std::size_t want = (4 * std::size_t(sm_count()) + Q - 1) / Q;
      S = unsigned(std::max<std::size_t>(1, std::min<std::size_t>({want, 64, N / 1024})));
    }
    const std::size_t slice = (N + S - 1) / S;
    float* cand_d = cand;
    int* cand_i = reinterpret_cast<int*>(cand + Q * S * K);
    topk_fixed<M, FROM_DOT, K><<<dim3(S, unsigned(Q)), TOPK_THREADS, 0, s>>>(
        scores, N, slice, q_norms, x_norms, ids, dists, cand_d, cand_i);
    if (S > 1) topk_merge<K><<<unsigned(Q), TOPK_THREADS, 0, s>>>(cand_d, cand_i, int(S), ids, dists);
  };
  if (select != topk_algo::runtime_k) switch (k) {
      case 1: return fixed(std::integral_constant<int, 1>{});
      case 4: return fixed(std::integral_constant<int, 4>{});
      case 8: return fixed(std::integral_constant<int, 8>{});
      case 10: return fixed(std::integral_constant<int, 10>{});
      case 16: return fixed(std::integral_constant<int, 16>{});
    }
  topk_rows<M, FROM_DOT><<<unsigned(Q), TOPK_THREADS, 0, s>>>(scores, N, k, q_norms, x_norms, ids, dists);
}

// scratch holds Q × N distances, then room for the split_rows candidates (≤ Q × 64 × 16 pairs).
template <class M>
void launch_topk(gpu_algo algo, topk_algo select, std::size_t Q, std::size_t N,
                 const float* db_norms, const float* query_norms, float* scratch,
                 matrix_view<std::int64_t> ids, matrix_view<float> dists, cudaStream_t s) {
  const int k = int(ids.cols);
  float* cand = scratch + Q * N;
  if (algo == gpu_algo::naive)
    launch_select<M, false>(select, Q, N, k, scratch, nullptr, nullptr, ids.data, dists.data, cand, s);
  else
    launch_select<M, true>(select, Q, N, k, scratch, M::needs_norms ? query_norms : nullptr,
                           M::needs_norms ? db_norms : nullptr, ids.data, dists.data, cand, s);
  CK(cudaGetLastError());
}

// The run-time metric becomes a type here, once per call — the same move as in cpu.cpp.
template <class F>
void with_metric(metric m, F&& f) {
  if (m == metric::l2) f(L2{});
  else f(InnerProduct{});
}

double ms_since(std::chrono::steady_clock::time_point t0) {
  return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count();
}

}  // namespace

bool has_cuda() {
  int n = 0;
  return cudaGetDeviceCount(&n) == cudaSuccess && n > 0;
}

void search_async(metric m, gpu_algo algo, topk_algo select, matrix_view<const float> queries,
                  matrix_view<const float> db, const float* db_norms, float* query_norms,
                  float* scratch, matrix_view<std::int64_t> ids, matrix_view<float> dists,
                  CUstream_st* stream) {
  validate_device(queries, db, ids, dists);
  with_metric(m, [&](auto metric_tag) {
    using M = decltype(metric_tag);
    if (algo == gpu_algo::gemm && M::needs_norms && (!db_norms || !query_norms))
      throw std::invalid_argument("search_async: this metric needs norms for the gemm path");
    launch_distances<M>(algo, queries, db, query_norms, scratch, stream);
    launch_topk<M>(algo, select, queries.rows, db.rows, db_norms, query_norms, scratch, ids, dists, stream);
  });
}

struct gpu_index::impl {
  metric m;
  std::size_t n, dim, chunk = 256, k_alloc = 0;
  device_buffer<float> db, db_norms, q, q_norms, scratch;
  device_buffer<std::int64_t> ids;
  device_buffer<float> dists;
  pinned_buffer<float> q_stage;
  pinned_buffer<std::int64_t> ids_stage;
  pinned_buffer<float> dists_stage;
  cudaStream_t stream = nullptr;
  cudaEvent_t ev[5] = {};
  std::size_t pending_rows = 0, pending_k = 0;

  impl(matrix_view<const float> host_db, metric metric_)
      : m(metric_), n(host_db.rows), dim(host_db.cols) {
    if (host_db.where != memory::host) throw std::invalid_argument("gpu_index: db must be host memory");
    db = device_buffer<float>(n * dim);
    CK(cudaMemcpy2D(db.data(), dim * sizeof(float), host_db.data, host_db.stride * sizeof(float),
                    dim * sizeof(float), n, cudaMemcpyHostToDevice));
    db_norms = device_buffer<float>(n);
    row_norms<<<unsigned((n + 255) / 256), 256>>>(db.data(), n, dim, db_norms.data());
    CK(cudaGetLastError());
    q = device_buffer<float>(chunk * dim);
    q_norms = device_buffer<float>(chunk);
    scratch = device_buffer<float>(chunk * (n + 2048));
    q_stage = pinned_buffer<float>(chunk * dim);
    CK(cudaDeviceSynchronize());
    CK(cudaStreamCreate(&stream));
    for (auto& e : ev) CK(cudaEventCreate(&e));
  }
  ~impl() {
    if (stream) cudaStreamSynchronize(stream);
    for (auto& e : ev)
      if (e) cudaEventDestroy(e);
    if (stream) cudaStreamDestroy(stream);
  }

  void ensure_k(std::size_t k) {
    if (k < 1 || k > std::size_t(KMAX))
      throw std::invalid_argument("gpu_index: k must be between 1 and " + std::to_string(KMAX));
    if (k <= k_alloc) return;
    ids = device_buffer<std::int64_t>(chunk * k);
    dists = device_buffer<float>(chunk * k);
    ids_stage = pinned_buffer<std::int64_t>(chunk * k);
    dists_stage = pinned_buffer<float>(chunk * k);
    k_alloc = k;
  }

  // Queue one chunk: copy in, compute, copy out. Events mark the stage boundaries.
  // Returns the host time of the two parts: {staging memcpy, issuing the GPU work}.
  std::pair<double, double> queue_chunk(matrix_view<const float> rows, std::size_t k, gpu_algo algo,
                                        topk_algo select) {
    const std::size_t R = rows.rows;
    // The staging buffers are reused: the previous chunk's copies must be finished first.
    CK(cudaStreamSynchronize(stream));
    const auto t0 = std::chrono::steady_clock::now();
    for (std::size_t i = 0; i < R; ++i)
      std::memcpy(q_stage.data() + i * dim, rows.row(i), dim * sizeof(float));
    const double staged = ms_since(t0);
    const auto t1 = std::chrono::steady_clock::now();
    CK(cudaEventRecord(ev[0], stream));
    CK(cudaMemcpyAsync(q.data(), q_stage.data(), R * dim * sizeof(float), cudaMemcpyHostToDevice, stream));
    CK(cudaEventRecord(ev[1], stream));
    auto dq = make_view<const float>(q.data(), R, dim, memory::device);
    auto ddb = make_view<const float>(db.data(), n, dim, memory::device);
    auto did = make_view(ids.data(), R, k, memory::device);
    auto ddi = make_view(dists.data(), R, k, memory::device);
    with_metric(m, [&](auto tag) {
      using M = decltype(tag);
      launch_distances<M>(algo, dq, ddb, q_norms.data(), scratch.data(), stream);
      CK(cudaEventRecord(ev[2], stream));
      launch_topk<M>(algo, select, R, n, db_norms.data(), q_norms.data(), scratch.data(), did, ddi, stream);
    });
    CK(cudaEventRecord(ev[3], stream));
    CK(cudaMemcpyAsync(ids_stage.data(), ids.data(), R * k * sizeof(std::int64_t), cudaMemcpyDeviceToHost, stream));
    CK(cudaMemcpyAsync(dists_stage.data(), dists.data(), R * k * sizeof(float), cudaMemcpyDeviceToHost, stream));
    CK(cudaEventRecord(ev[4], stream));
    return {staged, ms_since(t1)};
  }
};

gpu_index::gpu_index(matrix_view<const float> db, metric m) : p_(std::make_unique<impl>(db, m)) {}
gpu_index::~gpu_index() = default;
gpu_index::gpu_index(gpu_index&&) noexcept = default;
gpu_index& gpu_index::operator=(gpu_index&&) noexcept = default;
std::size_t gpu_index::size() const { return p_->n; }
std::size_t gpu_index::dim() const { return p_->dim; }
std::size_t gpu_index::chunk() const { return p_->chunk; }

void gpu_index::search(matrix_view<const float> queries, matrix_view<std::int64_t> ids,
                       matrix_view<float> dists, gpu_algo algo, topk_algo select, trace* t) {
  auto& P = *p_;
  if (queries.where != memory::host || ids.where != memory::host || dists.where != memory::host)
    throw std::invalid_argument("gpu_index::search: queries and outputs are host memory");
  if (queries.cols != P.dim) throw std::invalid_argument("gpu_index::search: dimension mismatch");
  if (ids.rows != queries.rows || dists.rows != queries.rows || ids.cols != dists.cols)
    throw std::invalid_argument("gpu_index::search: output shape must be Q × k");
  const std::size_t k = ids.cols;
  P.ensure_k(k);
  const auto t_call = std::chrono::steady_clock::now();
  if (t) *t = trace{};
  for (std::size_t c0 = 0; c0 < queries.rows; c0 += P.chunk) {
    const std::size_t c1 = std::min(queries.rows, c0 + P.chunk), R = c1 - c0;
    const auto host = P.queue_chunk(queries.slice_rows(c0, c1), k, algo, select);
    if (t) {
      t->stage_in += host.first;
      t->queue += host.second;
    }
    CK(cudaStreamSynchronize(P.stream));
    if (t) {
      float a, b, c, d;
      CK(cudaEventElapsedTime(&a, P.ev[0], P.ev[1]));
      CK(cudaEventElapsedTime(&b, P.ev[1], P.ev[2]));
      CK(cudaEventElapsedTime(&c, P.ev[2], P.ev[3]));
      CK(cudaEventElapsedTime(&d, P.ev[3], P.ev[4]));
      t->h2d += a;
      t->distances += b;
      t->topk += c;
      t->d2h += d;
    }
    auto t1 = std::chrono::steady_clock::now();
    for (std::size_t i = 0; i < R; ++i) {
      std::memcpy(ids.row(c0 + i), P.ids_stage.data() + i * k, k * sizeof(std::int64_t));
      std::memcpy(dists.row(c0 + i), P.dists_stage.data() + i * k, k * sizeof(float));
    }
    if (t) t->stage_out += ms_since(t1);
  }
  if (t) t->total = ms_since(t_call);
}

void gpu_index::enqueue(matrix_view<const float> queries, std::size_t k, gpu_algo algo,
                        topk_algo select) {
  auto& P = *p_;
  if (queries.where != memory::host || queries.cols != P.dim)
    throw std::invalid_argument("gpu_index::enqueue: host queries of the index's dimension");
  if (queries.rows > P.chunk)
    throw std::invalid_argument("gpu_index::enqueue: at most chunk() queries");
  P.ensure_k(k);
  P.queue_chunk(queries, k, algo, select);
  P.pending_rows = queries.rows;
  P.pending_k = k;
}
void gpu_index::wait() { CK(cudaStreamSynchronize(p_->stream)); }
const std::int64_t* gpu_index::result_ids() const { return p_->ids_stage.data(); }
const float* gpu_index::result_dists() const { return p_->dists_stage.data(); }

}  // namespace tinyknn

%%writefile tinyknn/python/bindings.cpp
// python/bindings.cpp — the Python module _tinyknn. Turns NumPy arrays into views, calls the
// library, turns the results back into NumPy arrays. No numerical code lives here.
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <chrono>
#include <cstdint>
#include <string>

#include "tinyknn/api.hpp"

namespace py = pybind11;
using clk = std::chrono::steady_clock;

// "Give me C-contiguous float32 — and if the array isn't, convert it."
// The conversion happens before the function body runs: a float64 array arrives as a fresh
// float32 copy, silently. `strict` arrays refuse instead (see probe_strict).
using farray = py::array_t<float, py::array::c_style | py::array::forcecast>;
using strict = py::array_t<float, py::array::c_style>;

static tinyknn::matrix_view<const float> view_of(const farray& a, const char* name) {
  if (a.ndim() != 2) throw std::invalid_argument(std::string(name) + " must be 2-D");
  return tinyknn::make_view(a.data(), std::size_t(a.shape(0)), std::size_t(a.shape(1)));
}

static tinyknn::metric metric_of(const std::string& m) {
  if (m == "l2") return tinyknn::metric::l2;
  if (m == "ip") return tinyknn::metric::inner_product;
  throw std::invalid_argument("metric must be 'l2' or 'ip'");
}

static tinyknn::gpu_algo algo_of(const std::string& a) {
  if (a == "naive") return tinyknn::gpu_algo::naive;
  if (a == "gemm") return tinyknn::gpu_algo::gemm;
  throw std::invalid_argument("algo must be 'naive' or 'gemm'");
}

static tinyknn::topk_algo topk_of(const std::string& a) {
  if (a == "runtime_k") return tinyknn::topk_algo::runtime_k;
  if (a == "compile_time_k") return tinyknn::topk_algo::compile_time_k;
  if (a == "split_rows") return tinyknn::topk_algo::split_rows;
  throw std::invalid_argument("topk must be 'runtime_k', 'compile_time_k' or 'split_rows'");
}

static double ms(clk::time_point a, clk::time_point b) {
  return std::chrono::duration<double, std::milli>(b - a).count();
}

PYBIND11_MODULE(_tinyknn, m) {
  m.doc() = "tinyknn: brute-force nearest-neighbour search, CPU and CUDA";

  m.def("has_cuda", &tinyknn::has_cuda);

  m.def(
      "search",
      [](farray queries, farray db, int k, const std::string& metric, int threads,
         std::size_t db_block, bool fixed_dim, bool release_gil) {
        auto qv = view_of(queries, "queries"), dv = view_of(db, "db");
        const auto Q = qv.rows;
        py::array_t<std::int64_t> ids({Q, std::size_t(k)});
        py::array_t<float> dists({Q, std::size_t(k)});
        auto iv = tinyknn::make_view(ids.mutable_data(), Q, std::size_t(k));
        auto sv = tinyknn::make_view(dists.mutable_data(), Q, std::size_t(k));
        auto run = [&] { tinyknn::search(metric_of(metric), qv, dv, iv, sv, threads, db_block, fixed_dim); };
        if (release_gil) {
          // From here to the end of the scope this thread touches no Python object, so other
          // Python threads may run. The arrays stay alive: this frame holds references to them.
          py::gil_scoped_release unlocked;
          run();
        } else {
          run();
        }
        return py::make_tuple(ids, dists);
      },
      py::arg("queries"), py::arg("db"), py::arg("k") = 10, py::arg("metric") = "l2",
      py::arg("threads") = 1, py::arg("db_block") = 512, py::arg("fixed_dim") = true,
      py::arg("release_gil") = true);

  // The address C++ ends up reading from. Equal to arr.ctypes.data → no copy was made.
  m.def("probe", [](farray a) { return reinterpret_cast<std::uintptr_t>(a.data()); });
  m.def("probe_strict", [](strict a) { return reinterpret_cast<std::uintptr_t>(a.data()); },
        py::arg("a").noconvert());

  // The cheapest possible call, and one that does n floats of work: the price of crossing
  // from Python into C++ once, and how fast it is paid back.
  m.def("noop", [] {});
  m.def("sum", [](farray a) {
    const float* p = a.data();
    const auto n = std::size_t(a.size());
    float acc[16] = {};
    std::size_t i = 0;
    for (; i + 16 <= n; i += 16)
      for (int l = 0; l < 16; ++l) acc[l] += p[i + l];
    for (; i < n; ++i) acc[0] += p[i];
    float s = 0;
    for (float v : acc) s += v;
    return s;
  });

#ifdef TINYKNN_WITH_CUDA
  py::class_<tinyknn::gpu_index>(m, "Index")
      .def(py::init([](farray db, const std::string& metric) {
             return tinyknn::gpu_index(view_of(db, "db"), metric_of(metric));
           }),
           py::arg("db"), py::arg("metric") = "l2")
      .def_property_readonly("size", &tinyknn::gpu_index::size)
      .def_property_readonly("dim", &tinyknn::gpu_index::dim)
      .def(
          "search",
          [](tinyknn::gpu_index& self, farray queries, int k, const std::string& algo,
             bool trace, const std::string& topk) -> py::tuple {
            const auto t_enter = clk::now();
            auto qv = view_of(queries, "queries");
            const auto Q = qv.rows;
            py::array_t<std::int64_t> ids({Q, std::size_t(k)});
            py::array_t<float> dists({Q, std::size_t(k)});
            auto iv = tinyknn::make_view(ids.mutable_data(), Q, std::size_t(k));
            auto sv = tinyknn::make_view(dists.mutable_data(), Q, std::size_t(k));
            const auto t_ready = clk::now();
            tinyknn::trace t;
            {
              py::gil_scoped_release unlocked;
              self.search(qv, iv, sv, algo_of(algo), topk_of(topk), trace ? &t : nullptr);
            }
            const auto t_done = clk::now();
            if (!trace) return py::make_tuple(ids, dists);
            py::dict d;
            d["bind_in"] = ms(t_enter, t_ready);  // views + output arrays allocated
            d["stage_in"] = t.stage_in;
            d["queue"] = t.queue;
            d["h2d"] = t.h2d;
            d["distances"] = t.distances;
            d["topk"] = t.topk;
            d["d2h"] = t.d2h;
            d["stage_out"] = t.stage_out;
            d["library_total"] = t.total;
            d["cpp_total"] = ms(t_enter, t_done);
            return py::make_tuple(ids, dists, d);
          },
          py::arg("queries"), py::arg("k") = 10, py::arg("algo") = "gemm",
          py::arg("trace") = false, py::arg("topk") = "split_rows")
      .def(
          "enqueue",
          [](tinyknn::gpu_index& self, farray queries, int k, const std::string& algo,
             const std::string& topk) {
            self.enqueue(view_of(queries, "queries"), std::size_t(k), algo_of(algo), topk_of(topk));
          },
          py::arg("queries"), py::arg("k") = 10, py::arg("algo") = "gemm",
          py::arg("topk") = "split_rows")
      .def("wait", &tinyknn::gpu_index::wait, py::call_guard<py::gil_scoped_release>())
      // A NumPy array that *views* the index's pinned result buffer — no copy. Passing the
      // index as the array's base keeps the index (and so the buffer) alive as long as the
      // array is: ownership stays with C++, lifetime is extended by Python's reference count.
      .def(
          "peek_ids",
          [](py::object self_obj, std::size_t rows, std::size_t k) {
            auto& self = self_obj.cast<tinyknn::gpu_index&>();
            if (rows > self.chunk()) throw std::invalid_argument("rows > chunk");
            return py::array_t<std::int64_t>({rows, k}, self.result_ids(), self_obj);
          },
          py::arg("rows"), py::arg("k"));
#endif
}

def find_cuda_lib(name):
    """Directory and file name of lib<name>.so.* — the CUDA toolkit first, then pip packages."""
    roots = [f"{CUDA_HOME}/lib64", f"{CUDA_HOME}/targets/x86_64-linux/lib"] if CUDA_HOME else []
    roots += glob.glob(f"{sysconfig.get_paths()['purelib']}/nvidia/*/lib")
    for r in roots:
        hits = sorted(glob.glob(f"{r}/lib{name}.so*"), key=len)
        if hits:
            return r, Path(hits[0]).name
    return None, None


EXT = sysconfig.get_config_var("EXT_SUFFIX")
PYINC = sh(f"{sys.executable} -m pybind11 --includes").stdout.strip()
HAVE_CUDA = False
with experiment("build"):
    cublas_dir, cublas_file = find_cuda_lib("cublas")
    cudart_dir, cudart_file = find_cuda_lib("cudart")
    HAVE_CUDA = bool(NVCC and cublas_dir and cudart_dir)
    cuda_def = "-DTINYKNN_WITH_CUDA" if HAVE_CUDA else ""
    cublas_inc = " ".join(f"-I{p}" for p in glob.glob(f"{sysconfig.get_paths()['purelib']}/nvidia/cublas/include"))
    steps = [("cpu.cpp → cpu.o",
              f"g++ -std=c++20 -O3 -march=native -fPIC {cuda_def} -Itinyknn/include -c tinyknn/src/cpu.cpp -o build/cpu.o",
              "build/cpu.o")]
    if HAVE_CUDA:
        steps.append(("cuda.cu → cuda.o",
                      f"{NVCC} -std=c++17 -O3 -arch=sm_75 -Xptxas -v -Xcompiler -fPIC -Itinyknn/include {cublas_inc} "
                      f"-c tinyknn/src/cuda.cu -o build/cuda.o", "build/cuda.o"))
        steps.append(("link libtinyknn.so",
                      f"g++ -shared build/cpu.o build/cuda.o -L{cudart_dir} -l:{cudart_file} -L{cublas_dir} "
                      f"-l:{cublas_file} -Wl,-rpath,{cudart_dir} -Wl,-rpath,{cublas_dir} -o build/libtinyknn.so",
                      "build/libtinyknn.so"))
    else:
        steps.append(("link libtinyknn.so", "g++ -shared build/cpu.o -o build/libtinyknn.so", "build/libtinyknn.so"))
    steps.append(("bindings.cpp → _tinyknn module",
                  f"g++ -std=c++20 -O3 -fPIC -shared {cuda_def} {PYINC} -Itinyknn/include tinyknn/python/bindings.cpp "
                  f"-Lbuild -ltinyknn -Wl,-rpath,'$ORIGIN' -o build/_tinyknn{EXT}", f"build/_tinyknn{EXT}"))
    rows, kernels = [], []
    for name, cmd, out in steps:
        t0 = time.perf_counter()
        p = sh(cmd)
        rows.append({"step": name, "seconds": round(time.perf_counter() - t0, 2),
                     "output_bytes": Path(out).stat().st_size})
        if "-Xptxas -v" in cmd:
            # ptxas reports, per compiled kernel, its registers and its stack frame: bytes of
            # per-thread memory that live outside the registers, in slow "local" memory.
            for block in p.stderr.split("Compiling entry function")[1:]:
                mangled = block.split("'")[1]
                stack = re.search(r"(\d+) bytes stack frame", block)
                regs = re.search(r"Used (\d+) registers", block)
                kernels.append({"kernel": sh(f"echo {mangled} | c++filt").stdout.strip(),
                                "registers": int(regs.group(1)) if regs else None,
                                "stack_frame_bytes": int(stack.group(1)) if stack else None})
    RESULTS["build"] = {"cuda": HAVE_CUDA, "steps": rows, "kernels": kernels}
    display(pd.DataFrame(rows))
    if kernels:
        kd = pd.DataFrame(kernels)
        kd["kernel"] = kd.kernel.str.replace(r"tinyknn::\(anonymous namespace\)::", "", regex=True).str.slice(0, 90)
        display(kd)

sys.path.insert(0, str(Path("build").resolve()))
import _tinyknn as tk
HAVE_CUDA = HAVE_CUDA and tk.has_cuda()        # built for CUDA, and a GPU is actually there
print("GPU support:", HAVE_CUDA)

def exact_knn(q, db, k, metric="l2"):
    q, db = q.astype(np.float64), db.astype(np.float64)
    d = ((q * q).sum(1)[:, None] + (db * db).sum(1)[None] - 2 * q @ db.T) if metric == "l2" else -(q @ db.T)
    kk = min(k, db.shape[0])
    i = np.argpartition(d, kk - 1, axis=1)[:, :kk]
    order = np.take_along_axis(d, i, 1).argsort(1)
    return np.take_along_axis(i, order, 1)


def recall(ids, ref):
    k = ref.shape[1]
    return float(np.mean([len(set(a[:k]) & set(b)) / k for a, b in zip(ids, ref)]))


N, D, K = 100_000, 128, 10
rng = np.random.default_rng(0)
DB = rng.standard_normal((N, D), dtype=np.float32)
QUERIES = rng.standard_normal((1000, D), dtype=np.float32)

with experiment("tests"):
    rows = []
    for metric in ["l2", "ip"]:
        for dim in [128, 100]:                                  # 128 is specialised, 100 is not
            q, db = QUERIES[:37, :dim], DB[:20_000, :dim]      # 37: not a multiple of 4
            ref = exact_knn(q, db, K, metric)
            ids, dists = tk.search(q, db, K, metric)
            rows.append({"path": "cpu", "metric": metric, "dim": dim, "recall": recall(ids, ref),
                         "sorted": bool((np.diff(dists, axis=1) >= 0).all())})
            if HAVE_CUDA:
                index = tk.Index(db, metric)
                for algo in ["naive", "gemm"]:
                    for sel in ["runtime_k", "compile_time_k", "split_rows"]:
                        ids, dists = index.search(q, K, algo, topk=sel)
                        rows.append({"path": f"gpu {algo} / {sel}", "metric": metric, "dim": dim,
                                     "recall": recall(ids, ref), "sorted": bool((np.diff(dists, axis=1) >= 0).all())})
                # k = 7 has no compiled kernel: it must fall back to the general one and still be right.
                ids7, _ = index.search(q[:5], 7, "gemm", topk="split_rows")
                rows.append({"path": "gpu gemm / k=7 fallback", "metric": metric, "dim": dim,
                             "recall": recall(ids7, exact_knn(q[:5], db, 7, metric)), "sorted": True})
                del index
    # Edge cases: fewer rows than k, a non-contiguous query array, threads.
    ids, _ = tk.search(QUERIES[:3], DB[:5], K)
    edge_short = bool((ids[:, 5:] == -1).all() and (np.sort(ids[:, :5], 1) == np.arange(5)).all())
    strided = QUERIES[:40:2]
    edge_strided = recall(tk.search(strided, DB[:20_000], K)[0], exact_knn(strided, DB[:20_000], K)) == 1.0
    edge_threads = bool((tk.search(QUERIES[:50], DB[:20_000], K, threads=4)[0] ==
                         tk.search(QUERIES[:50], DB[:20_000], K, threads=1)[0]).all())
    RESULTS["tests"] = {"rows": rows, "fewer_rows_than_k": edge_short, "strided_queries": edge_strided,
                        "threads_identical": edge_threads}
    display(pd.DataFrame(rows))
    print("fewer rows than k:", edge_short, "| strided queries:", edge_strided, "| 1 vs 4 threads identical:", edge_threads)

def python_search(q, rows, k):
    best = []
    for i, row in enumerate(rows):
        s = 0.0
        for a, b in zip(q, row):
            d = a - b
            s += d * d
        best.append((s, i))
    best.sort()
    return [i for _, i in best[:k]]


DB_NORMS = (DB * DB).sum(1)                       # computed once, as an index would


def numpy_broadcast(q, k=K):                      # one query: every difference materialised
    d = ((DB - q) ** 2).sum(1)
    i = np.argpartition(d, k)[:k]
    return i[np.argsort(d[i])]


def numpy_gemm(qs, k=K):                          # many queries: one matrix multiplication
    d = (qs * qs).sum(1)[:, None] + DB_NORMS[None] - 2.0 * (qs @ DB.T)
    i = np.argpartition(d, k, axis=1)[:, :k]
    return np.take_along_axis(i, np.take_along_axis(d, i, 1).argsort(1), 1)


with experiment("one_search"):
    rows = []
    db_list, q_list = DB.tolist(), QUERIES[0].tolist()
    t = best_time(lambda: python_search(q_list, db_list, K), reps=1, warmup=0)
    rows.append({"impl": "pure Python", "queries": 1, "ms_total": t * 1e3})
    del db_list

    q1 = QUERIES[:1]
    cases = [("NumPy, broadcast", lambda q: np.stack([numpy_broadcast(x) for x in q])),
             ("NumPy, matrix multiply", numpy_gemm),
             ("tinyknn CPU, 1 thread", lambda q: tk.search(q, DB, K, threads=1)),
             (f"tinyknn CPU, {os.cpu_count()} threads", lambda q: tk.search(q, DB, K, threads=os.cpu_count()))]
    if HAVE_CUDA:
        INDEX = tk.Index(DB)
        cases.append(("tinyknn GPU", lambda q: INDEX.search(q, K)))
    for name, fn in cases:
        for nq in ([1, 16] if "broadcast" in name else [1, 1000]):   # broadcast: ~50 ms a query
            q = QUERIES[:nq]
            t = best_time(lambda: fn(q), reps=5)
            rows.append({"impl": name, "queries": nq, "ms_total": t * 1e3})
    df = pd.DataFrame(rows)
    df["ms_per_query"] = df.ms_total / df.queries
    RESULTS["one_search"] = df.to_dict("records")
    display(df.round(3))

with experiment("python_objects"):
    import ctypes
    out = {"sizeof_float": sys.getsizeof(1.5), "sizeof_int_small": sys.getsizeof(7),
           "sizeof_empty_list": sys.getsizeof([]), "pointer_bytes": ctypes.sizeof(ctypes.c_void_p)}
    n = 1_000_000
    L = [random.random() for _ in range(n)]
    A64 = np.array(L)
    A32 = A64.astype(np.float32)
    out["list_bytes_1m"] = sys.getsizeof(L) + n * sys.getsizeof(1.5)
    out["numpy64_bytes_1m"] = A64.nbytes
    out["numpy32_bytes_1m"] = A32.nbytes

    # Where the objects are. Sixteen floats created one after another; the first sixteen of a
    # million; the first sixteen after the million are shuffled.
    fresh = [float(i) + 0.5 for i in range(16)]
    out["fresh_addresses"] = [hex(id(v)) for v in fresh]
    first = [id(v) for v in L[:16]]
    random.shuffle(L)
    shuffled = [id(v) for v in L[:16]]
    out["list_addresses"] = [hex(a) for a in first]
    out["list_addresses_after_shuffle"] = [hex(a) for a in shuffled]
    out["numpy_addresses"] = [hex(A64.ctypes.data + 8 * i) for i in range(16)]
    out["numpy_strides"] = list(A64.reshape(1000, 1000).strides)

    # Time per element to add up a million numbers, five ways.
    def loop_list():
        s = 0.0
        for v in L:
            s += v
        return s

    def loop_numpy():
        s = 0.0
        for v in A64:                  # each element becomes a new numpy.float64 object
            s += v
        return s

    per = {
        "Python loop over list": best_time(loop_list, reps=3),
        "sum(list)": best_time(lambda: sum(L), reps=3),
        "Python loop over NumPy array": best_time(loop_numpy, reps=3),
        "np.sum (float64)": best_time(lambda: A64.sum(), reps=10),
        "tinyknn.sum (float32, C++)": best_time(lambda: tk.sum(A32), reps=10),
    }
    out["ns_per_element"] = {k: v / n * 1e9 for k, v in per.items()}
    RESULTS["python_objects"] = out
    print(json.dumps({k: v for k, v in out.items() if "addresses" not in k}, indent=1))
    print("fresh floats:  ", out["fresh_addresses"][:5])
    print("list order:    ", out["list_addresses"][:5])
    print("after shuffle: ", out["list_addresses_after_shuffle"][:5])
    print("numpy:         ", out["numpy_addresses"][:5])

%%writefile bench/bench.hpp
// bench/bench.hpp — timing and output helpers shared by the small benchmark programs.
// Every program prints one JSON object on stdout; the notebook parses it.
#pragma once
#include <chrono>
#include <cstdio>
#include <string>
#include <vector>

namespace bench {

inline double now() {
  return std::chrono::duration<double>(std::chrono::steady_clock::now().time_since_epoch()).count();
}

// Fastest of `reps` runs, in seconds. The fastest run is the one least disturbed by everything
// else on the machine, which is the quantity these experiments are after.
template <class F>
double best_of(int reps, F&& f) {
  double best = 1e30;
  for (int r = 0; r < reps; ++r) {
    double t0 = now();
    f();
    double t = now() - t0;
    if (t < best) best = t;
  }
  return best;
}

// Keeps a value "used" so the compiler cannot delete the computation that produced it.
template <class T>
inline void keep(T const& v) {
  asm volatile("" : : "r,m"(v) : "memory");
}

// Minimal JSON writer: fields are appended in order, nesting is written by hand.
struct json {
  std::string s = "{";
  bool first = true;
  json& key(const std::string& k) {
    s += (first ? "\"" : ",\"") + k + "\":";
    first = false;
    return *this;
  }
  json& num(const std::string& k, double v) {
    char b[64];
    std::snprintf(b, sizeof b, "%.15g", v);
    key(k).s += b;
    return *this;
  }
  json& str(const std::string& k, const std::string& v) {
    key(k).s += "\"" + v + "\"";
    return *this;
  }
  json& raw(const std::string& k, const std::string& v) {
    key(k).s += v;
    return *this;
  }
  std::string done() { return s + "}"; }
};

inline std::string array(const std::vector<double>& v) {
  std::string s = "[";
  char b[64];
  for (std::size_t i = 0; i < v.size(); ++i) {
    std::snprintf(b, sizeof b, "%s%.9g", i ? "," : "", v[i]);
    s += b;
  }
  return s + "]";
}

}  // namespace bench