cache=True 可在不同的 NUMBA_SLP_VECTORIZE/NUMBA_OPT 设置下重复使用结果

import json, os, pathlib, pickle, sys, time os.environ["NUMBA_SLP_VECTORIZE"] = sys.argv[1] if len(sys.argv) > 1 else "0" import numba import numpy as np from numba import njit @njit(cache=True, fastmath=True, boundscheck=False, error_model="numpy") def kernel(x, out, A, C): """Short fixed-length loops with local scratch: the shape SLP vectorises.""" for e in range(x.shape[0]): u = np.empty(24); f = np.zeros(24); a = np.empty(6); b = np.empty(6) for i in range(24): u[i] = x[e, i] for g in range(8): for k in range(6): acc = 0.0 for i in range(24): acc += A[g, k, i] * u[i] a[k] = acc for k in range(6): acc = 0.0 for j in range(6): acc += C[k, j] * a[j] b[k] = acc for k in range(6): s = b[k] for i in range(24): f[i] += A[g, k, i] * s for i in range(24): out[e, i] = f[i] def index_entries(): files = sorted(pathlib.Path("pycache").glob("*.nbi")) if not files: return 0 with open(files[0], "rb") as handle: pickle.load(handle) # version stamp _, overloads = pickle.loads(handle.read()) return len(overloads) rng = np.random.default_rng(0) A = np.ascontiguousarray(rng.standard_normal((8, 6, 24))) C = np.ascontiguousarray(rng.standard_normal((6, 6))) x = np.ascontiguousarray(rng.standard_normal((200_000, 24))) out = np.empty_like(x) started = time.perf_counter() kernel(x, out, A, C) first_call = time.perf_counter() - started times = [] for _ in range(5): moment = time.perf_counter() kernel(x, out, A, C) times.append(time.perf_counter() - moment) print(json.dumps({ "NUMBA_SLP_VECTORIZE": os.environ["NUMBA_SLP_VECTORIZE"], …