Kernels

Packed XNOR/popcount GEMM with runtime SIMD dispatch. See Portable SIMD kernel for the design.

Runtime dispatch

kernel_name

kernel_name() -> str

Name of the SIMD path the native kernel selected at runtime.

"numpy" when no native library is loaded, "unknown" for a native library built before runtime dispatch existed.

Source code in bnn/kernels/packed.py

def kernel_name() -> str:
    """Name of the SIMD path the native kernel selected at runtime.

    ``"numpy"`` when no native library is loaded, ``"unknown"`` for a native
    library built before runtime dispatch existed.
    """
    lib = _try_load_native()
    if not lib:
        return "numpy"
    if not hasattr(lib, "binary_gemm_kernel_id"):
        return "unknown"
    return _KERNEL_NAMES.get(int(lib.binary_gemm_kernel_id()), "unknown")

available_kernels

available_kernels() -> list[str]

Kernel paths usable on this machine, slowest first (always ≥ scalar).

Source code in bnn/kernels/packed.py

def available_kernels() -> list[str]:
    """Kernel paths usable on this machine, slowest first (always ≥ scalar)."""
    feats = cpu_features()
    out = ["scalar"]
    if feats["avx2"]:
        out.append("avx2")
    if feats["avx512_vpopcntdq"]:
        out.append("avx512")
    if feats["neon"]:
        out.append("neon")
    return out

cpu_features

cpu_features() -> dict[str, bool]

Which accelerated paths this CPU can actually run.

Source code in bnn/kernels/packed.py

def cpu_features() -> dict[str, bool]:
    """Which accelerated paths this CPU can actually run."""
    lib = _try_load_native()
    if not lib or not hasattr(lib, "binary_gemm_cpu_features"):
        return {"avx2": False, "avx512_vpopcntdq": False, "neon": False}
    bits = int(lib.binary_gemm_cpu_features())
    return {
        "avx2": bool(bits & 1),
        "avx512_vpopcntdq": bool(bits & 2),
        "neon": bool(bits & 4),
    }

set_kernel

set_kernel(name: str | None) -> str

Force a kernel path (None re-runs auto-detection).

Falls back to scalar if the requested path is unsupported here. Returns the path actually in effect. Intended for validation and reproducibility — every path must produce identical results.

Source code in bnn/kernels/packed.py

def set_kernel(name: str | None) -> str:
    """Force a kernel path (``None`` re-runs auto-detection).

    Falls back to ``scalar`` if the requested path is unsupported here. Returns
    the path actually in effect. Intended for validation and reproducibility —
    every path must produce identical results.
    """
    lib = _try_load_native()
    if not lib or not hasattr(lib, "binary_gemm_set_kernel"):
        return kernel_name()
    if name is None:
        lib.binary_gemm_set_kernel(-1)
        return kernel_name()
    ids = {v: k for k, v in _KERNEL_NAMES.items()}
    if name not in ids:
        raise ValueError(f"unknown kernel {name!r}; choose from {sorted(ids)} or None")
    lib.binary_gemm_set_kernel(ids[name])
    return kernel_name()

Packing

pack_binary_pm1

pack_binary_pm1(
    x: ndarray, axis: int = -1
) -> tuple[np.ndarray, int]

Pack ±1 values along axis into uint64 words. bit1 => -1/non-positive.

Uses NumPy packbits (little bit-order) — much faster than per-bit multiply-sum.

Source code in bnn/kernels/packed.py

def pack_binary_pm1(x: np.ndarray, axis: int = -1) -> tuple[np.ndarray, int]:
    """Pack ±1 values along `axis` into uint64 words. bit1 => -1/non-positive.

    Uses NumPy ``packbits`` (little bit-order) — much faster than per-bit multiply-sum.
    """
    x = np.ascontiguousarray(np.asarray(x))
    if x.size == 0:
        raise ValueError("pack_binary_pm1: empty array")
    if not np.issubdtype(x.dtype, np.floating) and not np.issubdtype(x.dtype, np.integer):
        raise TypeError(f"pack_binary_pm1: expected numeric dtype, got {x.dtype}")
    x = np.moveaxis(x, axis, -1)
    n = int(x.shape[-1])
    return pack_bits_u64(np.less_equal(x, 0)), n

