#!/usr/bin/env/python

# TODO

# - hmm what is "normalized" boolean at KNN indexing time? -- COSINE similarity sets this to true

# - try turning diversity off -- faster forceMerge? better recall?

# - why force merge 12X slower

# - why only one thread

# - report net concurrency utilized in the table

# - report total cpu for all indexing threads too

# - hmm how come so much faster to compute exact NN at queryStartIndex=0 than 10000, 20000? 60 sec vs ~470 sec!?

# - not always the first run! sometimes 2nd run is super-fast

import argparse

import itertools

import mmap

import multiprocessing

import os

import random

import re

import shlex

import shutil

import statistics

import struct

import subprocess

import sys

import time

from pathlib import Path

try:

import numpy as np

except ImportError:

print("\nERROR: numpy is required but not installed.\n")

print("To fix, run from the luceneutil root directory:\n")

print(" make env")

print(" source .venv/bin/activate")

print(" python -u src/python/knnPerfTest.py\n")

raise SystemExit(1) from None

import autologger

import benchUtil

import constants

import knnExactNN

import ps_head

from benchUtil import GNUPLOT_PATH, PERF_EXE

from common import getLuceneDirFromGradleProperties

# toggle between 'pread' and 'mmap' for concurrent random vector reads when smelling vectors -- pread is

# maybe a bit faster?

IO_METHOD = "pread"

# IO_METHOD = "mmap"

def advise_will_need(file_name, offset_bytes=0, length_bytes=0):

"""Proactively hint to the OS to load a range of file into RAM, using the configured IO_METHOD."""

if not os.path.exists(file_name):

return

file_size = os.path.getsize(file_name)

if length_bytes <= 0 or offset_bytes + length_bytes > file_size:

length_bytes = file_size - offset_bytes

if length_bytes <= 0:

return

with open(file_name, "rb") as f:

if IO_METHOD == "pread":

os.posix_fadvise(f.fileno(), offset_bytes, length_bytes, os.POSIX_FADV_WILLNEED)

elif IO_METHOD == "mmap":

# map the part of the file we need

mm = mmap.mmap(f.fileno(), length_bytes, offset=offset_bytes, access=mmap.ACCESS_READ)

try:

mm.madvise(mmap.MADV_WILLNEED)

finally:

mm.close()

# see also https://share.google/aimode/IDYCxtTyGhFUwC1pX for clean-ish

# ways to use io_uring-like async io from Python

# Measure vector search recall and latency while exploring hyperparameters

# SETUP:

### Download and extract data files: Wikipedia line docs + GloVe

# python src/python/initial_setup.py -download OR curl -O https://downloads.cs.stanford.edu/nlp/data/glove.6B.zip -k

# cd ../data

# unzip glove.6B.zip

# unlzma enwiki-20120502-lines-1k.txt.lzma OR xz enwiki-20120502-lines-1k.txt.lzma

### Create document and task vectors

# ./gradlew vectors-100

#

# change the parameters below and then run (you can still manually run this file, but using gradle command

# below will auto recompile if you made any changes to java files in luceneutils)

# ./gradlew runKnnPerfTest

#

# for the median result of n runs with the same parameters:

# ./gradlew runKnnPerfTest -Pruns=n

#

# you may want to modify the following settings:

# uses CPUTime sampling (newly available/experimental in Java 25, seems to work on the tasks benchmark)

DO_PROFILING = False

DO_PS = True

DO_VMSTAT = True

# precompute exact NN using numpy (multi-threaded BLAS matmul). when False,

# KnnGraphTester.java computes exact NN itself (slower, single-threaded Java).

USE_NUMPY_EXACT_NN = True

# Set this to True to use perf tool to record instructions executed and confirm SIMD

# instructions were executed

# TODO: how much overhead / perf impact from this? can we always run?

CONFIRM_SIMD_ASM_MODE = False

# perf stat SIMD validation: uses hardware counters to confirm SIMD (SSE/AVX2/AVX-512)

# is actually used during vector operations. negligible overhead -- always on when perf

# is available.

DO_PERF_STAT_SIMD = PERF_EXE is not None

# set this to True to collect all HNSW traversal scores and generate a histogram

DO_HNSW_SCORE_HISTOGRAM = False

# sample 1 in every N HNSW traversal scores when building the histogram (to keep HTML size reasonable)

HNSW_SAMPLE_EVERY_N = 100

# set this to True to compute sampled all query x doc distances and generate a histogram

DO_ALL_DISTANCES_HISTOGRAM = False

# sample 1 in every N all-distances scores (sampling done in Java).

# with 400K docs and 10K queries (4B distances), 1000 -> ~4M samples -> ~40 MB HTML

ALL_DISTANCES_SAMPLE_EVERY_N = 1000

