Accelerate ternary v3.3 mask decoding with AVX2 and Arm NEON

README.md

@@ -170,6 +170,22 @@ Reports are written to:
 `benchmarking/bit_1/reports/results_shapes_{device}.csv`
 `benchmarking/bit_1_58/reports/results_shapes_{device}.csv`
 
+**Compare the ternary v3.3 SIMD and scalar kernels**
+
+The scalar comparison is compiled from the same source with SIMD disabled.
+Preprocessing is excluded from both measurements.
+
+```bash
+python -m benchmarking.bit_1_58.bench_v3_3_simd \
+  --rows 3200 \
+  --cols 8640 \
+  --k-values 6 8 \
+  --warmup 10 \
+  --repeats 41
+```
+
+See the recorded [Apple M4 results](benchmarking/bit_1_58/reports/apple_m4_v3_3_simd.md).
+
 **Benchmark end-to-end LLM inference**
 
 Pass either a single preprocessed model directory or a parent directory that

benchmarking/bit_1_58/bench_v3_3_simd.py

@@ -0,0 +1,184 @@
+"""Benchmark the v3.3 SIMD mask decoder against its scalar fallback.
+
+The scalar comparison library is compiled from the same source with
+``RSR_DISABLE_SIMD``. Weight preprocessing is excluded from both timings.
+"""
+
+import argparse
+import ctypes
+import json
+import os
+import platform
+import statistics
+import subprocess
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from multiplier.bit_1_58.cpu import rsr_nonsquare as rsr  # noqa: E402
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--rows", type=int, default=3200)
+    parser.add_argument("--cols", type=int, default=8640)
+    parser.add_argument("--k-values", type=int, nargs="+", default=[6, 8])
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--repeats", type=int, default=41)
+    parser.add_argument("--seed", type=int, default=20260724)
+    parser.add_argument("--cc", default=os.environ.get("CC", "cc"))
+    return parser.parse_args()
+
+
+def build_scalar_library(cc, output):
+    kernel_dir = ROOT / "kernels" / "bit_1_58" / "cpu"
+    command = [
+        cc,
+        "-O3",
+        "-march=native",
+        "-fomit-frame-pointer",
+        "-DRSR_DISABLE_SIMD",
+        "-shared",
+        "-fPIC",
+        str(kernel_dir / "rsr_ternary_v3_3.c"),
+        "-o",
+        str(output),
+    ]
+    subprocess.run(command, check=True)
+
+
+def configure_library(path):
+    library = ctypes.CDLL(str(path))
+    function = library.rsr_ternary_gemv_v3_3
+    function.restype = None
+    function.argtypes = rsr._v33_lib.rsr_ternary_gemv_v3_3.argtypes
+    return library
+
+
+def call_kernel(library, multiplier, vector, output):
+    library.rsr_ternary_gemv_v3_3(
+        multiplier._perms_u16_ptr,
+        multiplier._group_ends_u16_ptr,
+        multiplier._pos_masks_ptr,
+        multiplier._neg_masks_ptr,
+        multiplier._block_meta_ptr,
+        rsr.tensor_float_ptr(vector),
+        rsr.tensor_float_ptr(output),
+        multiplier.n_cols,
+        multiplier.k,
+        multiplier._num_blocks,
+    )
+
+
+def measure(library, multiplier, vector, output, warmup, repeats):
+    for _ in range(warmup):
+        call_kernel(library, multiplier, vector, output)
+
+    samples = []
+    for _ in range(repeats):
+        started = time.perf_counter_ns()
+        call_kernel(library, multiplier, vector, output)
+        samples.append((time.perf_counter_ns() - started) / 1e6)
+
+    return {
+        "median_ms": statistics.median(samples),
+        "p10_ms": float(np.percentile(samples, 10)),
+        "p90_ms": float(np.percentile(samples, 90)),
+    }
+
+
+def main():
+    args = parse_args()
+    if args.rows <= 0 or args.cols <= 0:
+        raise ValueError("rows and cols must be positive")
+    if args.warmup < 0 or args.repeats <= 0:
+        raise ValueError("warmup must be non-negative and repeats must be positive")
+
+    generator = torch.Generator().manual_seed(args.seed)
+    matrix = torch.randint(
+        -1,
+        2,
+        (args.rows, args.cols),
+        generator=generator,
+        dtype=torch.int8,
+    )
+    vector = torch.randn(args.cols, generator=generator, dtype=torch.float32)
+
+    result = {
+        "machine": platform.machine(),
+        "processor": platform.processor(),
+        "platform": platform.platform(),
+        "shape": [args.rows, args.cols],
+        "warmup": args.warmup,
+        "repeats": args.repeats,
+        "seed": args.seed,
+        "results": [],
+    }
+
+    with tempfile.TemporaryDirectory(prefix="rsr-v3-3-bench-") as temp_dir:
+        scalar_path = Path(temp_dir) / "rsr_ternary_v3_3_scalar.so"
+        build_scalar_library(args.cc, scalar_path)
+        scalar_library = configure_library(scalar_path)
+
+        for k in args.k_values:
+            multiplier = rsr.RSRTernaryNonSquareMultiplier(matrix, k)
+            if not multiplier._use_v33:
+                raise ValueError(
+                    f"shape {args.rows}x{args.cols} does not select v3.3"
+                )
+            optimized_output = torch.empty(
+                multiplier.n_rows_padded, dtype=torch.float32
+            )
+            scalar_output = torch.empty(
+                multiplier.n_rows_padded, dtype=torch.float32
+            )
+
+            call_kernel(
+                rsr._v33_lib, multiplier, vector, optimized_output
+            )
+            call_kernel(
+                scalar_library, multiplier, vector, scalar_output
+            )
+            torch.testing.assert_close(optimized_output, scalar_output)
+            max_abs_difference = float(
+                (optimized_output - scalar_output).abs().max()
+            )
+
+            optimized = measure(
+                rsr._v33_lib,
+                multiplier,
+                vector,
+                optimized_output,
+                args.warmup,
+                args.repeats,
+            )
+            scalar = measure(
+                scalar_library,
+                multiplier,
+                vector,
+                scalar_output,
+                args.warmup,
+                args.repeats,
+            )
+            result["results"].append(
+                {
+                    "k": k,
+                    "optimized": optimized,
+                    "scalar": scalar,
+                    "speedup": scalar["median_ms"] / optimized["median_ms"],
+                    "max_abs_difference": max_abs_difference,
+                }
+            )
+
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

