bench_python_np.py

#!/usr/bin/env python3
# Usage examples:
#   python3 bench_python_np.py
#   python3 bench_python_np.py --repeat 5 --sieve-limit 3000000 --matmul-n 220 --regex-len 8000000 --io-kb 32768
import argparse
import math
import os
import re
import sys
import tempfile
import time

def now():
    return time.perf_counter()

def time_call(fn, *args, **kwargs):
    t0 = now()
    result = fn(*args, **kwargs)
    t1 = now()
    return result, (t1 - t0)

# ----------------------- Optional NumPy backend -----------------------
try:
    import numpy as np
    _HAS_NUMPY = True
except Exception:  # pragma: no cover
    _HAS_NUMPY = False

MOD = 1_000_000_007

# ====================== Pure-Python baseline (your code) ======================
def sieve_primes_py(limit: int):
    if limit < 2:
        return []
    sieve = bytearray(b"\x01") * (limit + 1)
    sieve[0:2] = b"\x00\x00"
    m = int(limit ** 0.5)
    for p in range(2, m + 1):
        if sieve[p]:
            step = p
            start = p * p
            sieve[start:limit + 1:step] = b"\x00" * ((limit - start) // step + 1)
    return [i for i, v in enumerate(sieve) if v]

def matmul_naive_py(n: int):
    # Deterministic matrices, match original formulas
    A = [[(i * 31 + j * 17) % 97 for j in range(n)] for i in range(n)]
    B = [[(i * 13 + j * 7)  % 89 for j in range(n)] for i in range(n)]
    C = [[0]*n for _ in range(n)]
    for i in range(n):
        Ai = A[i]; Ci = C[i]
        for k in range(n):
            aik = Ai[k]; Bk = B[k]
            for j in range(n):
                Ci[j] += aik * Bk[j]
    return sum(c for row in C for c in row) % MOD

def regex_parse_count_py(total_len: int):
    def rnd_next(state):
        return (1103515245 * state + 12345) % 0x7fffffff
    state = 1234
    parts = []
    total = 0
    while total < total_len:
        state = rnd_next(state); v1 = state % 1_000_000
        state = rnd_next(state); v2 = state % 1_000_000
        state = rnd_next(state); v3 = state % 1_000_000
        state = rnd_next(state); v4 = state % 1_000_000
        chunk = f"val{v1}, foo{v2}; {v3}\nbar{v4} "
        parts.append(chunk)
        total += len(chunk)
    s = "".join(parts)[:total_len]
    nums = re.findall(r"\b\d+\b", s)
    return len(nums), sum(int(x) for x in nums) % MOD

def file_io_py(size_kb: int):
    line = b"The quick brown fox jumps over the lazy dog. 1234567890\n"
    target = size_kb * 1024
    # Repeat and slice using pure Python
    blob = (line * (target // len(line) + 1))[:target]
    checksum_w = sum(blob) % MOD
    with tempfile.NamedTemporaryFile(delete=False) as f:
        fname = f.name
        f.write(blob)
    try:
        with open(fname, "rb") as f:
            content = f.read()
        checksum_r = sum(content) % MOD
    finally:
        try:
            os.remove(fname)
        except OSError:
            pass
    return len(blob), checksum_w, checksum_r

# ============================ NumPy-optimized ===========================
def sieve_primes_np(limit: int):
    if limit < 2:
        return []
    # Boolean mask of size limit+1; True = prime, False = composite
    sieve = np.ones(limit + 1, dtype=bool)
    sieve[:2] = False
    m = int(limit ** 0.5)
    # Standard vectorized sieve: only loop primes up to sqrt(limit)
    for p in range(2, m + 1):
        if sieve[p]:
            sieve[p*p::p] = False
    return np.nonzero(sieve)[0].tolist()

def matmul_np(n: int):
    # Build matrices via broadcasting (I: n×1, J: 1×n)
    I = np.arange(n, dtype=np.int64)[:, None]
    J = np.arange(n, dtype=np.int64)[None, :]
    A = (I * 31 + J * 17) % 97  # n×n
    B = (I * 13 + J * 7)  % 89  # n×n
    C = A @ B                   # GEMM in C/BLAS
    return int(C.sum(dtype=np.int64) % MOD)

def regex_parse_count_np(total_len: int):
    """
    Keep the same LCG and exact text as the baseline, but:
      * build bytes (bytearray) for lower overhead
      * run a compiled *bytes* regex
      * sum with NumPy fromiter for a faster int conversion loop
    """
    def rnd_next(state):
        return (1103515245 * state + 12345) % 0x7fffffff
    state = 1234
    out = bytearray()
    line_append = out.extend  # local for speed
    while len(out) < total_len:
        state = rnd_next(state); v1 = state % 1_000_000
        state = rnd_next(state); v2 = state % 1_000_000
        state = rnd_next(state); v3 = state % 1_000_000
        state = rnd_next(state); v4 = state % 1_000_000
        # Construct the exact same ASCII bytes as the baseline
        line_append(f"val{v1}, foo{v2}; {v3}\nbar{v4} ".encode("ascii"))
    s = bytes(out[:total_len])

    pat = re.compile(rb"\b\d+\b")
    # Iterate matches (generator), convert to ints via NumPy for speed
    it = (int(m.group()) for m in pat.finditer(s))
    # dtype=int64 ensures large, fast accumulation in C
    arr = np.fromiter(it, dtype=np.int64, count=-1)
    count = int(arr.size)
    checksum = int(arr.sum(dtype=np.int64) % MOD)
    return count, checksum

def file_io_np(size_kb: int):
    """
    Same behavior as baseline, but compute checksums with NumPy and
    simplify blob construction using bytes repetition.
    """
    line = b"The quick brown fox jumps over the lazy dog. 1234567890\n"
    target = size_kb * 1024
    blob = (line * (target // len(line) + 1))[:target]
    # Fast checksum in C
    checksum_w = int(np.frombuffer(blob, dtype=np.uint8).sum(dtype=np.uint64) % MOD)

    with tempfile.NamedTemporaryFile(delete=False) as f:
        fname = f.name
        f.write(blob)
    try:
        with open(fname, "rb") as f:
            content = f.read()
        checksum_r = int(np.frombuffer(content, dtype=np.uint8).sum(dtype=np.uint64) % MOD)
    finally:
        try:
            os.remove(fname)
        except OSError:
            pass
    return len(blob), checksum_w, checksum_r

# =============================== Runner =================================
def pick_backend(name, backend):
    """
    Select NumPy version when available and requested, otherwise Python baseline.
    """
    if backend == "py" or not _HAS_NUMPY:
        if name == "sieve":  return sieve_primes_py
        if name == "matmul": return matmul_naive_py
        if name == "regex":  return regex_parse_count_py
        if name == "fileio": return file_io_py
    # default (np/auto with numpy available)
    if name == "sieve":  return sieve_primes_np
    if name == "matmul": return matmul_np
    if name == "regex":  return regex_parse_count_np
    if name == "fileio": return file_io_np
    raise ValueError(f"Unknown task: {name}")

def run_task(task_name, fn, *args):
    result, elapsed = time_call(fn, *args)
    return task_name, elapsed, result

def main():
    parser = argparse.ArgumentParser(description="Python micro-benchmarks (NumPy optimized).")
    parser.add_argument("--repeat", type=int, default=3, help="Timed iterations per task (default: 3)")
    parser.add_argument("--no-warmup", action="store_true", help="Disable warm-up run")
    parser.add_argument("--sieve-limit", type=int, default=2_000_000, help="Prime sieve limit")
    parser.add_argument("--matmul-n", type=int, default=180, help="Matrix size N for NxN multiply")
    parser.add_argument("--regex-len", type=int, default=5_000_000, help="Target string length for regex parse")
    parser.add_argument("--io-kb", type=int, default=16_384, help="File I/O size in KB (default: 16384 = 16 MB)")
    parser.add_argument("--verbose", action="store_true", help="Print per-step progress")
    parser.add_argument("--only", choices=["sieve", "matmul", "regex", "fileio"], nargs="*", help="Run only these tasks")
    parser.add_argument("--skip", choices=["sieve", "matmul", "regex", "fileio"], nargs="*", help="Skip these tasks")
    parser.add_argument("--quick", action="store_true", help="Use smaller sizes for a quick smoke test")
    parser.add_argument("--backend", choices=["auto", "np", "py"], default="auto",
                        help="Select backend: auto (prefer NumPy), np (force NumPy), py (force pure Python)")
    args = parser.parse_args()

    if args.quick:
        args.repeat      = min(args.repeat, 1)
        args.sieve_limit = min(args.sieve_limit, 200_000)
        args.matmul_n    = min(args.matmul_n, 100)
        args.regex_len   = min(args.regex_len, 1_000_000)
        args.io_kb       = min(args.io_kb, 1024)

    ver = sys.version.split()[0]
    hdr = f"Python {ver}  repeats={args.repeat}  warmup={'off' if args.no_warmup else 'on'}"
    if _HAS_NUMPY:
        hdr += f"  (NumPy {np.__version__}, backend={args.backend})"
    else:
        hdr += f"  (NumPy not available, backend=py)"
    print(hdr, flush=True)

    # Select functions per backend
    backend = args.backend
    if backend == "np" and not _HAS_NUMPY:
        print("[warn] --backend=np requested but NumPy not found; using pure Python backend.", flush=True)
        backend = "py"

    task_specs = [
        ("sieve",  pick_backend("sieve",  backend), (args.sieve_limit,)),
        ("matmul", pick_backend("matmul", backend), (args.matmul_n,)),
        ("regex",  pick_backend("regex",  backend), (args.regex_len,)),
        ("fileio", pick_backend("fileio", backend), (args.io_kb,)),
    ]

    # Filter by --only / --skip
    if args.only:
        only = set(args.only)
        task_specs = [t for t in task_specs if t[0] in only]
    if args.skip:
        skip = set(args.skip)
        task_specs = [t for t in task_specs if t[0] not in skip]

    if not task_specs:
        print("No tasks selected.", flush=True)
        return

    if args.verbose:
        print(f"[info] tasks: {', '.join(t[0] for t in task_specs)}", flush=True)

    # Warm-up
    if not args.no_warmup:
        for name, fn, fn_args in task_specs:
            if args.verbose:
                print(f"[warmup] {name}...", flush=True)
            _ = fn(*fn_args)

    # Timed runs
    results = {}
    for name, fn, fn_args in task_specs:
        times = []
        last_result = None
        if args.verbose:
            print(f"[run] {name} x{args.repeat}", flush=True)
        for r in range(args.repeat):
            if args.verbose:
                print(f" - iter {r+1}/{args.repeat}", flush=True)
            last_result, elapsed = time_call(fn, *fn_args)
            times.append(elapsed)
        results[name] = {
            "min": min(times),
            "avg": sum(times) / len(times),
            "max": max(times),
            "last_result": last_result
        }

    # Report
    print("\n=== Results (seconds) ===", flush=True)
    width = max(len(n) for n,_,_ in task_specs) + 2
    print(f"{'task'.ljust(width)} min       avg       max", flush=True)
    print("-" * (width + 28), flush=True)
    for name in [t[0] for t in task_specs]:
        r = results[name]
        print(f"{name.ljust(width)} {r['min']:.6f} {r['avg']:.6f} {r['max']:.6f}", flush=True)

    # Correctness crumbs (same as your baseline)
    if any(t[0] == "sieve" for t in task_specs):
        # count with whichever backend got used for sieve
        sieve_count = len(pick_backend("sieve", backend)(args.sieve_limit))
        print(f"\n(check) primes <= {args.sieve_limit}: {sieve_count}", flush=True)

    if "matmul" in results:
        print(f"(check) matmul N={args.matmul_n} checksum: {results['matmul']['last_result']}", flush=True)

    if "regex" in results:
        rr = results['regex']['last_result']
        print(f"(check) regex: count={rr[0]}, checksum={rr[1]}", flush=True)

    if "fileio" in results:
        fio = results['fileio']['last_result']
        print(f"(check) fileio bytes={fio[0]}, checksum write={fio[1]}, read={fio[2]}", flush=True)

if __name__ == "__main__":
    main()