if CONFIRM_SIMD_ASM_MODE and PERF_EXE is None:

raise RuntimeError("CONFIRM_SIMD_ASM_MODE is True but PERF_EXE is not found; install 'perf' tool and rerun?")

# e.g. to compile KnnIndexer:

#

# javac -d build -cp /l/trunk/lucene/core/build/libs/lucene-core-10.0.0-SNAPSHOT.jar:/l/trunk/lucene/join/build/libs/lucene-join-10.0.0-SNAPSHOT.jar src/main/knn/*.java src/main/WikiVectors.java src/main/perf/VectorDictionary.java

#

NOISY = True

# TODO

# - can we expose greediness (global vs local queue exploration in KNN search) here?

# test parameters. This script will run KnnGraphTester on every combination of these parameters

PARAMS = {

# "ndoc": (10_000_000,),

#'ndoc': (10000, 100000, 200000, 500000),

#'ndoc': (10000, 100000, 200000, 500000),

#'ndoc': (2_000_000,),

#'ndoc': (1_000_000,),

"ndoc": (500_000,),

#'ndoc': (50_000,),

"maxConn": (64,),

# "maxConn": (64,),

#'maxConn': (32,),

"beamWidthIndex": (250,),

# "beamWidthIndex": (250,),

#'beamWidthIndex': (50,),

"fanout": (100,),

# "fanout": (50,),

#'quantize': None,

#'quantizeBits': (32, 7, 4),

"numMergeWorker": (24,),

"numMergeThread": (8,),

"numSearchThread": (4,),

#'numMergeWorker': (1,),

#'numMergeThread': (1,),

"encoding": ("float32",),

# "metric": ("cosine",), # default is angular (dot_product)

"metric": ("dot_product",),

# 'metric': ('mip',),

#'quantize': (True,),

"quantizeBits": (32, 8, 7, 4, 2),

# "quantizeBits": (1,),

# "overSample": (5,), # extra ratio of vectors to retrieve, for testing approximate scoring, e.g. quantized indices

#'fanout': (0,),

"topK": (100,),

# "bp": ("false", "true"),

#'quantizeCompress': (True, False),

"quantizeCompress": (True,),

# "indexType": ("flat", "hnsw"), # index type,

# "queryStartIndex": (0, 10000, 20000, 30000, 40000, 50000, 60000, 70000, 80000, 90000, 100000), # seek to this start vector before searching, to sample different vectors

# "queryStartIndex": (0, 200000, 400000, 600000),

"forceMerge": (True,),

"niter": (10000,),

# "filterStrategy": ("query-time-pre-filter", "query-time-post-filter", "index-time-filter"),

# "filterSelectivity": ("0.5", "0.2", "0.1", "0.01",),

# "searchType": ("radius",),

# "resultSimilarity": ("0.8", "0.9",),

# "decay": ("0", "0.5", "0.8",),

}

OUTPUT_HEADERS = [

"recall",

"latency(ms)",

"netCPU",

"avgCpuCount",

"nDoc",

"searchType",

"topK",

"fanout",

"resultSimilarity",

"decay",

"resultCount",

"maxConn",

"beamWidth",

"quantized",

"visited",

"index(s)",

"index_docs/s",

"force_merge(s)",

"num_segments",

"index_size(MB)",

"filterStrategy",

"filterSelectivity",

"overSample",

"vec_disk(MB)",

"vec_RAM(MB)",

"bp-reorder",

"indexType",

]

# TODO: "bp",

def advance(ix, values):

for i in reversed(range(len(ix))):

# scary to rely on dict key enumeration order? but i guess if dict never changes while we do this, it's stable?

param = list(values.keys())[i]

# print("advance " + param)

if type(values[param]) in (list, tuple) and ix[i] == len(values[param]) - 1:

ix[i] = 0

else:

ix[i] += 1

return True

return False

_SPARK_CHARS = " ▁▂▃▄▅▆▇█"

_NUM_SPARK_BINS = 20

_NUM_DIM_SAMPLE_VECS = 2000

_THRESH_CONSTANT_STD = 1e-6

_THRESH_SPARSE_PCT_ZEROS = 0.50

_THRESH_SKEWED_ABS = 1.0

_THRESH_HEAVY_TAILS_KURTOSIS = 3.0

_THRESH_FLAT_KURTOSIS = -1.0

_THRESH_OUTLIER_SPREAD_SIGMA = 3.0

def _sparklines_2row(counts):

"""Returns (top_row_str, bot_row_str) for a two-row histogram using unicode block chars.

Block chars fill from the bottom of the cell, so with 8 levels per row the two rows

connect seamlessly: a partial char in the top row sits at the bottom of its cell,

flush against the full block below it. Total resolution: 16 height levels.

"""

max_count = max(counts) if max(counts) > 0 else 1

top_chars = []

bot_chars = []