benchmarking/bit_1_58/reports/apple_m4_v3_3_simd.md

@@ -0,0 +1,29 @@
+# Ternary v3.3 SIMD benchmark on Apple M4
+
+Measured on July 24, 2026, using a 10-core Apple M4 MacBook Air with 24 GB
+memory and macOS 26.2. Both kernels ran single-threaded. Matrix preprocessing
+was excluded.
+
+The optimized and scalar libraries were compiled from the same
+`rsr_ternary_v3_3.c` source. The scalar library used `RSR_DISABLE_SIMD`.
+
+```bash
+python -m benchmarking.bit_1_58.bench_v3_3_simd \
+  --rows 3200 \
+  --cols 8640 \
+  --k-values 6 8 \
+  --warmup 10 \
+  --repeats 41 \
+  --seed 20260724
+```
+
+| `k` | Optimized median | Scalar median | Speedup | Optimized p10–p90 | Scalar p10–p90 |
+|---:|---:|---:|---:|---:|---:|
+| 6 | 5.653 ms | 8.653 ms | 1.53× | 5.608–5.817 ms | 8.515–8.942 ms |
+| 8 | 14.466 ms | 31.164 ms | 2.15× | 14.335–14.683 ms | 30.945–31.675 ms |
+
+These results apply to this matrix shape and hardware. They are kernel-level
+latencies, not end-to-end model throughput. On Apple Silicon, `k=4` remains on
+the scalar path because the SIMD table decoder did not improve latency at this
+production-sized shape. Optimized and scalar outputs were identical for both
+recorded runs (`max_abs_difference = 0`).

docs/OPTIMIZATION.md

@@ -170,6 +170,12 @@ Over v1.4:
 
 Replaces the variable-length signed scatter arrays with two fixed-size `uint16` masks per group: `pos_mask` and `neg_mask`. The kernel iterates set bits with `__builtin_ctz` and `mask &= mask - 1`. Requires `k ≤ 16`.
 
+For `k ≤ 8`, AVX2 can decode the masks through a coefficient table and
+accumulate all output rows as one vector. Arm NEON uses the same approach for
+`k=5..8`; `k=4` retains the scalar mask walk because it is faster at
+production-sized Apple Silicon shapes. Other sizes use the original scalar
+path.
+
 This is **the key ternary CPU optimization**: metadata per group stops depending on how many rows are active. Two compact masks replace a variable-length scatter list.
 
 ### `RSRTernaryNonSquareMultiplier`

kernels/bit_1_58/cpu/Makefile

@@ -1,34 +1,44 @@
 CC = gcc
 CFLAGS = -O3 -Wall -march=native -fomit-frame-pointer