theoretical_ops

theoretical_ops(m: int, n: int, k: int) -> dict

Source code in bnn/kernels/packed.py

def theoretical_ops(m: int, n: int, k: int) -> dict:
    fp32_macs = m * k * n
    binary_word_ops = m * k * math.ceil(n / 64)
    return {
        "fp32_macs": fp32_macs,
        "binary_word_xnor_popcount": binary_word_ops,
        "theoretical_word_reduction": fp32_macs / max(binary_word_ops, 1),
        "weight_bytes_fp32": k * n * 4,
        "weight_bytes_binary": k * math.ceil(n / 8),
        "weight_compression": (k * n * 4) / max(k * math.ceil(n / 8), 1),
    }

GEMM

binary_gemm_packed

binary_gemm_packed(
    x_pm1: ndarray,
    w_pm1: ndarray,
    *,
    prepacked_w: tuple[ndarray, int] | None = None,
) -> np.ndarray

Compute Y = X @ W.T for ±1 matrices using packed XNOR-popcount.

Native SIMD when the library loads. Otherwise packed NumPy for small batch and dequant+BLAS at/above :func:numpy_packed_blas_crossover_batch so the no-native path is never 5–11× slower than FP32 at B=64 (docs/45 P1).

Source code in bnn/kernels/packed.py

def binary_gemm_packed(
    x_pm1: np.ndarray,
    w_pm1: np.ndarray,
    *,
    prepacked_w: tuple[np.ndarray, int] | None = None,
) -> np.ndarray:
    """Compute Y = X @ W.T for ±1 matrices using packed XNOR-popcount.

    Native SIMD when the library loads. Otherwise packed NumPy for small batch
    and dequant+BLAS at/above :func:`numpy_packed_blas_crossover_batch` so the
    no-native path is never 5–11× slower than FP32 at B=64 (docs/45 P1).
    """
    x_pm1 = np.asarray(x_pm1)
    if x_pm1.ndim != 2:
        raise ValueError(f"x_pm1 must be 2D, got shape {x_pm1.shape}")
    n_feat = int(x_pm1.shape[1])
    batch = int(x_pm1.shape[0])

    w_arr: np.ndarray | None
    wp: np.ndarray | None
    if prepacked_w is None:
        w_arr = np.asarray(w_pm1)
        if w_arr.ndim != 2:
            raise ValueError(f"w_pm1 must be 2D, got shape {w_arr.shape}")
        if w_arr.shape[1] != n_feat:
            raise ValueError(
                f"in_features mismatch: x {n_feat} vs w {w_arr.shape[1]}"
            )
        wp = None
    else:
        wp, n2 = prepacked_w
        if n_feat != n2:
            raise ValueError(f"packed n mismatch: {n_feat} vs {n2}")
        w_arr = None

    # Skip packing when native is absent and BLAS wins. Inputs are ±1 (same
    # contract as fp32_gemm); do not copy through _as_pm1 — a 4096×4096 where()
    # is ~3× the GEMM. Non-±1 values are used as-is here; the packed path still
    # signs via pack_binary_pm1.
    if _try_load_native() is None and prefer_numpy_blas_fallback(batch):
        if w_arr is not None:
            return fp32_gemm(x_pm1, w_arr)
        if wp is None:
            raise ValueError("prepacked_w missing packed weights")
        return fp32_gemm(x_pm1, unpack_binary_pm1(np.asarray(wp), n_feat))

    xp, n = pack_binary_pm1(x_pm1, axis=1)
    if prepacked_w is None:
        if w_arr is None:
            raise ValueError("w_pm1 is required when prepacked_w is omitted")
        wp, n2 = pack_binary_pm1(w_arr, axis=1)
        if n != n2:
            raise ValueError(f"packed n mismatch: {n} vs {n2}")
    else:
        wp, n2 = prepacked_w
        if n != n2:
            raise ValueError(f"packed n mismatch: {n} vs {n2}")

    native = binary_gemm_native_prepacked(xp, wp, n)
    if native is not None:
        return native
    return binary_gemm_numpy_prepacked(xp, wp, n)