for c in counts:

level = round(c / max_count * 16)

if level <= 8:

bot_chars.append(_SPARK_CHARS[level])

top_chars.append(" ")

else:

bot_chars.append("█")

top_chars.append(_SPARK_CHARS[level - 8])

return "".join(top_chars), "".join(bot_chars)

def _print_dim_line(d, mean, std, pct_zeros, counts, labels, dim_idx_width):

"""Prints sparkline + stats, with any smell labels (with details) on separate lines below."""

top_row, bot_row = _sparklines_2row(counts)

prefix = f" dim {d:0{dim_idx_width}d} μ={mean:+8.3f} σ={std:8.3f} zeros={pct_zeros * 100:3.0f}% " # noqa: RUF001 sigma (std deviation) is intentional

print(f"{' ' * len(prefix)}[{top_row}]")

print(f"{prefix}[{bot_row}]")

indent = f" {' ' * dim_idx_width} "

for name, detail in labels:

print(f"{indent}-> {name}: {detail}")

print()

def _read_vectors_pread(file_name, sample_indices, vec_size_bytes):

"""Generator that issues concurrent readahead hints and yields vectors via pread."""

with open(file_name, "rb") as f:

fd = f.fileno()

# hint random access for the whole file, to suppress wasteful readahead

os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_RANDOM)

# concurrently send all requests to the OS as hints

for vec_idx in sample_indices:

os.posix_fadvise(fd, vec_idx * vec_size_bytes, vec_size_bytes, os.POSIX_FADV_WILLNEED)

# yield vectors; they should be pre-fetched by the kernel

for vec_idx in sample_indices:

yield vec_idx, np.frombuffer(os.pread(fd, vec_size_bytes, vec_idx * vec_size_bytes), dtype="<f4")

# hint sequential access for the subsequent (sequential) indexing test

os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_SEQUENTIAL)

def _read_vectors_mmap(file_name, sample_indices, vec_size_bytes, dim):

"""Generator that issues concurrent readahead hints and yields vectors via mmap."""

with open(file_name, "rb") as f:

# mmap.PAGESIZE is typically 4096 on Linux

pagesize = mmap.PAGESIZE

# map the entire file

mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)

# hint random access to the whole mmap, to suppress wasteful readahead

mm.madvise(mmap.MADV_RANDOM)

try:

# concurrently send all requests to the OS as hints

for vec_idx in sample_indices:

offset = vec_idx * vec_size_bytes