+UNAME_S := $(shell uname -s)
+UNAME_M := $(shell uname -m)
+
+ifeq ($(UNAME_M),x86_64)
+SIMD_FLAGS = -mavx2
+endif
+
+ifneq ($(UNAME_S),Darwin)
+OPENMP_FLAGS = -fopenmp
+endif
 
 all: bitnet_ternary.so tmac_ternary.so rsr_ternary.so rsr_ternary_prep.so rsr_ternary_prep_nonsquare.so rsr_ternary_v3_1.so rsr_ternary_v3_1_batch.so rsr_ternary_v3_3.so rsr_ternary_v3_3_batch.so
 
 bitnet_ternary.so: bitnet_ternary.c
-	$(CC) $(CFLAGS) -mavx2 -shared -fPIC -o bitnet_ternary.so bitnet_ternary.c -lm
+	$(CC) $(CFLAGS) $(SIMD_FLAGS) -shared -fPIC -o bitnet_ternary.so bitnet_ternary.c -lm
 
 tmac_ternary.so: tmac_ternary.c
-	$(CC) $(CFLAGS) -mavx2 -shared -fPIC -o tmac_ternary.so tmac_ternary.c -lm
+	$(CC) $(CFLAGS) $(SIMD_FLAGS) -shared -fPIC -o tmac_ternary.so tmac_ternary.c -lm
 
 rsr_ternary.so: rsr_ternary.c
-	$(CC) $(CFLAGS) -mavx2 -fopenmp -shared -fPIC -o rsr_ternary.so rsr_ternary.c
+	$(CC) $(CFLAGS) $(SIMD_FLAGS) $(OPENMP_FLAGS) -shared -fPIC -o rsr_ternary.so rsr_ternary.c
 
 rsr_ternary_prep.so: rsr_ternary_prep.c
-	$(CC) $(CFLAGS) -fopenmp -shared -fPIC -o rsr_ternary_prep.so rsr_ternary_prep.c
+	$(CC) $(CFLAGS) $(OPENMP_FLAGS) -shared -fPIC -o rsr_ternary_prep.so rsr_ternary_prep.c
 
 rsr_ternary_prep_nonsquare.so: rsr_ternary_prep_nonsquare.c
-	$(CC) $(CFLAGS) -fopenmp -shared -fPIC -o rsr_ternary_prep_nonsquare.so rsr_ternary_prep_nonsquare.c
+	$(CC) $(CFLAGS) $(OPENMP_FLAGS) -shared -fPIC -o rsr_ternary_prep_nonsquare.so rsr_ternary_prep_nonsquare.c
 
 rsr_ternary_v3_1.so: rsr_ternary_v3_1.c
-	$(CC) $(CFLAGS) -mavx2 -fopenmp -shared -fPIC -o rsr_ternary_v3_1.so rsr_ternary_v3_1.c
+	$(CC) $(CFLAGS) $(SIMD_FLAGS) $(OPENMP_FLAGS) -shared -fPIC -o rsr_ternary_v3_1.so rsr_ternary_v3_1.c
 
 rsr_ternary_v3_1_batch.so: rsr_ternary_v3_1_batch.c
-	$(CC) $(CFLAGS) -mavx2 -fopenmp -shared -fPIC -o rsr_ternary_v3_1_batch.so rsr_ternary_v3_1_batch.c -lm
+	$(CC) $(CFLAGS) $(SIMD_FLAGS) $(OPENMP_FLAGS) -shared -fPIC -o rsr_ternary_v3_1_batch.so rsr_ternary_v3_1_batch.c -lm
 
-rsr_ternary_v3_3.so: rsr_ternary_v3_3.c
-	$(CC) $(CFLAGS) -mavx2 -fopenmp -shared -fPIC -o rsr_ternary_v3_3.so rsr_ternary_v3_3.c
+rsr_ternary_v3_3.so: rsr_ternary_v3_3.c rsr_mask_coeff_lut.h
+	$(CC) $(CFLAGS) $(SIMD_FLAGS) $(OPENMP_FLAGS) -shared -fPIC -o rsr_ternary_v3_3.so rsr_ternary_v3_3.c
 
 rsr_ternary_v3_3_batch.so: rsr_ternary_v3_3_batch.c
-	$(CC) $(CFLAGS) -mavx2 -fopenmp -shared -fPIC -o rsr_ternary_v3_3_batch.so rsr_ternary_v3_3_batch.c -lm
+	$(CC) $(CFLAGS) $(SIMD_FLAGS) $(OPENMP_FLAGS) -shared -fPIC -o rsr_ternary_v3_3_batch.so rsr_ternary_v3_3_batch.c -lm
 
 clean:
 	rm -f *.so