binary_gemm_numpy_prepacked

binary_gemm_numpy_prepacked(
    xp: ndarray, wp: ndarray, n: int
) -> np.ndarray

Y = binary GEMM from pre-packed uint64 matrices (NumPy path).

Source code in bnn/kernels/packed.py

def binary_gemm_numpy_prepacked(
    xp: np.ndarray, wp: np.ndarray, n: int
) -> np.ndarray:
    """Y = binary GEMM from pre-packed uint64 matrices (NumPy path)."""
    B, M, _words = _validate_prepacked(xp, wp, n)
    out = np.empty((B, M), dtype=np.float32)
    # Row-at-a-time to keep temporaries small and cache-friendly
    for b in range(B):
        xor = np.bitwise_xor(xp[b : b + 1], wp)  # (M, words) via broadcast
        dist = bitwise_count(xor).sum(axis=1).astype(np.int32)
        out[b] = n - 2 * dist
    return out

binary_gemm_native_prepacked

binary_gemm_native_prepacked(
    xp: ndarray, wp: ndarray, n: int
) -> np.ndarray | None

Source code in bnn/kernels/packed.py

def binary_gemm_native_prepacked(
    xp: np.ndarray, wp: np.ndarray, n: int
) -> np.ndarray | None:
    lib = _try_load_native()
    if not lib:
        return None
    ensure_native_threads()
    B, M, words = _validate_prepacked(xp, wp, n)
    out = np.empty((B, M), dtype=np.float32)
    lib.binary_gemm_u64(
        xp.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
        wp.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
        out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
        B,
        M,
        words,
        n,
    )
    return out

binary_gemm_native_scaled

binary_gemm_native_scaled(
    xp: ndarray,
    wp: ndarray,
    n: int,
    alpha: ndarray | None = None,
    bias: ndarray | None = None,
) -> np.ndarray | None

Native GEMM with alpha / bias folded into the kernel epilogue.

Y = alpha * (n - 2*hamming) + bias in one pass. Returns None when the native library is missing or predates the fused entry point, so callers can fall back to the unfused path.

Source code in bnn/kernels/packed.py

def binary_gemm_native_scaled(
    xp: np.ndarray,
    wp: np.ndarray,
    n: int,
    alpha: np.ndarray | None = None,
    bias: np.ndarray | None = None,
) -> np.ndarray | None:
    """Native GEMM with ``alpha`` / ``bias`` folded into the kernel epilogue.

    ``Y = alpha * (n - 2*hamming) + bias`` in one pass. Returns ``None`` when
    the native library is missing or predates the fused entry point, so callers
    can fall back to the unfused path.
    """
    lib = _try_load_native()
    if not lib or not hasattr(lib, "binary_gemm_u64_scaled"):
        return None
    ensure_native_threads()
    B, M, words = _validate_prepacked(xp, wp, n)

    fptr = ctypes.POINTER(ctypes.c_float)
    null = fptr()  # NULL pointer

    def _vec(v: np.ndarray | None, name: str):
        if v is None:
            return null, None
        arr = np.ascontiguousarray(v, dtype=np.float32).reshape(-1)
        if arr.size != M:
            raise ValueError(f"{name} must have {M} elements, got {arr.size}")
        # Keep a reference alive until the call returns.
        return arr.ctypes.data_as(fptr), arr

    a_ptr, _a_keep = _vec(alpha, "alpha")
    b_ptr, _b_keep = _vec(bias, "bias")

    out = np.empty((B, M), dtype=np.float32)
    lib.binary_gemm_u64_scaled(
        xp.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
        wp.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
        out.ctypes.data_as(fptr),
        a_ptr,
        b_ptr,
        B,
        M,
        words,
        n,
    )
    return out

fp32_gemm

fp32_gemm(x: ndarray, w: ndarray) -> np.ndarray

FP32 reference GEMM: Y = X @ W.T.