page_offset = (offset // pagesize) * pagesize

length = offset + vec_size_bytes - page_offset

mm.madvise(mmap.MADV_WILLNEED, page_offset, length)

# yield vectors; they should be pre-fetched by the kernel

for vec_idx in sample_indices:

offset = vec_idx * vec_size_bytes

yield vec_idx, np.frombuffer(mm, dtype="<f4", count=dim, offset=offset)

# hint sequential access for the subsequent (sequential) indexing test

mm.madvise(mmap.MADV_SEQUENTIAL)

finally:

mm.close()

print()

def _check_dim_distributions(dim, file_name, num_vectors, vec_size_bytes):

"""Samples vectors and computes per-dim statistics to detect degenerate dimensions."""

if num_vectors == 0:

return

num_sample = min(_NUM_DIM_SAMPLE_VECS, num_vectors)

sample_indices = random.sample(range(num_vectors), num_sample)

if NOISY:

print(f"smell: sampling {num_sample} of {num_vectors} vectors for per-dim distribution...")

# load all sampled vectors concurrently into a (num_sample, dim) float32 array

samples = np.empty((num_sample, dim), dtype=np.float32)

t0_sec = time.monotonic()

if IO_METHOD == "pread":

reader = _read_vectors_pread(file_name, sample_indices, vec_size_bytes)

elif IO_METHOD == "mmap":

reader = _read_vectors_mmap(file_name, sample_indices, vec_size_bytes, dim)

else:

raise ValueError(f'unknown IO_METHOD "{IO_METHOD}"')

not_norm_count = 0

for i, (vec_idx, vec) in enumerate(reader):

samples[i] = vec

# CPU work: check norm immediately as vector arrives

norm = np.linalg.norm(vec)

if not np.isclose(norm, 1.0, rtol=0.0001, atol=0.0001):

# print warning on new line so it doesn't get overwritten by next progress \r

print(f'\nWARNING: vec {vec_idx} in "{file_name}" has norm={norm} (not normalized)')

not_norm_count += 1

completed = i + 1

if completed % max(1, num_sample // 20) == 0 or completed == num_sample:

elapsed_sec = time.monotonic() - t0_sec

if NOISY:

print(f"\rsmell: {completed}/{num_sample} ({100 * completed / num_sample:3.0f}%) {elapsed_sec:.1f}s", end="", flush=True)

if NOISY:

print()

if not_norm_count > 0:

print(f'WARNING: dimension or vector file name might be wrong? {not_norm_count} of {num_sample} randomly checked vectors are not normalized in "{file_name}"')

# per-dim stats, all vectorized over axis=0, result shape: (dim,)

mean = samples.mean(axis=0)

std = samples.std(axis=0)

pct_zeros = (samples == 0.0).mean(axis=0)

centered = samples - mean

m3 = (centered**3).mean(axis=0)

m4 = (centered**4).mean(axis=0)

with np.errstate(invalid="ignore", divide="ignore"):

skewness = m3 / std**3

excess_kurtosis = m4 / std**4 - 3.0

# zero out stats that are undefined for constant dims

non_const = std > _THRESH_CONSTANT_STD

skewness = np.where(non_const, skewness, 0.0)

excess_kurtosis = np.where(non_const, excess_kurtosis, 0.0)

# OUTLIER_SPREAD: flag dims whose std is an outlier across all dims' stds

mean_of_stds = std.mean()

std_of_stds = std.std()

# use global range so histogram widths are directly comparable across dims:

# a dim with larger sigma spans more bins visually, matching the printed sigma

global_min = float(samples.min())

global_max = float(samples.max())

# per-dim histogram (loop is over dims not over samples, so it's cheap)

dim_counts = []

for d in range(dim):

col = samples[:, d]

if col.min() == col.max():

counts = [0] * _NUM_SPARK_BINS

counts[_NUM_SPARK_BINS // 2] = num_sample

else:

counts = np.histogram(col, bins=_NUM_SPARK_BINS, range=(global_min, global_max))[0].tolist()

dim_counts.append(counts)

# assign labels per dim -- each label is (name, detail_string)

dim_labels = []

for d in range(dim):

labels = []

if std[d] <= _THRESH_CONSTANT_STD:

labels.append(("CONSTANT", f"std={std[d]:.6f}, threshold={_THRESH_CONSTANT_STD}"))

if pct_zeros[d] > _THRESH_SPARSE_PCT_ZEROS:

labels.append(("SPARSE", f"{pct_zeros[d] * 100:.1f}% zeros, threshold={_THRESH_SPARSE_PCT_ZEROS * 100:.0f}%"))

if non_const[d]:

if abs(skewness[d]) > _THRESH_SKEWED_ABS:

labels.append(("SKEWED", f"skew={skewness[d]:+.2f}, threshold=|{_THRESH_SKEWED_ABS}|"))

if excess_kurtosis[d] > _THRESH_HEAVY_TAILS_KURTOSIS:

labels.append(("HEAVY_TAILS", f"kurtosis={excess_kurtosis[d]:+.2f}, threshold>{_THRESH_HEAVY_TAILS_KURTOSIS}"))

if excess_kurtosis[d] < _THRESH_FLAT_KURTOSIS:

labels.append(("FLAT", f"kurtosis={excess_kurtosis[d]:+.2f}, threshold<{_THRESH_FLAT_KURTOSIS}"))

if std_of_stds > 0 and abs(std[d] - mean_of_stds) > _THRESH_OUTLIER_SPREAD_SIGMA * std_of_stds:

z = (std[d] - mean_of_stds) / std_of_stds

labels.append(("OUTLIER_SPREAD", f"this_std={std[d]:.4f}, mean_std={mean_of_stds:.4f}, {z:+.1f}sigma vs threshold={_THRESH_OUTLIER_SPREAD_SIGMA}sigma"))

dim_labels.append(labels)

# format and print output

dim_idx_width = len(str(dim - 1))

bad_dims = [d for d in range(dim) if len(dim_labels[d]) > 0]

elapsed_sec = time.monotonic() - t0_sec

if bad_dims:

print(f"smell: {len(bad_dims)} degenerate dim(s) found in {elapsed_sec:.1f}s:")

if any(name == "OUTLIER_SPREAD" for lbs in dim_labels for name, _ in lbs):

print(f" (OUTLIER_SPREAD: std of all dims: μ={mean_of_stds:.3f}, σ={std_of_stds:.3f}, threshold={_THRESH_OUTLIER_SPREAD_SIGMA}σ)") # noqa: RUF001 sigma (std deviation) is intentional

for d in bad_dims:

_print_dim_line(d, float(mean[d]), float(std[d]), float(pct_zeros[d]), dim_counts[d], dim_labels[d], dim_idx_width)

if NOISY:

print(f"smell: all {dim} dims:")

for d in range(dim):

_print_dim_line(d, float(mean[d]), float(std[d]), float(pct_zeros[d]), dim_counts[d], dim_labels[d], dim_idx_width)

elif NOISY:

print(f"smell: no degenerate dims found in {elapsed_sec:.1f}s")

def smell_vectors(dim, file_name):

"""Runs some simple sanity checks on the vector source file, because we don't store any

self-describing metadata in the .vec source file.

"""

size_bytes = os.path.getsize(file_name)

vec_size_bytes = dim * 4

# cool, i didn't know about divmod!

num_vectors, leftover = divmod(size_bytes, vec_size_bytes)

if leftover != 0:

raise RuntimeError(

f'vector file "{file_name}" cannot be dimension {dim}: its size is not a multiple of each vector\'s size in bytes ({vec_size_bytes}); wrong vector source file or dimensionality?'

)

if NOISY:

print(f"smell vectors from {file_name}")

_check_dim_distributions(dim, file_name, num_vectors, vec_size_bytes)

# all candidate perf stat counters for SIMD validation -- probed at import time

# to discover which ones this CPU actually supports.

#

# two families:

# fp_arith_inst_retired.* -- counts float SIMD instructions (float32 vectors)

# int_vec_retired.* -- counts integer SIMD instructions (quantized vectors)

# available on Ice Lake+ (different naming on older CPUs)

_FP_SIMD_CANDIDATE_COUNTERS = (

"fp_arith_inst_retired.scalar_single",

"fp_arith_inst_retired.128b_packed_single",

"fp_arith_inst_retired.256b_packed_single",

"fp_arith_inst_retired.512b_packed_single",

)

_INT_SIMD_CANDIDATE_COUNTERS = (

"int_vec_retired.128bit",

"int_vec_retired.256bit",

"int_vec_retired.512bit",

)

# label, counter name, weight (proportional to SIMD width so FP and INT are comparable)

_FP_SIMD_LEVEL_DEFS = (

("FP-scalar", "fp_arith_inst_retired.scalar_single", 1),

("FP-SSE", "fp_arith_inst_retired.128b_packed_single", 4),

("FP-AVX2", "fp_arith_inst_retired.256b_packed_single", 8),

("FP-AVX512", "fp_arith_inst_retired.512b_packed_single", 16),

)

_INT_SIMD_LEVEL_DEFS = (

("INT-SSE", "int_vec_retired.128bit", 4),

("INT-AVX2", "int_vec_retired.256bit", 8),

("INT-AVX512", "int_vec_retired.512bit", 16),

)

def _probe_perf_counter(counter):

"""Return True if perf recognizes this counter on the current CPU."""

try:

result = subprocess.run(

[PERF_EXE, "stat", "-e", counter, "--", "true"],

capture_output=True,

text=True,

timeout=5,

check=False,

)

# perf exits 0 and prints the counter (possibly "<not counted>" for short

# commands) when the counter is valid. it exits non-zero or prints

# "<not supported>" when the counter doesn't exist on this CPU.

return result.returncode == 0 and "<not supported>" not in result.stderr

except (subprocess.TimeoutExpired, OSError):

return False

def _probe_simd_perf_counters():

"""Probe which SIMD perf counters this CPU supports.

Probes each counter individually because perf refuses to run at all if any

counter name is unrecognized (e.g. 512b on a non-AVX-512 CPU).

Returns (available_counters, fp_levels, int_levels).

"""

if PERF_EXE is None:

return (), (), ()

available = []

for counter in _FP_SIMD_CANDIDATE_COUNTERS + _INT_SIMD_CANDIDATE_COUNTERS:

if _probe_perf_counter(counter):

available.append(counter)

available_set = set(available)

counters = tuple(available)

fp_levels = tuple(t for t in _FP_SIMD_LEVEL_DEFS if t[1] in available_set)

int_levels = tuple(t for t in _INT_SIMD_LEVEL_DEFS if t[1] in available_set)

return counters, fp_levels, int_levels

# probe once at import time

_AVAILABLE_SIMD_COUNTERS, FP_SIMD_LEVELS, INT_SIMD_LEVELS = _probe_simd_perf_counters()

if _AVAILABLE_SIMD_COUNTERS:

fp_names = [c for c in _AVAILABLE_SIMD_COUNTERS if c.startswith("fp_")]

int_names = [c for c in _AVAILABLE_SIMD_COUNTERS if c.startswith("int_")]

parts = []

if fp_names:

parts.append(f"FP: {', '.join(fp_names)}")

if int_names:

parts.append(f"INT: {', '.join(int_names)}")

print(f"NOTE: perf SIMD counters available: {'; '.join(parts)}")

if not int_names:

print("NOTE: integer SIMD counters (int_vec_retired.*) not available on this CPU; quantized runs will only show FP counters")

elif DO_PERF_STAT_SIMD:

print("WARNING: perf SIMD counters not available on this CPU; disabling SIMD validation")

DO_PERF_STAT_SIMD = False

def wrap_cmd_with_perf_stat_simd(cmd, output_file):

"""Prepend perf stat with SIMD counters to a command, writing stats to output_file."""

return [

PERF_EXE,

"stat",

"-o",

output_file,

"-e",

",".join(_AVAILABLE_SIMD_COUNTERS),

"--",

] + cmd

def parse_perf_stat_file(path):

"""Parse a perf stat output file and return dict of counter_name -> count (int).

Returns empty dict if file is missing or unparseable.

"""

result = {}

try:

text = Path(path).read_text()

except OSError:

return result

for line in text.splitlines():

# format: " 12,345,678 cpu_core/fp_arith_inst_retired.256b_packed_single/ (88.34%)"

# or: " 12345678 fp_arith_inst_retired.256b_packed_single"

# or: " 12345678 int_vec_retired.256bit"

m = re.match(r"^\s+([\d,]+)\s+(?:\S+/)?((?:fp_arith_inst_retired|int_vec_retired)\.\S+?)(?:/|\s)", line)

if m:

count = int(m.group(1).replace(",", ""))

counter = m.group(2)

result[counter] = count

return result

def _summarize_levels(counters, level_defs):

"""Compute per-level breakdown for a set of SIMD level defs.

Returns (parts_list, total_ops, dominant_label) where parts_list has

(label, count, ops) tuples for levels with count > 0.

"""

parts = []

total_ops = 0

dominant_label = None

dominant_ops = 0

for label, counter, ops_per_insn in level_defs:

count = counters.get(counter, 0)

ops = count * ops_per_insn

total_ops += ops

if count > 0:

parts.append((label, count, ops))

if ops > dominant_ops:

dominant_ops = ops

dominant_label = label

return parts, total_ops, dominant_label

def format_simd_report(counters, is_quantized=False):

"""Format a one-line SIMD usage summary from parsed perf stat counters.

Combines FP (fp_arith_inst_retired) and integer (int_vec_retired) counters

into a single flat percentage breakdown weighted by SIMD width, so all

percentages sum to 100%.

Returns (report_string, dominant_level) where dominant_level is

"FP-scalar", "FP-SSE", "FP-AVX2", "FP-AVX512",

"INT-SSE", "INT-AVX2", "INT-AVX512", or "none".

"""

if not counters:

return "SIMD: no perf data", "none"

all_levels = list(FP_SIMD_LEVELS) + list(INT_SIMD_LEVELS)

all_parts, grand_total, dominant = _summarize_levels(counters, all_levels)

if grand_total == 0:

if is_quantized and not INT_SIMD_LEVELS:

return "SIMD: no FP instructions (expected for quantized); integer SIMD counters not available on this CPU", "none"

return "SIMD: no instructions detected", "none"

pct_parts = []

for label, _count, ops in all_parts:

pct_parts.append(f"{label} {100.0 * ops / grand_total:.1f}%")

report = f"SIMD: {' | '.join(pct_parts)} (dominant: {dominant})"

# warn only when the relevant family shows no SIMD

if dominant == "FP-scalar" and not is_quantized:

report += " WARNING: no SIMD vectorization detected!"

elif is_quantized and not any(label.startswith("INT-") for label, _, _ in all_parts):

if not INT_SIMD_LEVELS:

report += " (int_vec_retired counters not available on this CPU)"

else:

report += " WARNING: no integer SIMD detected!"

return report, dominant

def get_unique_log_name(log_path, sub_tool):

log_dir_name, log_base_name = log_path

upto = 0

while True:

log_file_name = f"{log_dir_name}/{log_base_name}-{sub_tool}"

if upto > 0:

log_file_name += f"-{upto}"

log_file_name += ".log"

if not os.path.exists(log_file_name):

return log_file_name

upto += 1

def print_run_summary(values):

options = []

fixed = []

combos = 1

# print these important params first, in this order:

print_order = ["forceMerge", "ndoc", "niter", "topK", "quantizeBits"]

key_to_ord = {}

other_keys = []

max_key_len = 0

for key, value in values.items():

max_key_len = max(max_key_len, len(key))

try:

key_ord = print_order.index(key)

except ValueError:

other_keys.append(key)

for key_ord, key in enumerate(print_order):

key_to_ord[key] = key_ord

other_keys.sort()

for key_ord, key in enumerate(other_keys):

key_to_ord[key] = len(print_order) + key_ord

ord_to_key = [-1] * len(key_to_ord)

for key, key_ord in key_to_ord.items():

ord_to_key[key_ord] = key

for key in ord_to_key:

value = values[key]

if len(value) > 1:

# yay, it turns out you can nest {...} in f-strings!

options.append(f" {key:<{max_key_len}s}: {','.join(str(v) for v in value)}")

combos *= len(value)

else:

# yay, it turns out you can nest {...} in f-strings!

fixed.append(f" {key:<{max_key_len}s}: {value[0]}")

if len(options) == 0:

assert combos == 1

print("NOTE: will run single warmup+test with these params:")

else:

print(f"NOTE: will run {combos} total warmups+tests with all combinations of:")

for s in options:

print(s)

for s in fixed:

print(s)

METRIC_LABELS = {

"dot_product": ("dot_product similarity", "higher --->"),

"angular": ("dot_product similarity", "higher --->"),

"cosine": ("cosine similarity", "higher --->"),

"euclidean": ("euclidean distance", "<--- lower"),

"mip": ("max inner product", "higher --->"),

}

def generate_exact_nn_histogram(scores_path, output_dir, log_base_name, metric=None):

"""Read the binary exact NN scores file and generate an HTML histogram using Google Charts."""

if not os.path.exists(scores_path):

print(f"WARNING: exact NN scores file not found: {scores_path}")

return

file_size = os.path.getsize(scores_path)

num_floats = file_size // 4

if num_floats == 0:

print("WARNING: exact NN scores file is empty")

return

all_scores = struct.unpack(f"<{num_floats}f", Path(scores_path).read_bytes())

metric_name, direction = METRIC_LABELS.get(metric or "", ("similarity score", ""))

min_score = min(all_scores)

max_score = max(all_scores)

if min_score == max_score:

print(f"WARNING: all exact NN scores are identical ({min_score}); skipping histogram")

return

# Sort scores so JS can use binary search for fast range queries

sorted_scores = sorted(all_scores)

# Emit as a compact JSON array (6 decimal places keeps file reasonable)

scores_js_lines = []

chunk_size = 500

for i in range(0, len(sorted_scores), chunk_size):

chunk = sorted_scores[i : i + chunk_size]

scores_js_lines.append(",".join(f"{s:.6f}" for s in chunk))

scores_js = ",\n".join(scores_js_lines)

html = f"""<!DOCTYPE html>

<html>

<head>

<script type="text/javascript" src="https://www.google.com/jsapi"></script>

<script type="text/javascript">

google.load("visualization", "1", {{packages:["corechart"]}});

google.setOnLoadCallback(init);

// sorted scores for fast range queries via binary search

var scores = [

{scores_js}

];

var globalMin = {min_score:.6f};

var globalMax = {max_score:.6f};

var curMin = globalMin;

var curMax = globalMax;

var chart, chartDiv;

var zoomStack = [];

function lowerBound(arr, val) {{

var lo = 0, hi = arr.length;

while (lo < hi) {{

var mid = (lo + hi) >> 1;

if (arr[mid] < val) lo = mid + 1; else hi = mid;

}}

return lo;

}}

function buildHistogram(lo, hi, numBins) {{

var range = hi - lo;

var binWidth = range / numBins;

var bins = new Array(numBins).fill(0);

var startIdx = lowerBound(scores, lo);

var endIdx = lowerBound(scores, hi);

var count = 0;

for (var i = startIdx; i < scores.length && scores[i] <= hi; i++) {{

var idx = Math.floor((scores[i] - lo) / binWidth);

if (idx >= numBins) idx = numBins - 1;

if (idx < 0) idx = 0;

bins[idx]++;

count++;

}}

return {{bins: bins, binWidth: binWidth, count: count}};

}}

function drawChart(lo, hi) {{

var numBins = 50;

var h = buildHistogram(lo, hi, numBins);

var rows = [['{metric_name} ({direction})', 'Count']];

for (var i = 0; i < numBins; i++) {{

var center = lo + (i + 0.5) * h.binWidth;

rows.push([center, h.bins[i]]);

}}

var data = google.visualization.arrayToDataTable(rows);

var zoomLabel = (lo === globalMin && hi === globalMax) ? '' : ' [zoomed]';

var options = {{

title: 'Exact NN {metric_name} distribution (' + h.count + ' of {num_floats} scores)' + zoomLabel,

legend: {{position: 'none'}},

hAxis: {{

title: '{metric_name} ({direction})',

viewWindow: {{min: lo, max: hi}}

}},

vAxis: {{title: 'Count'}},

bar: {{groupWidth: '95%'}},

chartArea: {{left: 80, right: 20, top: 40, bottom: 60}}

}};

chart.draw(data, options);

document.getElementById('info').innerHTML =

'Range: ' + lo.toFixed(6) + ' .. ' + hi.toFixed(6) +

' &nbsp; Bin width: ' + h.binWidth.toFixed(6) +

' &nbsp; Scores in view: ' + h.count;

}}

function init() {{

chartDiv = document.getElementById('chart_div');

chart = new google.visualization.ColumnChart(chartDiv);

drawChart(globalMin, globalMax);

// drag-to-zoom

var dragStart = null;

var overlay = document.getElementById('overlay');

chartDiv.addEventListener('mousedown', function(e) {{

var rect = chartDiv.getBoundingClientRect();

dragStart = {{x: e.clientX - rect.left, clientX: e.clientX}};

overlay.style.left = dragStart.x + 'px';

overlay.style.width = '0px';

overlay.style.display = 'block';

}});

chartDiv.addEventListener('mousemove', function(e) {{

if (!dragStart) return;

var rect = chartDiv.getBoundingClientRect();

var curX = e.clientX - rect.left;

var left = Math.min(dragStart.x, curX);

var width = Math.abs(curX - dragStart.x);

overlay.style.left = left + 'px';

overlay.style.width = width + 'px';

}});

chartDiv.addEventListener('mouseup', function(e) {{

if (!dragStart) return;

overlay.style.display = 'none';

var rect = chartDiv.getBoundingClientRect();

var endX = e.clientX - rect.left;

var x0 = Math.min(dragStart.x, endX);

var x1 = Math.max(dragStart.x, endX);

dragStart = null;

// need at least 5px drag to count as zoom

if (x1 - x0 < 5) return;

var cli = chart.getChartLayoutInterface();

var val0 = cli.getHAxisValue(x0);

var val1 = cli.getHAxisValue(x1);

if (val0 === null || val1 === null) return;

var newMin = Math.max(Math.min(val0, val1), globalMin);

var newMax = Math.min(Math.max(val0, val1), globalMax);

if (newMax - newMin < (globalMax - globalMin) * 0.001) return;

zoomStack.push({{min: curMin, max: curMax}});

curMin = newMin;

curMax = newMax;

drawChart(curMin, curMax);

document.getElementById('resetBtn').style.display = 'inline';

document.getElementById('backBtn').style.display = 'inline';

}});

document.getElementById('resetBtn').addEventListener('click', function() {{

zoomStack = [];

curMin = globalMin;

curMax = globalMax;

drawChart(curMin, curMax);

this.style.display = 'none';

document.getElementById('backBtn').style.display = 'none';

}});

document.getElementById('backBtn').addEventListener('click', function() {{

if (zoomStack.length === 0) return;

var prev = zoomStack.pop();

curMin = prev.min;

curMax = prev.max;

drawChart(curMin, curMax);

if (zoomStack.length === 0) {{

document.getElementById('resetBtn').style.display = 'none';

this.style.display = 'none';

}}

}});

}}

</script>

<style>

#chart_container {{ position: relative; width: 1200px; height: 600px; }}

#chart_div {{ width: 100%; height: 100%; }}

#overlay {{ position: absolute; top: 0; height: 100%; background: rgba(66,133,244,0.15);

border-left: 1px solid rgba(66,133,244,0.5); border-right: 1px solid rgba(66,133,244,0.5);

display: none; pointer-events: none; z-index: 10; }}

button {{ margin: 4px 4px 4px 0; padding: 4px 12px; }}

</style>

</head>

<body>

<div id="chart_container">

<div id="chart_div"></div>

<div id="overlay"></div>

</div>

<button id="backBtn" style="display:none">Back</button>

<button id="resetBtn" style="display:none">Reset zoom</button>

<p id="info"></p>

<p style="color:#888">Click and drag on the chart to zoom in. Source: {scores_path}</p>

</body>

</html>

"""

output_file = f"{output_dir}/{log_base_name}-knnDistanceHistogram.html"

Path(output_file).write_text(html)

print(f"Wrote exact NN distance histogram to {output_file}")

def generate_all_distances_histogram(scores_path, output_dir, log_base_name, metric=None, sample_every_n=None):

"""Read the sampled all-distances binary file and generate an HTML histogram.

The binary format is raw little-endian float32 scores (no per-query structure).

"""

if not os.path.exists(scores_path):

print(f"WARNING: all-distances scores file not found: {scores_path}")

return

file_size = os.path.getsize(scores_path)

num_floats = file_size // 4

if num_floats == 0:

print("WARNING: all-distances scores file is empty")

return

all_scores = struct.unpack(f"<{num_floats}f", Path(scores_path).read_bytes())

sample_label = ""

if sample_every_n is not None:

sample_label = f" (sampled 1-in-{sample_every_n})"

metric_name, direction = METRIC_LABELS.get(metric or "", ("similarity score", ""))

min_score = min(all_scores)

max_score = max(all_scores)

if min_score == max_score:

print(f"WARNING: all distances scores are identical ({min_score}); skipping histogram")

return

sorted_scores = sorted(all_scores)

scores_js_lines = []

chunk_size = 500

for i in range(0, len(sorted_scores), chunk_size):

chunk = sorted_scores[i : i + chunk_size]

scores_js_lines.append(",".join(f"{s:.6f}" for s in chunk))

scores_js = ",\n".join(scores_js_lines)

html = f"""<!DOCTYPE html>

<html>