Digital Drosophila Archive
575 lines
31 KiB
Python
575 lines
31 KiB
Python
"""Independent CUDA eligibility candidate; never installs weights or runs a brain.
|
|
|
|
The caller authenticates frozen sources, configuration, the complete KC list,
|
|
and the candidate edge mapping. This module does not infer D from DAN activity.
|
|
CuPy is imported only when CudaEligibility is instantiated. Set CUPY_CACHE_DIR
|
|
to a dedicated candidate directory before doing so. The only files this module
|
|
writes are explicit nvcc build evidence and cubins beneath that cache directory.
|
|
|
|
API: prepare_events(...) once, then CudaEligibility(...).run(prepared, ...).
|
|
All commands are sampled AFTER the current tick; only the LAST command snapshot
|
|
feeds the single epoch-boundary update. With no commands sampled_trace is +0.
|
|
cpu_reference has the same run arguments for host-only validation, and is not a
|
|
replacement for comparison against the SHA-pinned frozen implementation.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import time
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
|
|
COMPARTMENTS = ("B2B'2a", "a1", "y2a'1", "y5B'2a")
|
|
_INT32_MAX = np.iinfo(np.int32).max
|
|
CUDA_OPTIONS = ("--std=c++11", "--fmad=false", "--ftz=false")
|
|
DEFAULT_NVCC = "/usr/local/cuda/bin/nvcc"
|
|
CUDA_SOURCE = r'''
|
|
extern "C" __global__
|
|
void eligibility_trace(const long long* ptr, const int* ticks,
|
|
const int* commands, const int n_commands,
|
|
const int n_kc, const int n_steps, const float decay,
|
|
float* snapshots, float* final_trace, float* sampled) {
|
|
const int k = blockDim.x * blockIdx.x + threadIdx.x;
|
|
if (k >= n_kc) return;
|
|
long long cursor = ptr[k];
|
|
const long long end = ptr[k + 1];
|
|
int command = 0;
|
|
float trace = 0.0f;
|
|
float last_sample = 0.0f;
|
|
// Each KC follows the FULL sequential recurrence, including empty ticks.
|
|
for (int tick = 0; tick < n_steps; ++tick) {
|
|
float spike = 0.0f;
|
|
if (cursor < end && ticks[cursor] == tick) {
|
|
spike = 1.0f;
|
|
++cursor;
|
|
}
|
|
trace = __fmul_rn(trace, decay);
|
|
trace = __fadd_rn(trace, spike);
|
|
trace = trace < 1.0f ? trace : 1.0f;
|
|
if (command < n_commands && commands[command] == tick) {
|
|
snapshots[(long long)command * n_kc + k] = trace;
|
|
last_sample = trace;
|
|
++command;
|
|
}
|
|
}
|
|
final_trace[k] = trace;
|
|
sampled[k] = last_sample;
|
|
}
|
|
|
|
extern "C" __global__
|
|
void eligibility_delta(const float* sampled, const int* source_local,
|
|
const unsigned char* selected, const float* delta,
|
|
const float* base, const int n_edges,
|
|
const float negative_eta_times_d,
|
|
const float negative_floor, const int nonzero_d,
|
|
const int frozen, float* output) {
|
|
const int edge = blockDim.x * blockIdx.x + threadIdx.x;
|
|
if (edge >= n_edges) return;
|
|
if (frozen) {
|
|
output[edge] = delta[edge];
|
|
return;
|
|
}
|
|
float update = 0.0f;
|
|
if (nonzero_d && selected[edge])
|
|
update = __fmul_rn(negative_eta_times_d, sampled[source_local[edge]]);
|
|
const float out = __fadd_rn(delta[edge], update);
|
|
const float lower = __fmul_rn(negative_floor, base[edge]);
|
|
// np.maximum chooses its SECOND operand on equal finite values, including
|
|
// +0/-0. CUDA fmaxf is deliberately not substituted for that tie behavior.
|
|
output[edge] = out > lower ? out : lower;
|
|
}
|
|
'''
|
|
|
|
|
|
def _file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
for block in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _save_json_new(path: Path, value: Any) -> None:
|
|
with path.open("x", encoding="utf-8", newline="\n") as stream:
|
|
json.dump(value, stream, sort_keys=True, indent=2, allow_nan=False)
|
|
stream.write("\n")
|
|
|
|
|
|
def _nvcc_cubin(cache_root: Path, capability: str, nvcc_path: Path) -> tuple[Path, dict, dict]:
|
|
"""Compile outside CuPy's compiler path, which injects -ftz=true in 13.6.
|
|
|
|
A cache entry is immutable. Failed/incomplete entries are retained and cause
|
|
an error on reuse; the caller can select a fresh candidate cache directory.
|
|
Only the subprocess's temporary-directory environment is changed.
|
|
"""
|
|
probe_started = time.perf_counter()
|
|
_require(capability.isdigit() and 2 <= len(capability) <= 3, "invalid CUDA compute capability")
|
|
_require(cache_root.is_absolute(), "candidate cache directory must be absolute")
|
|
_require(nvcc_path.is_absolute() and nvcc_path.is_file(), "nvcc must be an existing absolute compiler path")
|
|
_require(not any(os.environ.get(key, "").strip()
|
|
for key in ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS")),
|
|
"NVCC injected flags are incompatible with explicit candidate compiler options")
|
|
cache_root = cache_root.resolve()
|
|
nvcc_path = nvcc_path.resolve()
|
|
version = subprocess.run([str(nvcc_path), "--version"], capture_output=True, text=True,
|
|
timeout=30, check=False)
|
|
_require(version.returncode == 0, "nvcc --version failed: " + version.stderr)
|
|
options = ("--cubin", "-arch=sm_" + capability, *CUDA_OPTIONS, "--keep")
|
|
identity = dict(schema="conditioning-cuda-nvcc-build-v1",
|
|
source_sha256=hashlib.sha256(CUDA_SOURCE.encode("utf-8")).hexdigest(),
|
|
nvcc_path=str(nvcc_path), nvcc_sha256=_file_sha256(nvcc_path),
|
|
nvcc_version_stdout=version.stdout, nvcc_version_stderr=version.stderr,
|
|
options=list(options), compute_capability=capability,
|
|
nvcc_ccbin=os.environ.get("NVCC_CCBIN", ""))
|
|
key = hashlib.sha256(json.dumps(identity, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
|
|
build_dir = cache_root / "eligibility-nvcc" / key
|
|
_require(not build_dir.is_symlink() and build_dir.resolve().is_relative_to(cache_root),
|
|
"candidate build cache path escapes the requested directory")
|
|
cubin = build_dir / "eligibility.cubin"
|
|
receipt_path = build_dir / "build-receipt.json"
|
|
probe_seconds = time.perf_counter() - probe_started
|
|
if build_dir.exists():
|
|
_require(receipt_path.is_file(), "incomplete nvcc candidate build retained: " + str(build_dir))
|
|
receipt = json.loads(receipt_path.read_text(encoding="utf-8"))
|
|
_require(receipt.get("complete") is True and receipt["identity"] == identity,
|
|
"failed or mismatched nvcc candidate build retained: " + str(build_dir))
|
|
for name, expected in receipt["files_sha256"].items():
|
|
path = build_dir / name
|
|
_require(path.name == name and not path.is_symlink()
|
|
and path.resolve().is_relative_to(build_dir)
|
|
and _file_sha256(path) == expected, "nvcc cache artifact changed: " + name)
|
|
_require(receipt["files_sha256"].get("eligibility.cubin") == _file_sha256(cubin)
|
|
and receipt["files_sha256"].get("eligibility.cu") == identity["source_sha256"],
|
|
"nvcc cache lacks authenticated source/cubin")
|
|
return cubin, dict(cache_key=key, cache_directory=str(build_dir), cache_hit=True,
|
|
receipt_sha256=_file_sha256(receipt_path), **receipt), dict(
|
|
compiler_probe_seconds=probe_seconds, nvcc_compile_seconds=0.0)
|
|
build_dir.mkdir(parents=True, exist_ok=False)
|
|
source_path = build_dir / "eligibility.cu"
|
|
with source_path.open("x", encoding="utf-8", newline="\n") as stream:
|
|
stream.write(CUDA_SOURCE)
|
|
command = [str(nvcc_path), *options, "--keep-dir", str(build_dir),
|
|
str(source_path), "-o", str(cubin)]
|
|
_save_json_new(build_dir / "compile-command.json", dict(argv=command, cwd=str(build_dir),
|
|
temp_directory=str(build_dir), inherited_nvcc_injected_flags=False))
|
|
_save_json_new(build_dir / "compiler-identity.json", identity)
|
|
environment = dict(os.environ, TMPDIR=str(build_dir), TMP=str(build_dir), TEMP=str(build_dir))
|
|
compile_started = time.perf_counter()
|
|
returncode = None
|
|
failure = None
|
|
try:
|
|
with (build_dir / "compiler.stdout.log").open("xb") as stdout, \
|
|
(build_dir / "compiler.stderr.log").open("xb") as stderr:
|
|
process = subprocess.run(command, cwd=build_dir, env=environment,
|
|
stdout=stdout, stderr=stderr, timeout=120, check=False)
|
|
returncode = process.returncode
|
|
except (OSError, subprocess.TimeoutExpired) as exc:
|
|
failure = f"{type(exc).__name__}: {exc}"
|
|
compile_seconds = time.perf_counter() - compile_started
|
|
complete = returncode == 0 and cubin.is_file() and cubin.stat().st_size > 0
|
|
receipt = dict(complete=complete, identity=identity, command=command,
|
|
returncode=returncode, failure=failure,
|
|
nvcc_compile_seconds=compile_seconds,
|
|
files_sha256={path.name: _file_sha256(path) for path in sorted(build_dir.iterdir())
|
|
if path.is_file() and not path.is_symlink()})
|
|
_save_json_new(receipt_path, receipt)
|
|
if not complete:
|
|
raise RuntimeError("explicit nvcc candidate compilation failed; evidence retained at " + str(build_dir))
|
|
return cubin, dict(cache_key=key, cache_directory=str(build_dir), cache_hit=False,
|
|
receipt_sha256=_file_sha256(receipt_path), **receipt), dict(
|
|
compiler_probe_seconds=probe_seconds, nvcc_compile_seconds=compile_seconds)
|
|
|
|
|
|
def _require(condition: Any, message: str) -> None:
|
|
if not condition:
|
|
raise ValueError(message)
|
|
|
|
|
|
def _integer(value: Any, name: str, minimum: int = 1) -> int:
|
|
_require(isinstance(value, (int, np.integer)) and not isinstance(value, (bool, np.bool_)),
|
|
f"{name} must be an integer")
|
|
_require(minimum <= int(value) <= _INT32_MAX, f"{name} outside supported int32 range")
|
|
return int(value)
|
|
|
|
|
|
def _array(value: Any, dtype: Any, name: str) -> np.ndarray:
|
|
_require(isinstance(value, np.ndarray) and value.ndim == 1 and value.dtype == np.dtype(dtype),
|
|
f"{name} must be a one-dimensional {np.dtype(dtype)} ndarray")
|
|
return value
|
|
|
|
|
|
def _owned(value: np.ndarray) -> np.ndarray:
|
|
result = np.array(value, copy=True, order="C")
|
|
result.setflags(write=False)
|
|
return result
|
|
|
|
|
|
def _scalar(value: Any, name: str) -> float:
|
|
_require(isinstance(value, (float, int, np.floating, np.integer))
|
|
and not isinstance(value, (bool, np.bool_)), f"{name} must be a real scalar")
|
|
try:
|
|
result = float(value)
|
|
except (ValueError, OverflowError) as exc:
|
|
raise ValueError(f"{name} must be a finite scalar") from exc
|
|
_require(math.isfinite(result), f"{name} must be finite")
|
|
return result
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PreparedEvents:
|
|
kc: np.ndarray
|
|
indptr: np.ndarray
|
|
ticks: np.ndarray
|
|
command_steps: np.ndarray
|
|
steps: int
|
|
n_neurons: int
|
|
input_event_count: int
|
|
preparation_seconds: float
|
|
|
|
@property
|
|
def kc_event_count(self) -> int:
|
|
return len(self.ticks)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EligibilityResult:
|
|
sampled_trace: np.ndarray
|
|
final_trace: np.ndarray
|
|
snapshots: np.ndarray
|
|
delta_after: np.ndarray
|
|
timing: dict[str, Any]
|
|
audit: dict[str, Any]
|
|
|
|
|
|
def prepare_events(step: np.ndarray, neuron_idx: np.ndarray, kc: np.ndarray,
|
|
steps: int, command_steps: Any, *, n_neurons: int = 166700) -> PreparedEvents:
|
|
"""Validate original sparse events and pack sorted per-KC event CSR on CPU.
|
|
|
|
Inputs must be sorted by (tick, global neuron ID), unique and int32. No
|
|
[steps, neurons] raster is allocated. Original input storage is not retained.
|
|
Empty, one, or multiple sorted unique command ticks are supported.
|
|
"""
|
|
started = time.perf_counter()
|
|
steps = _integer(steps, "steps")
|
|
n_neurons = _integer(n_neurons, "n_neurons")
|
|
step = _array(step, np.int32, "step")
|
|
neuron_idx = _array(neuron_idx, np.int32, "neuron_idx")
|
|
kc = _array(kc, np.int32, "kc")
|
|
_require(0 < len(kc) <= _INT32_MAX and np.all((kc >= 0) & (kc < n_neurons)),
|
|
"KC IDs empty or outside graph")
|
|
_require(np.all(kc[1:] > kc[:-1]), "KC IDs must be sorted unique global IDs")
|
|
_require(len(step) == len(neuron_idx), "event columns have different lengths")
|
|
_require(np.all((step >= 0) & (step < steps)), "event tick outside trial")
|
|
_require(np.all((neuron_idx >= 0) & (neuron_idx < n_neurons)), "event neuron outside graph")
|
|
_require(np.all((step[1:] > step[:-1]) |
|
|
((step[1:] == step[:-1]) & (neuron_idx[1:] > neuron_idx[:-1]))),
|
|
"events must be sorted by tick/global ID and unique")
|
|
commands = np.asarray(command_steps)
|
|
_require(commands.ndim == 1, "command_steps must be one-dimensional")
|
|
if commands.size == 0:
|
|
commands = np.empty(0, np.int32)
|
|
_require(commands.dtype.kind in "iu", "command_steps must be integers")
|
|
_require(np.all((commands >= 0) & (commands < steps)), "command tick outside trial")
|
|
_require(np.all(commands[1:] > commands[:-1]), "command_steps must be sorted unique")
|
|
local = np.searchsorted(kc, neuron_idx)
|
|
keep = (local < len(kc)) & (kc[np.minimum(local, len(kc) - 1)] == neuron_idx)
|
|
local = local[keep]
|
|
order = np.argsort(local, kind="stable")
|
|
ticks = step[keep][order]
|
|
counts = np.bincount(local, minlength=len(kc)).astype(np.int64, copy=False)
|
|
ptr = np.empty(len(kc) + 1, np.int64)
|
|
ptr[0] = 0
|
|
np.cumsum(counts, out=ptr[1:])
|
|
return PreparedEvents(_owned(kc), _owned(ptr), _owned(ticks),
|
|
_owned(commands.astype(np.int32)), steps, n_neurons,
|
|
len(step), time.perf_counter() - started)
|
|
|
|
|
|
def _validate_prepared(prepared: PreparedEvents) -> None:
|
|
_require(isinstance(prepared, PreparedEvents), "prepared must come from prepare_events")
|
|
_integer(prepared.steps, "steps")
|
|
_integer(prepared.n_neurons, "n_neurons")
|
|
kc = _array(prepared.kc, np.int32, "prepared.kc")
|
|
ptr = _array(prepared.indptr, np.int64, "prepared.indptr")
|
|
ticks = _array(prepared.ticks, np.int32, "prepared.ticks")
|
|
commands = _array(prepared.command_steps, np.int32, "prepared.command_steps")
|
|
_require(0 < len(kc) <= _INT32_MAX and np.all((kc >= 0) & (kc < prepared.n_neurons))
|
|
and np.all(kc[1:] > kc[:-1]), "invalid prepared KC IDs")
|
|
_require(len(ptr) == len(kc) + 1 and ptr[0] == 0 and ptr[-1] == len(ticks)
|
|
and np.all(ptr[1:] >= ptr[:-1]), "invalid prepared CSR offsets")
|
|
_require(np.all((ticks >= 0) & (ticks < prepared.steps)), "invalid prepared event ticks")
|
|
# A decrease/equality is permitted only at a CSR row boundary. Checking all
|
|
# adjacent pairs stays vectorized and also catches a forged PreparedEvents.
|
|
nonincreasing = ticks[1:] <= ticks[:-1]
|
|
if len(ticks) > 1:
|
|
boundaries = ptr[1:-1]
|
|
boundaries = boundaries[(boundaries > 0) & (boundaries < len(ticks))] - 1
|
|
nonincreasing[boundaries] = False
|
|
_require(not np.any(nonincreasing), "prepared KC ticks must be sorted unique")
|
|
_require(np.all((commands >= 0) & (commands < prepared.steps))
|
|
and np.all(commands[1:] > commands[:-1]), "invalid prepared command ticks")
|
|
|
|
|
|
def _validate_run(prepared, candidate_sources, candidate_compartments, compartment,
|
|
delta_before, base_candidate, d, frozen, eta_mv, floor_fraction, decay):
|
|
_validate_prepared(prepared)
|
|
sources = _array(candidate_sources, np.int32, "candidate_sources")
|
|
delta = _array(delta_before, np.float32, "delta_before")
|
|
base = _array(base_candidate, np.float32, "base_candidate")
|
|
comps = candidate_compartments
|
|
_require(isinstance(comps, np.ndarray) and comps.ndim == 1 and comps.dtype.kind == "U",
|
|
"candidate_compartments must be a Unicode vector")
|
|
_require(len(delta) <= _INT32_MAX and len(sources) == len(delta) == len(base) == len(comps),
|
|
"candidate columns have different lengths or exceed int32 capacity")
|
|
_require(compartment in COMPARTMENTS and np.all(np.isin(comps, COMPARTMENTS)),
|
|
"unrecognized exact frozen compartment")
|
|
local = np.searchsorted(prepared.kc, sources)
|
|
_require(np.all(local < len(prepared.kc)), "all candidate sources must be canonical KCs")
|
|
_require(np.all(prepared.kc[local] == sources), "all candidate sources must be canonical KCs")
|
|
d = _scalar(d, "D")
|
|
_require(d in (0.0, 1.0), "D must be the fixed artificial scalar 0 or 1")
|
|
_require(isinstance(frozen, (bool, np.bool_)), "frozen must be Boolean")
|
|
eta = _scalar(eta_mv, "eta_mv")
|
|
floor = _scalar(floor_fraction, "floor_fraction")
|
|
decay_value = _scalar(decay, "decay")
|
|
_require(0 < eta <= np.finfo(np.float32).max, "eta_mv must be positive finite float32 representable")
|
|
_require(0 < floor <= 1, "floor_fraction must lie in (0,1]")
|
|
_require(0 < decay_value <= 1, "decay must lie in (0,1]")
|
|
_require(np.float32(eta) > 0 and np.float32(floor) > 0 and np.float32(decay_value) > 0,
|
|
"rule scalar underflows float32")
|
|
_require(np.all(np.isfinite(base)) and np.all(base > 0),
|
|
"original base must be finite positive float32 mV")
|
|
lower = np.float32(-floor) * base
|
|
_require(np.all(np.isfinite(delta)) and np.all((delta >= lower) & (delta <= 0)),
|
|
"delta violates original frozen floor")
|
|
_require(not np.any((delta == 0) & np.signbit(delta)), "delta zero must be canonical +0")
|
|
return dict(local=np.array(local, dtype=np.int32, order="C"),
|
|
selected=np.asarray(comps == compartment, dtype=np.uint8),
|
|
delta=np.array(delta, copy=True, order="C"),
|
|
base=np.array(base, copy=True, order="C"), d=d, frozen=bool(frozen),
|
|
eta=eta, floor=floor, decay=decay_value,
|
|
negative_eta_times_d=np.float32(-eta * d), negative_floor=np.float32(-floor))
|
|
|
|
|
|
def _audit(prepared: PreparedEvents, work: dict, backend: str) -> dict[str, Any]:
|
|
return dict(backend=backend, steps=prepared.steps, kc_count=len(prepared.kc),
|
|
events=prepared.input_event_count, actual_kc_events=prepared.kc_event_count,
|
|
command_count=len(prepared.command_steps), candidate_count=len(work["base"]),
|
|
d_command=work["d"], frozen=work["frozen"], eta_mv=work["eta"],
|
|
floor_fraction=work["floor"], decay_python_float=work["decay"],
|
|
decay_float32_bits=int(np.asarray(np.float32(work["decay"])).view(np.uint32)),
|
|
initial_trace="float32 canonical +0", sampling="after_current_tick_spike",
|
|
update_sample="last command snapshot, or canonical +0 without commands",
|
|
trace_tick_skipping=False, dense_tick_neuron_raster=False,
|
|
original_source_executed=False, caller_must_verify_source_identity=True,
|
|
source_identity_verified_by_this_module=False,
|
|
complete_mapping_verified_by_this_module=False,
|
|
d_semantics="artificial scheduled command; not measured DAN activity",
|
|
actual_dan_activity_used_to_gate=False, brain_dynamics_executed=False,
|
|
training_executed=False, weights_installed=False, learning_pass_claimed=False)
|
|
|
|
|
|
def cpu_reference(prepared: PreparedEvents, candidate_sources: np.ndarray,
|
|
candidate_compartments: np.ndarray, compartment: str,
|
|
delta_before: np.ndarray, base_candidate: np.ndarray, *, d: float,
|
|
frozen: bool, eta_mv: float, floor_fraction: float, decay: float) -> EligibilityResult:
|
|
"""Host arithmetic mirror for tests; frozen-source validation is external."""
|
|
started = time.perf_counter()
|
|
work = _validate_run(prepared, candidate_sources, candidate_compartments, compartment,
|
|
delta_before, base_candidate, d, frozen, eta_mv, floor_fraction, decay)
|
|
nkc = len(prepared.kc)
|
|
trace = np.zeros(nkc, np.float32)
|
|
snapshots = np.empty((len(prepared.command_steps), nkc), np.float32)
|
|
positions = prepared.indptr[:-1].copy()
|
|
ends = prepared.indptr[1:]
|
|
spike = np.zeros(nkc, np.float32)
|
|
command = 0
|
|
for tick in range(prepared.steps):
|
|
spike.fill(0)
|
|
active = np.flatnonzero(positions < ends)
|
|
active = active[prepared.ticks[positions[active]] == tick]
|
|
spike[active] = np.float32(1)
|
|
positions[active] += 1
|
|
np.multiply(trace, np.float32(work["decay"]), out=trace)
|
|
np.add(trace, spike, out=trace)
|
|
np.minimum(trace, np.float32(1), out=trace)
|
|
if command < len(prepared.command_steps) and prepared.command_steps[command] == tick:
|
|
snapshots[command] = trace
|
|
command += 1
|
|
sampled = snapshots[-1].copy() if len(snapshots) else np.zeros(nkc, np.float32)
|
|
output = work["delta"].copy()
|
|
if not work["frozen"]:
|
|
update = np.zeros(len(output), np.float32)
|
|
selected = work["selected"].astype(bool)
|
|
if work["d"] != 0:
|
|
update[selected] = work["negative_eta_times_d"] * sampled[work["local"][selected]]
|
|
output = output + update
|
|
output = np.maximum(output, work["negative_floor"] * work["base"]).astype(np.float32)
|
|
return EligibilityResult(_owned(sampled), _owned(trace), _owned(snapshots), _owned(output),
|
|
dict(end_to_end_seconds=time.perf_counter() - started,
|
|
preparation_seconds=prepared.preparation_seconds),
|
|
_audit(prepared, work, "numpy-host-math-mirror"))
|
|
|
|
|
|
class CudaEligibility:
|
|
"""Explicit nvcc cubin, loaded by CuPy, with separate float32 rounding.
|
|
|
|
Use a dedicated environment and set an absolute CUPY_CACHE_DIR outside any
|
|
frozen tree before construction. The first run builds or authenticates one
|
|
nvcc cubin and loads its kernels with RawModule(path=...). CuPy's compiler
|
|
is not called, so it cannot inject an incompatible FTZ option. Cold timing
|
|
separates nvcc compilation, compiler probing and driver module loading.
|
|
run includes validation, transfers and output ownership in end-to-end wall.
|
|
"""
|
|
|
|
def __init__(self, device_index: int = 0, *, nvcc_path: str = DEFAULT_NVCC):
|
|
started = time.perf_counter()
|
|
self.device_index = _integer(device_index, "device_index", 0)
|
|
cache_dir = os.environ.get("CUPY_CACHE_DIR", "")
|
|
_require(bool(cache_dir) and Path(cache_dir).is_absolute(),
|
|
"set CUPY_CACHE_DIR to an absolute candidate cache directory outside frozen trees")
|
|
self.cache_root = Path(cache_dir)
|
|
self.nvcc_path = Path(nvcc_path)
|
|
try:
|
|
import cupy as cp
|
|
except ImportError as exc:
|
|
raise RuntimeError("CuPy unavailable; use the dedicated CUDA candidate environment") from exc
|
|
self.cp = cp
|
|
self.device = cp.cuda.Device(self.device_index)
|
|
with self.device:
|
|
self.stream = cp.cuda.Stream(non_blocking=True)
|
|
self.device.synchronize()
|
|
self.initialization_seconds = time.perf_counter() - started
|
|
self._compiled = False
|
|
|
|
def run(self, prepared: PreparedEvents, candidate_sources: np.ndarray,
|
|
candidate_compartments: np.ndarray, compartment: str,
|
|
delta_before: np.ndarray, base_candidate: np.ndarray, *, d: float,
|
|
frozen: bool, eta_mv: float, floor_fraction: float, decay: float,
|
|
kernel_warmup: int = 0, kernel_repeats: int = 1) -> EligibilityResult:
|
|
"""Transfer one fixture and optionally benchmark resident kernel repeats.
|
|
|
|
Each repeat starts from zero trace and the SAME delta_before; repeats
|
|
never accumulate training updates. CUDA timing fields are medians across
|
|
kernel_repeats. end_to_end_seconds includes all warmups and repeats, so
|
|
use default counts for a single-call end-to-end benchmark.
|
|
"""
|
|
started = time.perf_counter()
|
|
work = _validate_run(prepared, candidate_sources, candidate_compartments, compartment,
|
|
delta_before, base_candidate, d, frozen, eta_mv, floor_fraction, decay)
|
|
kernel_warmup = _integer(kernel_warmup, "kernel_warmup", 0)
|
|
kernel_repeats = _integer(kernel_repeats, "kernel_repeats")
|
|
validation_seconds = time.perf_counter() - started
|
|
cp = self.cp
|
|
first_run = not self._compiled
|
|
compile_seconds = 0.0
|
|
compilation_timing = dict(compiler_probe_seconds=0.0, nvcc_compile_seconds=0.0,
|
|
module_load_seconds=0.0)
|
|
with self.device, self.stream:
|
|
if first_run:
|
|
compile_started = time.perf_counter()
|
|
cubin, self.compilation_audit, probe_timing = _nvcc_cubin(
|
|
self.cache_root, str(self.device.compute_capability), self.nvcc_path)
|
|
compilation_timing.update(probe_timing)
|
|
load_started = time.perf_counter()
|
|
self.module = cp.RawModule(path=str(cubin))
|
|
self.trace_kernel = self.module.get_function("eligibility_trace")
|
|
self.delta_kernel = self.module.get_function("eligibility_delta")
|
|
self.stream.synchronize()
|
|
compilation_timing["module_load_seconds"] = time.perf_counter() - load_started
|
|
compile_seconds = time.perf_counter() - compile_started
|
|
self._compiled = True
|
|
transfer_started = time.perf_counter()
|
|
ptr_gpu, ticks_gpu, commands_gpu, local_gpu, selected_gpu, delta_gpu, base_gpu = [
|
|
cp.asarray(a) for a in (prepared.indptr, prepared.ticks, prepared.command_steps,
|
|
work["local"], work["selected"], work["delta"], work["base"])]
|
|
nkc, nedges = len(prepared.kc), len(work["delta"])
|
|
snapshots_gpu = cp.empty((len(prepared.command_steps), nkc), cp.float32)
|
|
final_gpu = cp.empty(nkc, cp.float32)
|
|
sampled_gpu = cp.empty(nkc, cp.float32)
|
|
output_gpu = cp.empty(nedges, cp.float32)
|
|
self.stream.synchronize()
|
|
upload_seconds = time.perf_counter() - transfer_started
|
|
trace_args = (
|
|
ptr_gpu, ticks_gpu, commands_gpu, np.int32(len(prepared.command_steps)),
|
|
np.int32(nkc), np.int32(prepared.steps), np.float32(work["decay"]),
|
|
snapshots_gpu, final_gpu, sampled_gpu)
|
|
delta_args = (
|
|
sampled_gpu, local_gpu, selected_gpu, delta_gpu, base_gpu,
|
|
np.int32(nedges), work["negative_eta_times_d"], work["negative_floor"],
|
|
np.int32(work["d"] != 0), np.int32(work["frozen"]), output_gpu)
|
|
|
|
def launch_pair():
|
|
self.trace_kernel(((nkc + 127) // 128,), (128,), trace_args)
|
|
if nedges:
|
|
self.delta_kernel(((nedges + 255) // 256,), (256,), delta_args)
|
|
|
|
for _ in range(kernel_warmup):
|
|
launch_pair()
|
|
self.stream.synchronize()
|
|
resident_samples = []
|
|
for _ in range(kernel_repeats):
|
|
trace_start, trace_end, update_end = cp.cuda.Event(), cp.cuda.Event(), cp.cuda.Event()
|
|
trace_start.record(self.stream)
|
|
self.trace_kernel(((nkc + 127) // 128,), (128,), trace_args)
|
|
trace_end.record(self.stream)
|
|
if nedges:
|
|
self.delta_kernel(((nedges + 255) // 256,), (256,), delta_args)
|
|
update_end.record(self.stream)
|
|
update_end.synchronize()
|
|
resident_samples.append(dict(
|
|
trace_kernel_seconds=cp.cuda.get_elapsed_time(trace_start, trace_end) / 1000.0,
|
|
update_kernel_seconds=cp.cuda.get_elapsed_time(trace_end, update_end) / 1000.0,
|
|
kernel_seconds=cp.cuda.get_elapsed_time(trace_start, update_end) / 1000.0))
|
|
medians = {key: float(np.median([sample[key] for sample in resident_samples]))
|
|
for key in resident_samples[0]}
|
|
download_started = time.perf_counter()
|
|
sampled, final, snapshots, output = [
|
|
_owned(cp.asnumpy(a, stream=self.stream, blocking=True))
|
|
for a in (sampled_gpu, final_gpu, snapshots_gpu, output_gpu)]
|
|
download_seconds = time.perf_counter() - download_started
|
|
guard_started = time.perf_counter()
|
|
for name, array in (("sampled_trace", sampled), ("final_trace", final), ("snapshots", snapshots)):
|
|
if not (np.all(np.isfinite(array)) and np.all((array >= 0) & (array <= 1))):
|
|
raise RuntimeError(f"CUDA output guard failed: invalid {name}")
|
|
lower = work["negative_floor"] * work["base"]
|
|
if not (np.all(np.isfinite(output)) and np.all((output >= lower) & (output <= 0))):
|
|
raise RuntimeError("CUDA output guard failed: delta outside finite frozen floor")
|
|
selected = work["selected"].astype(bool)
|
|
if output[~selected].tobytes() != work["delta"][~selected].tobytes():
|
|
raise RuntimeError("CUDA output guard failed: off-compartment delta changed")
|
|
if (work["frozen"] or work["d"] == 0) and output.tobytes() != work["delta"].tobytes():
|
|
raise RuntimeError("CUDA output guard failed: frozen/zero-D identity")
|
|
guard_seconds = time.perf_counter() - guard_started
|
|
timing = dict(first_run=first_run, initialization_seconds=self.initialization_seconds,
|
|
cold_compile_seconds=compile_seconds,
|
|
**compilation_timing,
|
|
compile_timing_semantics="first-use nvcc probe/build or authenticated cache, then driver module load",
|
|
preparation_seconds=prepared.preparation_seconds,
|
|
run_validation_seconds=validation_seconds,
|
|
host_to_device_and_allocation_seconds=upload_seconds,
|
|
**medians, device_to_host_seconds=download_seconds,
|
|
output_guard_seconds=guard_seconds,
|
|
resident_kernel_warmup=kernel_warmup,
|
|
resident_kernel_repeats=kernel_repeats,
|
|
resident_kernel_samples=resident_samples,
|
|
kernel_statistic="median of resident timed repeats",
|
|
end_to_end_semantics="all requested warmups/repeats and transfers; preparation excluded",
|
|
end_to_end_seconds=time.perf_counter() - started)
|
|
audit = _audit(prepared, work, "nvcc-cubin-cupy-independent-candidate")
|
|
audit.update(numpy_version=str(np.__version__), cupy_version=str(cp.__version__),
|
|
device_index=self.device_index, cuda_options=list(CUDA_OPTIONS),
|
|
compiler_backend="external nvcc --cubin; CuPy RawModule path loading only",
|
|
compilation=self.compilation_audit,
|
|
float32_separate_rounding_intrinsics=True, tf32_used=False,
|
|
numerical_equivalence_established_by_this_call=False)
|
|
return EligibilityResult(sampled, final, snapshots, output, timing, audit)
|