Uses asarray rather than astype so already-float32 inputs are not copied. ndarray.astype copies unconditionally by default, which for a 4096x4096 baseline meant timing ~64 MB of memcpy alongside the GEMM and inflating every "vs FP32" speedup by ~2x.

Source code in bnn/kernels/packed.py

def fp32_gemm(x: np.ndarray, w: np.ndarray) -> np.ndarray:
    """FP32 reference GEMM: Y = X @ W.T.

    Uses ``asarray`` rather than ``astype`` so already-float32 inputs are *not*
    copied. ``ndarray.astype`` copies unconditionally by default, which for a
    4096x4096 baseline meant timing ~64 MB of memcpy alongside the GEMM and
    inflating every "vs FP32" speedup by ~2x.
    """
    return np.asarray(x, dtype=np.float32) @ np.asarray(w, dtype=np.float32).T

Threads

set_num_threads

set_num_threads(n: int | None) -> None

Set native OpenMP thread count (None / 0 = library default).

Also honors process env when first applied via ensure_native_threads().

Source code in bnn/kernels/packed.py

def set_num_threads(n: int | None) -> None:
    """Set native OpenMP thread count (None / 0 = library default).

    Also honors process env when first applied via ensure_native_threads().
    """
    global _THREADS_APPLIED
    lib = _try_load_native()
    if not lib:
        _THREADS_APPLIED = n if n and n > 0 else None
        return
    val = int(n) if n and n > 0 else 0
    if hasattr(lib, "binary_gemm_set_num_threads"):
        lib.binary_gemm_set_num_threads(val)
    _THREADS_APPLIED = val if val > 0 else None

get_num_threads

get_num_threads() -> int

Effective native thread count (1 if no OpenMP / no DLL).

Source code in bnn/kernels/packed.py

def get_num_threads() -> int:
    """Effective native thread count (1 if no OpenMP / no DLL)."""
    lib = _try_load_native()
    if lib and hasattr(lib, "binary_gemm_get_num_threads"):
        return int(lib.binary_gemm_get_num_threads())
    return 1

openmp_enabled

openmp_enabled() -> bool

Source code in bnn/kernels/packed.py

def openmp_enabled() -> bool:
    lib = _try_load_native()
    if lib and hasattr(lib, "binary_gemm_openmp_enabled"):
        return bool(lib.binary_gemm_openmp_enabled())
    return False

Building the native library

unix_compile_commands

unix_compile_commands(
    cc: str, out: Path, src: Path, openmp: bool
) -> list[list[str]]

Candidate compiler invocations, most-preferred first.

Deliberately no -march=native: the library selects AVX2 / AVX-512 / NEON at run time, so the object must stay portable to any CPU of the same architecture. Baking in build-host ISA would produce binaries that SIGILL on older machines — the opposite of what runtime dispatch is for.

Source code in bnn/kernels/compile_native.py

def unix_compile_commands(cc: str, out: Path, src: Path, openmp: bool) -> list[list[str]]:
    """Candidate compiler invocations, most-preferred first.

    Deliberately no ``-march=native``: the library selects AVX2 / AVX-512 /
    NEON at *run* time, so the object must stay portable to any CPU of the
    same architecture. Baking in build-host ISA would produce binaries that
    SIGILL on older machines — the opposite of what runtime dispatch is for.
    """
    base = ["-O3", "-shared", "-fPIC"]
    cmds: list[list[str]] = []
    if openmp:
        cmds.append([cc, *base, "-fopenmp", "-o", str(out), str(src)])
        libomp = _brew_libomp() if sys.platform == "darwin" else None
        if libomp is not None:
            # Apple clang needs libomp routed through the preprocessor.
            cmds.append([
                cc, *base,
                "-Xpreprocessor", "-fopenmp",
                f"-I{libomp / 'include'}",
                f"-L{libomp / 'lib'}",
                "-lomp",
                "-o", str(out), str(src),
            ])
    # Single-threaded fallback always builds; correctness never depends on OpenMP.
    cmds.append([cc, *base, "-o", str(out), str(src)])
    return cmds