From 30fd5324ced767c7f4f53fc6a4f6077ee94f03c5 Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Thu, 11 Jun 2026 15:44:18 +0200 Subject: [PATCH 1/3] Add real end-to-end example: YOLOv8 detection with fused GPU preprocessing examples/11_yolo_inference.py: JPEG -> ONE fused FKL kernel -> YOLOv8n. The preprocessing that ultralytics does in several numpy/torch steps is a single fused kernel: TensorRead (uchar3, any size) -> Resize(640x640, bilinear, PRESERVE_AR letterbox with the YOLO gray-114 padding) -> Div(255) -> TensorSplit (HWC -> planar CHW float, NCHW-ready). Standard decode+NMS postprocessing in numpy; inference via onnxruntime. Verified on bus.jpg (810x1080): 4 persons + 1 bus detected with correct boxes mapped back through the letterbox transform. Steady-state fused preprocess: ~0.55 ms (incl. D2H); kernel compiles once then disk-cached. examples/fetch_yolo_assets.sh downloads bus.jpg + a COCO-pretrained yolov8n.onnx; assets/ and .venv-yolo/ gitignored. --- .gitignore | 2 + examples/11_yolo_inference.py | 180 ++++++++++++++++++++++++++++++++++ examples/README.md | 3 + examples/fetch_yolo_assets.sh | 25 +++++ 4 files changed, 210 insertions(+) create mode 100644 examples/11_yolo_inference.py create mode 100755 examples/fetch_yolo_assets.sh diff --git a/.gitignore b/.gitignore index f6d0dfc..d54f302 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,5 @@ build/ .venv/ .cache/ err.log +examples/assets/ +.venv-yolo/ diff --git a/examples/11_yolo_inference.py b/examples/11_yolo_inference.py new file mode 100644 index 0000000..001b37a --- /dev/null +++ b/examples/11_yolo_inference.py @@ -0,0 +1,180 @@ +"""Example 11 — REAL end-to-end pipeline: YOLOv8 object detection with +fkl-python GPU preprocessing. + + JPEG -> [GPU: ONE fused kernel] -> YOLOv8n (onnxruntime) -> boxes + +The preprocessing that ultralytics does in several numpy/torch steps +(letterbox resize + RGB normalize + HWC->CHW) is ONE fused FKL kernel: + + TensorRead uchar3 frame (any size) + Resize(640, 640, bilinear letterbox, gray(114) padding -- + preserve AR) the exact YOLO convention + Div(255) normalize to [0,1] + TensorSplit packed HWC -> planar CHW float, NCHW-ready + +Requirements (see examples/assets/): + pip install numpy pillow onnxruntime + assets/bus.jpg (ultralytics sample image) + assets/yolov8n.onnx (COCO-pretrained YOLOv8n, 1x3x640x640 input) + +Run: python3 11_yolo_inference.py [image.jpg] +""" +import ctypes +import sys +import time +from pathlib import Path + +import numpy as np +from PIL import Image +import onnxruntime as ort + +import fkl + +ASSETS = Path(__file__).parent / "assets" +NET = 640 +CONF_THRES, IOU_THRES = 0.25, 0.45 + +COCO = ("person bicycle car motorcycle airplane bus train truck boat traffic-light " + "fire-hydrant stop-sign parking-meter bench bird cat dog horse sheep cow " + "elephant bear zebra giraffe backpack umbrella handbag tie suitcase frisbee " + "skis snowboard sports-ball kite baseball-bat baseball-glove skateboard " + "surfboard tennis-racket bottle wine-glass cup fork knife spoon bowl banana " + "apple sandwich orange broccoli carrot hot-dog pizza donut cake chair couch " + "potted-plant bed dining-table toilet tv laptop mouse remote keyboard " + "cell-phone microwave oven toaster sink refrigerator book clock vase " + "scissors teddy-bear hair-drier toothbrush").split() + + +# -------------------------------------------------------------------------- +# GPU preprocessing: ONE fused kernel, compiled once, reused every frame +# -------------------------------------------------------------------------- + +def build_preprocessor(): + return fkl.compose( + fkl.TensorRead(), # uchar3 HxW + fkl.Resize(NET, NET, interp="linear", + aspect_ratio="preserve", # letterbox + background=(114.0, 114.0, 114.0)), # YOLO gray pad + fkl.Div((255.0, 255.0, 255.0)), # [0,1] + fkl.TensorSplit(), # HWC -> CHW float + ) + + +def upload_image(img_np): + """HxWx3 uint8 numpy -> fkl.DeviceBuffer (zero further copies).""" + h, w, _ = img_np.shape + buf = fkl.DeviceBuffer(w, h, "uint8", channels=3) + buf.copy_from_host(img_np.tobytes()) + return buf + + +def gpu_preprocess(pipe, frame_dev): + """Run the fused kernel; return a (1,3,640,640) float32 numpy view.""" + out = pipe(frame_dev) # ONE kernel launch + chw = np.frombuffer(out.copy_to_host(), dtype=np.float32) + return chw.reshape(1, 3, NET, NET) + + +# -------------------------------------------------------------------------- +# YOLO postprocessing (standard: decode + NMS), CPU numpy +# -------------------------------------------------------------------------- + +def letterbox_params(w, h): + """Replicate FKL's PRESERVE_AR placement: scale + centered offsets.""" + scale = min(NET / w, NET / h) + nw, nh = round(w * scale), round(h * scale) + return scale, (NET - nw) // 2, (NET - nh) // 2 + + +def postprocess(pred, orig_w, orig_h): + """pred: (1, 84, 8400) -> [(x1,y1,x2,y2,conf,cls), ...] in image coords.""" + p = pred[0].T # (8400, 84) + boxes, scores = p[:, :4], p[:, 4:] + cls = scores.argmax(1) + conf = scores.max(1) + keep = conf >= CONF_THRES + boxes, conf, cls = boxes[keep], conf[keep], cls[keep] + + # xywh (letterbox space) -> xyxy (original image space) + scale, ox, oy = letterbox_params(orig_w, orig_h) + x, y, w, h = boxes.T + x1 = (x - w / 2 - ox) / scale + y1 = (y - h / 2 - oy) / scale + x2 = (x + w / 2 - ox) / scale + y2 = (y + h / 2 - oy) / scale + boxes = np.stack([x1.clip(0, orig_w), y1.clip(0, orig_h), + x2.clip(0, orig_w), y2.clip(0, orig_h)], 1) + + # class-aware NMS + order = conf.argsort()[::-1] + result = [] + while order.size: + i = order[0] + result.append(i) + if order.size == 1: + break + rest = order[1:] + same = cls[rest] == cls[i] + xx1 = np.maximum(boxes[i, 0], boxes[rest, 0]) + yy1 = np.maximum(boxes[i, 1], boxes[rest, 1]) + xx2 = np.minimum(boxes[i, 2], boxes[rest, 2]) + yy2 = np.minimum(boxes[i, 3], boxes[rest, 3]) + inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1) + area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1]) + area_r = (boxes[rest, 2] - boxes[rest, 0]) * (boxes[rest, 3] - boxes[rest, 1]) + iou = inter / (area_i + area_r - inter + 1e-9) + order = rest[~(same & (iou > IOU_THRES))] + return [(boxes[i], conf[i], cls[i]) for i in result] + + +# -------------------------------------------------------------------------- + +def main(): + img_path = Path(sys.argv[1]) if len(sys.argv) > 1 else ASSETS / "bus.jpg" + model_path = ASSETS / "yolov8n.onnx" + for p in (img_path, model_path): + if not p.exists(): + sys.exit(f"missing {p} — see the docstring for setup") + + img = np.asarray(Image.open(img_path).convert("RGB")) + H, W = img.shape[:2] + print(f"image: {img_path.name} {W}x{H}") + + session = ort.InferenceSession(str(model_path), + providers=["CPUExecutionProvider"]) + pipe = build_preprocessor() + frame = upload_image(img) + + # warmup (compiles the fused kernel once; cached on disk afterwards) + t0 = time.perf_counter() + blob = gpu_preprocess(pipe, frame) + t_first = (time.perf_counter() - t0) * 1e3 + + # steady-state preprocessing timing + t0 = time.perf_counter() + N = 50 + for _ in range(N): + blob = gpu_preprocess(pipe, frame) + t_pre = (time.perf_counter() - t0) * 1e3 / N + + t0 = time.perf_counter() + pred = session.run(None, {"images": blob})[0] + t_inf = (time.perf_counter() - t0) * 1e3 + + dets = postprocess(pred, W, H) + print(f"\npreprocess (fused GPU kernel + D2H): first {t_first:.1f} ms, " + f"steady {t_pre:.2f} ms") + print(f"inference (onnxruntime CPU): {t_inf:.1f} ms") + print(f"\n{len(dets)} detections:") + for box, conf, c in dets: + x1, y1, x2, y2 = (int(v) for v in box) + print(f" {COCO[c]:<14} {conf:.2f} [{x1},{y1} -> {x2},{y2}]") + + persons = sum(1 for _, _, c in dets if COCO[c] == "person") + buses = sum(1 for _, _, c in dets if COCO[c] == "bus") + print(f"\nsanity: {persons} persons + {buses} bus " + f"({'OK - expected for bus.jpg' if persons >= 3 and buses >= 1 else 'check image'})") + + +if __name__ == "__main__": + main() diff --git a/examples/README.md b/examples/README.md index 93349b8..e4da7a1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -31,6 +31,9 @@ disk cache in `~/.cache/fkl/` and start instantly. | `06_warping_and_borders.py` | affine rotation / perspective warp (animatable, no recompile), border policies on OOB crops | | `07_torch_interop_dlpack.py` | torch/cupy inputs, DLPack zero-copy outputs, external streams, preallocated outputs | | `08_performance.py` | cold compile vs cache hit vs hot launch; fused chain vs 6 separate kernels | +| `09_video_temporal_window.py` | CircularTensor: rolling N-frame window for temporal models, one fused update per frame | +| `10_torch_realworld.py` | torch-specific recipe (needs torch+cuda installed) | +| `11_yolo_inference.py` | **REAL end-to-end**: JPEG -> ONE fused GPU kernel (letterbox+normalize+CHW) -> YOLOv8n ONNX -> NMS -> boxes. Run `./fetch_yolo_assets.sh` first; needs `pip install numpy pillow onnxruntime` | ## The one-page mental model diff --git a/examples/fetch_yolo_assets.sh b/examples/fetch_yolo_assets.sh new file mode 100755 index 0000000..1257303 --- /dev/null +++ b/examples/fetch_yolo_assets.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Fetch the assets for examples/11_yolo_inference.py: +# bus.jpg ultralytics sample image (810x1080) +# yolov8n.onnx COCO-pretrained YOLOv8n exported to ONNX (1x3x640x640) +set -euo pipefail +cd "$(dirname "$0")/assets" 2>/dev/null || { mkdir -p "$(dirname "$0")/assets"; cd "$(dirname "$0")/assets"; } + +if [ ! -f bus.jpg ]; then + curl -sL -o bus.jpg \ + https://raw.githubusercontent.com/ultralytics/ultralytics/main/ultralytics/assets/bus.jpg + echo "downloaded bus.jpg" +fi + +if [ ! -f yolov8n.onnx ]; then + curl -sL -o yolov8n.onnx \ + "https://huggingface.co/salim4n/yolov8n-detect-onnx/resolve/main/yolov8n-onnx-web/yolov8n.onnx" + echo "downloaded yolov8n.onnx" +fi + +python3 - <<'EOF' +import struct +ok = open("yolov8n.onnx", "rb").read(4)[:2] == b"\x08\x08" +print("yolov8n.onnx looks like ONNX protobuf:", ok) +EOF +echo "assets ready" From eb93fd12c23b8bd7a17a61a1fe46e2ef46150078 Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Thu, 11 Jun 2026 16:08:43 +0200 Subject: [PATCH 2/3] CircularTensor: fix size-baked-in-cache bug; add race audit + coherence tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit REAL BUG found while hardening: the generated ct_update/ct_snapshot baked self.width/self.height as C literals, but the compile-cache signature only contains dtype/batch/order/layout/chain — NOT the size. Two CircularTensors with the same signature and different sizes would cache-hit the first .so and run with the WRONG hardcoded grid and snapshot byte count (out-of-bounds reads/writes; surfaced as 'invalid argument' in an unrelated test). Fix: PyCT now stores w/h at ct_create time and the kernel launch + snapshot use them at runtime, matching the values-vs-types rule everywhere else. CODEGEN_VERSION 6->7. New CircularTensor tests (21 checks total now): - t_race_audit: proves the divergent update is race-free by construction -- the temp slot written by seq1 is never among the slots read by seq2, for both orders and every rotation index. - t_temp_data_consistency: temp/data double-buffer coherence across 3 full rotation wraps, verified plane-by-plane via ct_snapshot_temp. - t_generate_source_no_compile: new generate_source() API (codegen without compiling) checked for all structural pieces. - t_size_not_in_signature_regression: the exact bug scenario. CI: the codegen-compile-check job now also compiles CircularTensor sources (packed/planar, newest/oldest, with preproc chains) via generate_source(). --- .github/workflows/wheels.yml | 28 +++++++++-- fkl/circular.py | 79 +++++++++++++++++++------------ fkl/codegen.py | 2 +- tests/test_circular_tensor.py | 87 +++++++++++++++++++++++++++++++++++ 4 files changed, 161 insertions(+), 35 deletions(-) diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml index 8fce8e7..01d3d2a 100644 --- a/.github/workflows/wheels.yml +++ b/.github/workflows/wheels.yml @@ -115,9 +115,8 @@ jobs: fklroot = os.environ["FKL_ROOT"] inc = os.environ["FKL_INCLUDE"] arch = os.environ.get("FKL_ARCH", "sm_90") - failed = [] - for name, ops, dt, shape, n in CASES: - cu = generate_cu(ops, dt, shape, n) + + def compile_cu(name, cu): with tempfile.NamedTemporaryFile("w", suffix=".cu", delete=False) as f: f.write(cu); path = f.name r = subprocess.run( @@ -129,9 +128,30 @@ jobs: print(f"{name}: {status}") if r.returncode != 0: print(r.stderr[-3000:]) + return r.returncode == 0 + + failed = [] + for name, ops, dt, shape, n in CASES: + if not compile_cu(name, generate_cu(ops, dt, shape, n)): failed.append(name) + + # CircularTensor: stateful, separate codegen path (PyCT + divergent + # launch). Compile-check both layouts/orders + a preproc chain. + CT_CASES = [ + ("ct_packed_f32", fkl.CircularTensor(32, 16, batch=4, dtype="float32"), + []), + ("ct_packed_oldest", fkl.CircularTensor(32, 16, batch=3, + dtype="float32", order="oldest_first"), [fkl.Mul(2.0)]), + ("ct_planar_u8c3", fkl.CircularTensor(16, 8, batch=2, dtype="uint8", + channels=3, layout="planar", out_dtype="float32"), + [fkl.Cast("float32"), fkl.Div(255.0)]), + ] + for name, ct, ops in CT_CASES: + if not compile_cu(name, ct.generate_source(ops)): + failed.append(name) + assert not failed, f"codegen compile failures: {failed}" - print("all codegen cases compiled") + print("all codegen cases compiled (chains + CircularTensor)") EOF release: diff --git a/fkl/circular.py b/fkl/circular.py index 042c06a..a7e6451 100644 --- a/fkl/circular.py +++ b/fkl/circular.py @@ -2,15 +2,25 @@ The only FKL feature that is NOT expressible as a chain op: it owns GPU state (the rolling tensor + a temp tensor + the rotation index) that -persists across calls. Python holds an opaque handle to a heap-allocated -C++ fk::CircularTensor; every update() is ONE fused kernel that: - - 1. runs your preprocessing chain on the incoming frame and writes it - into the correct rotation slot, and - 2. re-orders/copies the other BATCH-1 planes, - -both as a Divergent HF kernel (two sequences selected per plane), exactly -like the C++ `CircularTensor::update`. +persists across calls. The generated .so holds a heap-allocated PyCT +(data Tensor + temp Tensor + rotation idx) mirroring fk::CircularTensor's +double-buffer design; every update() is ONE fused kernel that: + + 1. runs your preprocessing chain on the incoming frame, writes it into + the temp tensor's rotation slot (CircularTensorWrite) AND into the + newest data plane, and + 2. rotate-copies the other BATCH-1 planes from temp into data + (CircularTensorRead with the order's direction), + +as a Divergent HF kernel: SequenceSelectorType routes plane z to sequence +1 (update) or 2 (copy), grid.z = BATCH exactly. + +Why not call fk::CircularTensor::update directly: the Executor's +getActiveThreads SUMS the z extents of all sequences (1 frame-read + +BATCH circular-read = BATCH+1 planes) and the extra plane writes past the +data tensor. Race-safety of our launch is auditable: the temp slot written +by seq1 is never among the slots read by seq2 in the same launch, for both +orders and every rotation index (see test_circular_tensor.py). Usage: ct = fkl.CircularTensor(64, 48, batch=4, dtype="uint8", channels=3, @@ -27,13 +37,11 @@ import ctypes from .backend import get_backend -from .codegen import plan, CODEGEN_VERSION -from .operations import READ, WRITE +from .codegen import CODEGEN_VERSION +from .operations import READ, WRITE, ChainState from .tensor import DeviceBuffer, as_device_view, stream_handle from .types import dtype as _dtype -_ARCH_ENV = "FKL_ARCH" - _ORDERS = {"newest_first": "CircularTensorOrder::NewestFirst", "oldest_first": "CircularTensorOrder::OldestFirst"} _LAYOUTS = ("packed", "planar") @@ -131,14 +139,13 @@ def __del__(self): # ---- compilation -------------------------------------------------------- - def _ensure_compiled(self, ops): - from .jit import _ARCH - states = [] + def _plan_chain(self, ops): + """Thread dtype/shape through the preproc chain; returns (states, key). + Raises if the chain output doesn't match the declared store dtype.""" + states, toks = [], [] dt = self.in_dtype shape = (self.width, self.height, 1) - toks = [] for op in ops: - from .operations import ChainState st = ChainState(dt, *shape) states.append(st) if hasattr(op, "bind"): @@ -149,7 +156,18 @@ def _ensure_compiled(self, ops): if str(dt) != str(self.store_dtype): raise TypeError(f"preproc chain outputs {dt}, CircularTensor " f"stores {self.store_dtype}") - key = "|".join(toks) + return states, "|".join(toks) + + def generate_source(self, ops=None): + """Return the generated .cu for this CircularTensor + preproc chain + WITHOUT compiling (debugging / CI compile checks).""" + ops = list(ops or []) + states, _ = self._plan_chain(ops) + return self._generate_cu(ops, states) + + def _ensure_compiled(self, ops): + from .jit import _ARCH + states, key = self._plan_chain(ops) if self._lib is not None and key == self._chain_key: return if self._lib is not None: @@ -195,14 +213,14 @@ def _generate_cu(self, ops, states) -> str: w_op = f"TensorSplit<{store_t}>" r_op = f"TensorPack<{store_t}>" cp = ch - snap_bytes = (f"(size_t){self.width} * {self.height} * {B} * " - f"{ch} * sizeof({base_t})") + snap_bytes = f"(size_t)ct->w * ct->h * {B} * {ch} * sizeof({base_t})" + else: w_op = f"TensorWrite<{store_t}>" r_op = f"TensorRead<{store_t}>" cp = 1 - snap_bytes = (f"(size_t){self.width} * {self.height} * {B} * " - f"sizeof({store_t})") + snap_bytes = f"(size_t)ct->w * ct->h * {B} * sizeof({store_t})" + tensor_t = base_t if self.layout == "planar" else store_t return f"""// AUTO-GENERATED by fkl-python (CircularTensor). Host+device TU. @@ -230,10 +248,11 @@ def _generate_cu(self, ops, states) -> str: struct PyCT {{ TensorTy data; TensorTy temp; + uint w, h; int idx; - PyCT(uint w, uint h) - : data(w, h, B, {cp}, MemType::Device), - temp(w, h, B, {cp}, MemType::Device), idx(0) {{}} + PyCT(uint w_, uint h_) + : data(w_, h_, B, {cp}, MemType::Device), + temp(w_, h_, B, {cp}, MemType::Device), w(w_), h(h_), idx(0) {{}} }}; template @@ -256,8 +275,8 @@ def _generate_cu(self, ops, states) -> str: void ct_update(void* h, void* d_frame, const float* params, void* ext_stream) {{ PyCT* ct = (PyCT*)h; - Ptr2D<{in_t}> frame(({in_t}*)d_frame, (uint){self.width}, (uint){self.height}, - (uint)({self.width} * sizeof({in_t})), MemType::Device); + Ptr2D<{in_t}> frame(({in_t}*)d_frame, ct->w, ct->h, + (uint)(ct->w * sizeof({in_t})), MemType::Device); MidWrite> toTemp; toTemp.params.first = ct->idx; @@ -276,11 +295,11 @@ def _generate_cu(self, ops, states) -> str: cudaStream_t s; if (ext_stream != nullptr) {{ s = reinterpret_cast(ext_stream); - launch_ct(seqUpdate, seqCopy, s, (uint){self.width}, (uint){self.height}); + launch_ct(seqUpdate, seqCopy, s, ct->w, ct->h); }} else {{ static Stream stream; s = stream.getCUDAStream(); - launch_ct(seqUpdate, seqCopy, s, (uint){self.width}, (uint){self.height}); + launch_ct(seqUpdate, seqCopy, s, ct->w, ct->h); stream.sync(); }} ct->idx = (ct->idx + 1) % B; diff --git a/fkl/codegen.py b/fkl/codegen.py index 11b8c1a..c9a90c7 100644 --- a/fkl/codegen.py +++ b/fkl/codegen.py @@ -21,7 +21,7 @@ from .types import DType # bump when generate_cu's emitted C++ changes for the SAME signature inputs -CODEGEN_VERSION = 6 +CODEGEN_VERSION = 7 def plan(ops: List[Op], in_dtype: DType, in_shape: Tuple[int, int, int], diff --git a/tests/test_circular_tensor.py b/tests/test_circular_tensor.py index cd4107b..651ac4f 100644 --- a/tests/test_circular_tensor.py +++ b/tests/test_circular_tensor.py @@ -108,6 +108,91 @@ def t_uchar3_planar_dnn_window(): tol=1e-6) +def t_race_audit(): + """Static audit of the divergent launch: within one update() kernel the + temp slot WRITTEN by seq1 must never be among the slots READ by seq2, + for both orders and every rotation index. This is what makes the + single-kernel update race-free (no __threadfence/global sync needed).""" + B = 4 + for order in ("newest_first", "oldest_first"): + for idx in range(B): + if order == "newest_first": + # selector: z==0 -> seq1 ; CTRead Descendent: slot = idx - z + seq1_z = {0} + write_slot = (0 + idx) % B # CTWrite Ascendent @ z=0 + read_slots = {(idx - z) % B for z in range(B) if z not in seq1_z} + else: + # selector: z==B-1 -> seq1 ; CTRead Ascendent: slot = idx + z + seq1_z = {B - 1} + write_slot = (B - 1 + idx) % B + read_slots = {(idx + z) % B for z in range(B) if z not in seq1_z} + check_true(f"race audit {order} idx={idx}", + write_slot not in read_slots, + f"w={write_slot} r={sorted(read_slots)}") + + +def t_temp_data_consistency(): + """After every update, temp must hold ALL frames of the window (in slot + space) and data must be temp re-ordered by the rotation — i.e. the + double-buffer never tears across many rotations (3 full wraps).""" + import ctypes + W, H, B = 8, 2, 3 + ct = fkl.CircularTensor(W, H, batch=B, dtype="float32") + n = W * H + for k in range(1, 3 * B + 1): # 9 pushes = 3 wraps + ct.update(_frame_f32(k, W, H)) + data = unf32(ct.snapshot().copy_to_host(), n * B) + tmp_buf = fkl.DeviceBuffer(W, H, "float32", planes=B) + ct._lib.ct_snapshot_temp(ct._handle, ctypes.c_void_p(tmp_buf.ptr), None) + temp = unf32(tmp_buf.copy_to_host(), n * B) + + idx = ct.frames_pushed % B # next write slot + # newest_first: data plane z == temp slot (idx-1-z) mod B + expect_data = [] + for z in range(B): + slot = (idx - 1 - z) % B + expect_data.extend(temp[slot * n:(slot + 1) * n]) + if data != expect_data: + check_true(f"temp/data coherent after push {k}", False) + return + check_true("temp/data coherent across 3 wraps (9 pushes)", True) + + +def t_generate_source_no_compile(): + """generate_source() returns valid-looking CUDA without compiling.""" + ct = fkl.CircularTensor(16, 8, batch=2, dtype="float32") + src = ct.generate_source([fkl.Mul(2.0)]) + needed = ("PyCT", "CircularTensorWrite", "CircularTensorRead", + "SequenceSelectorType", "launchDivergentBatchTransformDPP_Kernel", + "ct_update", "ct_snapshot", "ct_destroy") + missing = [s for s in needed if s not in src] + check_true("generate_source has all structural pieces", + not missing and ct._lib is None, f"missing={missing}") + + +def t_size_not_in_signature_regression(): + """REGRESSION (real bug): two CTs with identical signature + (dtype/B/order/layout/chain) but DIFFERENT sizes. The generated .so must + serve both — sizes are runtime values held in PyCT, never baked + literals. The bug: cache-hit on the first CT's .so with the first CT's + W/H hardcoded -> wrong grid + wrong snapshot byte count.""" + a = fkl.CircularTensor(6, 5, batch=3, dtype="float32") + for k in (1, 2, 3): + a.update(_frame_f32(k, 6, 5)) + b = fkl.CircularTensor(8, 2, batch=3, dtype="float32") # same sig, new size + for k in (10, 20, 30): + b.update(_frame_f32(k, 8, 2)) + + va = unf32(a.snapshot().copy_to_host(), 6 * 5 * 3) + vb = unf32(b.snapshot().copy_to_host(), 8 * 2 * 3) + ok_a = [va[p * 30] for p in range(3)] == [3.0, 2.0, 1.0] and all( + va[p * 30 + i] == va[p * 30] for p in range(3) for i in range(30)) + ok_b = [vb[p * 16] for p in range(3)] == [30.0, 20.0, 10.0] and all( + vb[p * 16 + i] == vb[p * 16] for p in range(3) for i in range(16)) + check_true("same-signature CTs with different sizes coexist", + ok_a and ok_b, f"a={[va[p*30] for p in range(3)]} b={[vb[p*16] for p in range(3)]}") + + def t_wrong_input_rejected(): W, H = 6, 4 ct = fkl.CircularTensor(W, H, batch=2, dtype="float32") @@ -127,4 +212,6 @@ def t_wrong_input_rejected(): run([t_basic_rotation_newest_first, t_basic_rotation_oldest_first, t_preproc_chain_fused, t_pixel_content_preserved, t_values_change_no_recompile, t_uchar3_planar_dnn_window, + t_race_audit, t_temp_data_consistency, t_generate_source_no_compile, + t_size_not_in_signature_regression, t_wrong_input_rejected], "circular-tensor") From 22ad4a56f7d20c518127c063612946a651119834 Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Thu, 11 Jun 2026 16:33:24 +0200 Subject: [PATCH 3/3] CPU backend, ThreadFusion, multi-GPU device routing, torch integration tests Closes out the pending-work checklist: - CPU backend (target='cpu'): generates a plain C++ TU running TransformDPP on MemType::Host pointers, compiled by clang++ (g++ rejects Stream_'s ctor spelling), numpy arrays in/out. 5 tests incl. CPU==GPU cross-validation. - ThreadFusion (thread_fusion=True): emits TransformDPP. Auto-falls back to scalar for shapes whose row bytes are not 16-aligned: external tight-pitch pointers fault on vectorized loads (found via real 'misaligned address' on a 251-wide image). 5 tests: exactness vs scalar baseline, odd widths, uchar3, distinct cache keys. - Multi-GPU: DeviceBuffer/CircularTensor take device=N (per-device primary contexts retained via the driver API); compose pipelines follow the input's device (torch cuda:N / cupy / DeviceBuffer); DLPack device id propagated. Honest caveat: single-GPU box, exercised on device 0 with clean errors for absent devices. - Torch integration suite (torch 2.12+cu130, real GPU): zero-copy both directions (mutation-visibility proven), external torch streams, out=, uchar3 HWC->CHW vs torch reference, CircularTensor->torch window, batch HF from torch list. Found+fixed _alloc_out shape bug for TensorSplit chains (was returning HWC-shaped tensor for planar CHW output). - fk::Equal / Tuple-input ops: documented as structural limit (needs a multi-source Read producing Tuple; upstream feature candidate). - spikes/repro_divergent_sum_z.cu: reproducer behind upstream issue #250 (Divergent Executor launches sum-of-z planes; compute-sanitizer verified OOB + canary corruption). Full matrix: 13 suites x 2 backends (nvcc+clang) + cpu + torch, all green. --- README.md | 48 ++++++++++++++- fkl/backend.py | 24 +++++++- fkl/circular.py | 14 ++++- fkl/codegen.py | 55 +++++++++++++---- fkl/jit.py | 100 +++++++++++++++++++++++++++---- fkl/tensor.py | 60 ++++++++++++++----- spikes/repro_divergent_sum_z.cu | 60 +++++++++++++++++++ tests/test_cpu_backend.py | 77 ++++++++++++++++++++++++ tests/test_thread_fusion.py | 60 +++++++++++++++++++ tests/test_torch_integration.py | 103 ++++++++++++++++++++++++++++++++ 10 files changed, 556 insertions(+), 45 deletions(-) create mode 100644 spikes/repro_divergent_sum_z.cu create mode 100644 tests/test_cpu_backend.py create mode 100644 tests/test_thread_fusion.py create mode 100644 tests/test_torch_integration.py diff --git a/README.md b/README.md index cbf9968..0e06298 100644 --- a/README.md +++ b/README.md @@ -83,10 +83,54 @@ backends (clang with auto-shims, nvcc): | test_warping_splitwrite | 8 | PASS | PASS | | test_niche_ops | 9 | PASS | PASS | | test_dlpack | 9 | PASS | PASS | -| test_circular_tensor | 10 | PASS | PASS | +| test_circular_tensor | 21 | PASS | PASS | +| test_thread_fusion | 5 | PASS | PASS | | test_e2e (timing/cache) | 1 | PASS | PASS | -112 checks per backend. Steady-state hot launch ~70 µs/call; cache-hit +Plus, separately: test_torch_integration (9 checks, real torch 2.12+cu130: +zero-copy both ways, external streams, CircularTensor->torch) and +test_cpu_backend (5 checks, ParArch::CPU with numpy in/out, CPU==GPU +cross-validated). + +## CPU backend (no CUDA required) + +```python +pipe = fkl.compose(fkl.TensorRead(), fkl.Mul(2.0), fkl.TensorWrite(), + target="cpu") +out = pipe(numpy_array) # numpy in -> numpy out, ParArch::CPU +``` +Same fused chains, FKL's CPU executor, plain C++ .so (clang++; g++ rejects +Stream_'s ctor spelling). Results cross-validated against +the GPU backend. + +## ThreadFusion + +```python +pipe = fkl.compose(..., thread_fusion=True) # TF::ENABLED, GPU only +``` +Vectorized multi-element threads. Auto-falls back to scalar for shapes +whose row bytes aren't 16-aligned (external tight-pitch pointers would +fault on float4 loads). + +## Multi-GPU + +`DeviceBuffer(..., device=N)`, `CircularTensor(..., device=N)`; `compose` +pipelines follow the input tensor's device (torch `cuda:N`, cupy device id, +DeviceBuffer.device). DLPack exports carry the right device id. NOTE: this +box has 1 GPU — device routing is implemented and exercised on device 0; +true multi-GPU runs still need a 2+ GPU machine. + +## Known structural limits + +- `fk::Equal` and other Tuple-input ops need a multi-source read (a read + producing `Tuple` from two pointers). FKL has no such Read op yet — + upstream feature candidate, not wrappable from here. +- Divergent HF goes through a direct kernel launch (not the Executor) + because of upstream issue + [#250](https://github.com/Libraries-Openly-Fused/FusedKernelLibrary/issues/250) + (grid.z = sum of sequence z-extents). + +128 checks per backend (+ 9 torch + 5 cpu separately). Steady-state hot launch ~70 µs/call; cache-hit compose ~0 ms (lazy); cold compile ~1-3 s once per chain signature. ## Temporal video: CircularTensor diff --git a/fkl/backend.py b/fkl/backend.py index 9cec033..b10fbb1 100644 --- a/fkl/backend.py +++ b/fkl/backend.py @@ -61,6 +61,7 @@ def _fkl_version() -> str: class CompilerBackend: CLANG = "clang" NVCC = "nvcc" + CPU = "cpu" # plain C++ (clang++/g++), ParArch::CPU, no CUDA at all def __init__(self, kind: str | None = None): self.kind = kind or self._auto() @@ -108,8 +109,12 @@ def compile(self, cu_source: str, sig: str, force: bool = False) -> Path: src_path = so.with_suffix(".cu") src_path.write_text(cu_source) - cmd = self._clang_cmd(src_path, so) if self.kind == self.CLANG \ - else self._nvcc_cmd(src_path, so) + if self.kind == self.CPU: + cmd = self._cpu_cmd(src_path, so) + elif self.kind == self.CLANG: + cmd = self._clang_cmd(src_path, so) + else: + cmd = self._nvcc_cmd(src_path, so) proc = subprocess.run(cmd, capture_output=True, text=True) # Automatic fallback: clang's CUDA support can lag the installed CUDA @@ -139,6 +144,21 @@ def compile(self, cu_source: str, sig: str, force: bool = False) -> Path: self._remember() return so + def _cpu_cmd(self, src: Path, so: Path): + # FKL's CPU path compiles as plain C++. NOTE: g++ rejects the + # Stream_() constructor spelling in stream.h; clang++ + # accepts it, so CPU mode requires clang++. + cxx = shutil.which("clang++") + if cxx is None: + raise RuntimeError("CPU backend requires clang++ (g++ chokes on " + "Stream_'s ctor spelling)") + cpp = src.with_suffix(".cpp") + if not cpp.exists(): + cpp.write_text(src.read_text()) + return [cxx, f"-std={_STD}", "-O2", "-shared", "-fPIC", + "-I", _FKL_INCLUDE, "-I", _FKL_ROOT, + str(cpp), "-o", str(so)] + def _nvcc_cmd(self, src: Path, so: Path): nvcc = str(Path(_CUDA_HOME) / "bin" / "nvcc") return [ diff --git a/fkl/circular.py b/fkl/circular.py index a7e6451..48bccb4 100644 --- a/fkl/circular.py +++ b/fkl/circular.py @@ -51,7 +51,7 @@ class CircularTensor: def __init__(self, width: int, height: int, batch: int, dtype="uint8", channels: int = 1, order: str = "newest_first", layout: str = "packed", - out_dtype=None): + out_dtype=None, device: int = 0): if order not in _ORDERS: raise ValueError(f"order must be one of {tuple(_ORDERS)}") if layout not in _LAYOUTS: @@ -67,6 +67,7 @@ def __init__(self, width: int, height: int, batch: int, self.store_dtype = self.store_dtype.with_channels(self.in_dtype.channels) self.width, self.height, self.batch = int(width), int(height), int(batch) self.order, self.layout = order, layout + self.device = int(device) self._lib = None self._handle = None self._chain_key = None @@ -89,6 +90,9 @@ def update(self, frame, ops=None, stream=None): raise TypeError(f"frame dtype {v.dtype} != declared {self.in_dtype}") self._ensure_compiled(ops) + if self.device != 0: + DeviceBuffer._load_driver() + DeviceBuffer._activate(self.device) params = [] shape = (self.width, self.height, 1) dt = self.in_dtype @@ -111,10 +115,11 @@ def snapshot(self, stream=None): if self.layout == "planar": out = DeviceBuffer(self.width, self.height, self.store_dtype.base, channels=1, - planes=self.batch * ch) + planes=self.batch * ch, device=self.device) else: out = DeviceBuffer(self.width, self.height, self.store_dtype.base, - channels=ch, planes=self.batch) + channels=ch, planes=self.batch, + device=self.device) self._lib.ct_snapshot(self._handle, ctypes.c_void_p(out.ptr), ctypes.c_void_p(stream_handle(stream))) return out @@ -188,6 +193,9 @@ def _ensure_compiled(self, ops): lib.ct_snapshot.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] lib.ct_destroy.argtypes = [ctypes.c_void_p] + if self.device != 0: + DeviceBuffer._load_driver() + DeviceBuffer._activate(self.device) # PyCT allocates on current dev self._handle = lib.ct_create(self.width, self.height) if not self._handle: raise RuntimeError("ct_create failed") diff --git a/fkl/codegen.py b/fkl/codegen.py index c9a90c7..cf6b732 100644 --- a/fkl/codegen.py +++ b/fkl/codegen.py @@ -67,7 +67,8 @@ def collect_params(ops: List[Op], in_dtype: DType, def generate_cu(ops: List[Op], in_dtype: DType, - in_shape: Tuple[int, int, int], n_inputs: int = 1) -> str: + in_shape: Tuple[int, int, int], n_inputs: int = 1, + target: str = "gpu", thread_fusion: bool = False) -> str: states, out_st = plan(ops, in_dtype, in_shape, n_inputs) in_st = states[0] @@ -95,6 +96,9 @@ def generate_cu(ops: List[Op], in_dtype: DType, read_op, write_op = ops[0], ops[-1] in_t, out_t = in_st.dtype.ctype, out_st.dtype.ctype + MEM = "MemType::Device" if target == "gpu" else "MemType::Host" + DPP = ("TransformDPP" if thread_fusion + else "TransformDPP<>") # ---- IO construction (host) ---- if read_op.name == "ReadSet": @@ -107,26 +111,26 @@ def generate_cu(ops: List[Op], in_dtype: DType, packed_t = ops[0].out_dtype(in_st.dtype).ctype in_decl = (f"Tensor<{base_in}> input(({base_in}*)d_in, (uint)dims->in_w, " f"(uint)dims->in_h, (uint)dims->in_planes, " - f"(uint){ops[0]._ch}, MemType::Device);") + f"(uint){ops[0]._ch}, {MEM});") read_expr = f"TensorPack<{packed_t}>::build(input)" elif n_inputs > 1: # HF over a batch of images: d_in is void** (array of device ptrs). # std::array, B> -> PerThreadRead inside FKL. elems = ", ".join( f"Ptr2D<{in_t}>((({in_t}**)d_in)[{i}], (uint)dims->in_w, " - f"(uint)dims->in_h, (uint)(dims->in_w * sizeof({in_t})), MemType::Device)" + f"(uint)dims->in_h, (uint)(dims->in_w * sizeof({in_t})), {MEM})" for i in range(n_inputs)) in_decl = (f"const std::array, {n_inputs}> " f"input{{ {elems} }};") read_expr = f"PerThreadRead::build(input)" elif in_st.planes > 1: in_decl = (f"Tensor<{in_t}> input(({in_t}*)d_in, (uint)dims->in_w, " - f"(uint)dims->in_h, (uint)dims->in_planes, 1, MemType::Device);") + f"(uint)dims->in_h, (uint)dims->in_planes, 1, {MEM});") read_expr = f"TensorRead<{in_t}>::build(input)" else: in_decl = (f"Ptr2D<{in_t}> input(({in_t}*)d_in, (uint)dims->in_w, " f"(uint)dims->in_h, (uint)(dims->in_w * sizeof({in_t})), " - f"MemType::Device);") + f"{MEM});") read_expr = f"PerThreadRead::build(input)" if write_op.name == "TensorSplit": @@ -134,7 +138,7 @@ def generate_cu(ops: List[Op], in_dtype: DType, # planes = batch (thread.z); color_planes = channels (split offsets) out_decl = (f"Tensor<{base_t}> output(({base_t}*)d_out, (uint)dims->out_w, " f"(uint)dims->out_h, (uint)dims->out_planes, " - f"(uint){out_st.dtype.channels}, MemType::Device);") + f"(uint){out_st.dtype.channels}, {MEM});") write_expr = f"TensorSplit<{out_t}>::build(output)" elif write_op.name == "TensorTSplit": base_t = out_st.dtype.base_ctype @@ -159,18 +163,18 @@ def generate_cu(ops: List[Op], in_dtype: DType, ptrs = ", ".join( f"Ptr2D<{base_t}>((({base_t}*)d_out) + {c} * {plane}, " f"(uint)dims->out_w, (uint)dims->out_h, " - f"(uint)(dims->out_w * sizeof({base_t})), MemType::Device)" + f"(uint)(dims->out_w * sizeof({base_t})), {MEM})" for c in range(ch)) out_decl = f"const std::vector> output{{ {ptrs} }};" write_expr = f"SplitWrite::build(output)" elif out_st.planes > 1: out_decl = (f"Tensor<{out_t}> output(({out_t}*)d_out, (uint)dims->out_w, " - f"(uint)dims->out_h, (uint)dims->out_planes, 1, MemType::Device);") + f"(uint)dims->out_h, (uint)dims->out_planes, 1, {MEM});") write_expr = f"TensorWrite<{out_t}>::build(output)" else: out_decl = (f"Ptr2D<{out_t}> output(({out_t}*)d_out, (uint)dims->out_w, " f"(uint)dims->out_h, (uint)(dims->out_w * sizeof({out_t})), " - f"MemType::Device);") + f"{MEM});") write_expr = f"PerThreadWrite::build(output)" if fuse_with_read is not None: @@ -192,6 +196,35 @@ def generate_cu(ops: List[Op], in_dtype: DType, else: all_iops = ",\n ".join([read_expr, *iop_exprs, write_expr]) + if target == "cpu": + return f"""// AUTO-GENERATED by fkl-python codegen v2 (CPU backend). Plain C++ TU. +#include +#include +#include + +using namespace fk; + +extern "C" {{ +struct FklDims {{ + int in_w, in_h, in_planes; + int out_w, out_h, out_planes; +}}; + +void fkl_entry(void* d_in, void* d_out, const FklDims* dims, + const float* params, void* ext_stream) +{{ + (void)ext_stream; // CPU executor is synchronous + {in_decl} + {out_decl} + + Stream_ stream; + executeOperations>(stream, + {all_iops}); + stream.sync(); +}} +}} // extern "C" +""" + return f"""// AUTO-GENERATED by fkl-python codegen v2. Single-step host+device TU. #include #include @@ -213,12 +246,12 @@ def generate_cu(ops: List[Op], in_dtype: DType, if (ext_stream != nullptr) {{ Stream stream(reinterpret_cast(ext_stream)); - executeOperations>(stream, + executeOperations<{DPP}>(stream, {all_iops}); // caller owns the stream: stay async }} else {{ static Stream stream; // persistent: avoid create/destroy per call - executeOperations>(stream, + executeOperations<{DPP}>(stream, {all_iops}); stream.sync(); }} diff --git a/fkl/jit.py b/fkl/jit.py index bc09689..2ae2711 100644 --- a/fkl/jit.py +++ b/fkl/jit.py @@ -36,24 +36,50 @@ class _FklDims(ctypes.Structure): class FusedKernel: """A composed chain. Compiles lazily on first call (input dtype/layout - is only known then) and caches per signature.""" + is only known then) and caches per signature. - def __init__(self, ops: List[Op]): + target="gpu" (default) JITs a CUDA .so; target="cpu" JITs a plain C++ + .so running FKL's ParArch::CPU executor on host memory (numpy in/out, + no CUDA required).""" + + def __init__(self, ops: List[Op], target: str = "gpu", + thread_fusion: bool = False): if not ops or ops[0].role != READ: raise ValueError("chain must start with TensorRead()") if ops[-1].role != WRITE: raise ValueError("chain must end with TensorWrite()/TensorSplit()") + if target not in ("gpu", "cpu"): + raise ValueError("target must be 'gpu' or 'cpu'") self.ops = ops + self.target = target + self.thread_fusion = bool(thread_fusion) and target == "gpu" self._variants = {} # signature -> (entry fn, params buffer, out_state) # ---- compile path (cold, once per signature) -------------------------- + def _tf_effective(self, dt: DType, shape) -> bool: + """ThreadFusion vectorizes row accesses (e.g. float4). With external + tight-pitched pointers every row must stay 16-byte aligned, i.e. + width * itemsize % 16 == 0. Otherwise fall back to the scalar DPP + for THIS shape (correctness first; the .so is per-signature).""" + if not self.thread_fusion: + return False + return (shape[0] * dt.itemsize) % 16 == 0 + def _get_variant(self, dt: DType, shape, n_inputs: int = 1): - sig = signature(self.ops, dt, shape, _ARCH, n_inputs) + arch = _ARCH if self.target == "gpu" else "host" + tf = self._tf_effective(dt, shape) + sig = (signature(self.ops, dt, shape, arch, n_inputs) + + f";t={self.target};tf={int(tf)}") hit = self._variants.get(sig) if hit is not None: return hit - cu = generate_cu(self.ops, dt, shape, n_inputs) - so = get_backend().compile(cu, sig) + cu = generate_cu(self.ops, dt, shape, n_inputs, target=self.target, + thread_fusion=tf) + if self.target == "cpu": + from .backend import CompilerBackend + so = CompilerBackend(CompilerBackend.CPU).compile(cu, sig) + else: + so = get_backend().compile(cu, sig) lib = ctypes.CDLL(str(so)) entry = lib.fkl_entry entry.restype = None @@ -69,9 +95,16 @@ def _get_variant(self, dt: DType, shape, n_inputs: int = 1): # ---- HOT PATH ---------------------------------------------------------- def __call__(self, x, out=None, stream=None): + if self.target == "cpu": + return self._call_cpu(x, out) if isinstance(x, (list, tuple)): return self._call_batch(x, out, stream) vin: DeviceView = as_device_view(x) + if vin.device != 0: + # make the input's device current for the launch + output alloc + from .tensor import DeviceBuffer + DeviceBuffer._load_driver() + DeviceBuffer._activate(vin.device) in_shape = (vin.width, vin.height, vin.planes) entry, pbuf, out_st, _ = self._get_variant(vin.dtype, in_shape) @@ -121,28 +154,69 @@ def _call_batch(self, xs, out=None, stream=None): # caller keeping tensors alive for async streams (documented). return out + def _call_cpu(self, x, out=None): + """CPU executor: numpy arrays in/out (host memory, synchronous).""" + import numpy as np + from .types import from_cai + x = np.ascontiguousarray(x) + ai = x.__array_interface__ + dt, w, h, p = from_cai(ai["shape"], ai["typestr"]) + entry, pbuf, out_st, _ = self._get_variant(dt, (w, h, p)) + + ch = out_st.dtype.channels + is_split = self.ops[-1].name in ("TensorSplit", "TensorTSplit") + if is_split: + shape = ((out_st.planes if out_st.planes > 1 else 1) * ch, + out_st.height, out_st.width) + elif out_st.planes > 1: + shape = (out_st.planes, out_st.height, out_st.width) + ((ch,) if ch > 1 else ()) + elif out_st.height > 1: + shape = (out_st.height, out_st.width) + ((ch,) if ch > 1 else ()) + else: + shape = (out_st.width,) + ((ch,) if ch > 1 else ()) + if out is None: + out = np.empty(shape, dtype=out_st.dtype.base) + out = np.ascontiguousarray(out) + + dims = _FklDims(w, h, p, out_st.width, out_st.height, out_st.planes) + entry(ctypes.c_void_p(x.ctypes.data), ctypes.c_void_p(out.ctypes.data), + ctypes.byref(dims), ctypes.cast(pbuf, ctypes.c_void_p), + ctypes.c_void_p(0)) + return out + # ---- output allocation -------------------------------------------------- def _alloc_out(self, vin: DeviceView, out_st): mod = type(vin.obj).__module__.split(".")[0] ch = out_st.dtype.channels - if out_st.planes > 1: + is_split = self.ops[-1].name in ("TensorSplit", "TensorTSplit") + if is_split: + # planar output: channels become leading planes (CHW / NCHW) + nplanes = (out_st.planes if out_st.planes > 1 else 1) * ch + shape = (nplanes, out_st.height, out_st.width) + alloc_dtype = out_st.dtype.base + elif out_st.planes > 1: shape = (out_st.planes, out_st.height, out_st.width) + ((ch,) if ch > 1 else ()) + alloc_dtype = out_st.dtype.base elif out_st.height > 1: shape = (out_st.height, out_st.width) + ((ch,) if ch > 1 else ()) + alloc_dtype = out_st.dtype.base else: shape = (out_st.width,) + ((ch,) if ch > 1 else ()) + alloc_dtype = out_st.dtype.base if mod == "torch": import torch tmap = {"float32": torch.float32, "float64": torch.float64, "uint8": torch.uint8, "int32": torch.int32, "int16": torch.int16, "uint16": torch.uint16, "int8": torch.int8} - return torch.empty(shape, device="cuda", dtype=tmap[out_st.dtype.base]) + return torch.empty(shape, device=f"cuda:{vin.device}", + dtype=tmap[out_st.dtype.base]) if mod == "cupy": import cupy - return cupy.empty(shape, dtype=out_st.dtype.base) + with cupy.cuda.Device(vin.device): + return cupy.empty(shape, dtype=out_st.dtype.base) from .tensor import DeviceBuffer - return DeviceBuffer.from_state(out_st) + return DeviceBuffer.from_state(out_st, device=vin.device) def source_for(self, dtype_spec, shape): """Debug: show the generated C++ for a given input.""" @@ -150,8 +224,12 @@ def source_for(self, dtype_spec, shape): return generate_cu(self.ops, _d(dtype_spec), tuple(shape)) -def compose(*ops: Op) -> FusedKernel: - return FusedKernel(list(ops)) +def compose(*ops: Op, target: str = "gpu", + thread_fusion: bool = False) -> FusedKernel: + """thread_fusion=True enables FKL's ThreadFusion (TF::ENABLED): + each thread processes multiple elements with vectorized accesses. + GPU-only; best for wide images with simple per-pixel chains.""" + return FusedKernel(list(ops), target=target, thread_fusion=thread_fusion) # ===================== Divergent Horizontal Fusion ======================== diff --git a/fkl/tensor.py b/fkl/tensor.py index 115b23f..19a1767 100644 --- a/fkl/tensor.py +++ b/fkl/tensor.py @@ -19,6 +19,7 @@ class DeviceView: planes: int dtype: DType obj: object # keeps the underlying buffer alive + device: int = 0 def as_device_view(x) -> "DeviceView": @@ -31,7 +32,22 @@ def as_device_view(x) -> "DeviceView": # require C-contiguous for now (zero-copy without repacking) raise ValueError("non-contiguous arrays not supported; call .contiguous() first") dt, w, h, p = from_cai(cai["shape"], cai["typestr"]) - return DeviceView(int(cai["data"][0]), w, h, p, dt, x) + return DeviceView(int(cai["data"][0]), w, h, p, dt, x, _device_of(x)) + + +def _device_of(x) -> int: + """Best-effort device index: DeviceBuffer.device, torch .device.index, + cupy .device.id; default 0 (CAI itself carries no device id).""" + d = getattr(x, "device", None) + if d is None: + return 0 + if isinstance(d, int): + return d + idx = getattr(d, "index", None) # torch.device + if idx is not None: + return int(idx) + did = getattr(d, "id", None) # cupy.cuda.Device + return int(did) if did is not None else 0 def stream_handle(stream) -> int: @@ -47,13 +63,18 @@ def stream_handle(stream) -> int: class DeviceBuffer: - """Dependency-free device buffer via the CUDA driver API.""" + """Dependency-free device buffer via the CUDA driver API. + + device=N allocates on GPU N (contexts are retained per device and the + allocation happens with that device's primary context current).""" _cuda = None - _ctx = None + _ctxs = {} def __init__(self, width: int, height: int = 1, dtype="float32", - channels: int = 1, planes: int = 1): + channels: int = 1, planes: int = 1, device: int = 0): self._load_driver() + self.device = int(device) + self._activate(self.device) from .types import dtype as _d base = _d(dtype) if isinstance(dtype, str) else dtype self.dtype = DType(base.base, channels if channels > 1 else base.channels) @@ -64,9 +85,23 @@ def __init__(self, width: int, height: int = 1, dtype="float32", ctypes.c_size_t(self._nbytes))) @classmethod - def from_state(cls, st): + def from_state(cls, st, device: int = 0): return cls(st.width, st.height, st.dtype.base, - channels=st.dtype.channels, planes=st.planes) + channels=st.dtype.channels, planes=st.planes, device=device) + + @classmethod + def _activate(cls, device: int): + """Make `device`'s primary context current for this thread.""" + if device not in cls._ctxs: + dev = ctypes.c_int(0) + if cls._cuda.cuDeviceGet(ctypes.byref(dev), device) != 0: + raise RuntimeError(f"cuDeviceGet({device}) failed — " + f"is GPU {device} present?") + ctx = ctypes.c_void_p() + if cls._cuda.cuDevicePrimaryCtxRetain(ctypes.byref(ctx), dev) != 0: + raise RuntimeError(f"cuDevicePrimaryCtxRetain({device}) failed") + cls._ctxs[device] = ctx + cls._cuda.cuCtxSetCurrent(cls._ctxs[device]) @property def ptr(self) -> int: @@ -87,15 +122,8 @@ def _load_driver(cls): lib.cuMemcpyDtoH_v2.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t] if lib.cuInit(0) != 0: raise RuntimeError("cuInit failed") - dev = ctypes.c_int(0) - if lib.cuDeviceGet(ctypes.byref(dev), 0) != 0: - raise RuntimeError("cuDeviceGet failed") - ctx = ctypes.c_void_p() - if lib.cuDevicePrimaryCtxRetain(ctypes.byref(ctx), dev) != 0: - raise RuntimeError("cuDevicePrimaryCtxRetain failed") - lib.cuCtxSetCurrent(ctx) cls._cuda = lib - cls._ctx = ctx + cls._activate(0) def _check(self, code): if code != 0: @@ -131,7 +159,7 @@ def __cuda_array_interface__(self): # no copy, no sync beyond what the consumer framework inserts. def __dlpack_device__(self): - return (2, 0) # (kDLCUDA, device 0) + return (2, self.device) # (kDLCUDA, device_id) def __dlpack__(self, stream=None): return _make_dlpack_capsule(self) @@ -207,7 +235,7 @@ def _make_dlpack_capsule(buf: "DeviceBuffer"): managed = _DLManagedTensor() managed.dl_tensor.data = ctypes.c_void_p(buf.ptr) - managed.dl_tensor.device = _DLDevice(2, 0) # kDLCUDA, dev 0 + managed.dl_tensor.device = _DLDevice(2, getattr(buf, "device", 0)) # kDLCUDA managed.dl_tensor.ndim = ndim managed.dl_tensor.dtype = _DLDataType(code, bits, 1) managed.dl_tensor.shape = shape_arr diff --git a/spikes/repro_divergent_sum_z.cu b/spikes/repro_divergent_sum_z.cu new file mode 100644 index 0000000..d525629 --- /dev/null +++ b/spikes/repro_divergent_sum_z.cu @@ -0,0 +1,60 @@ +// Reproducer: DivergentBatchTransformDPP Executor launches sum-of-z planes. +// Two sequences, EACH with a full-batch read covering B=3 planes and a +// selector routing plane 0 -> seq1, planes 1,2 -> seq2. Expected: grid.z = 3. +// Observed: grid.z = 6 (sum of both sequences' z extents); planes 3..5 +// execute with selector values beyond the intent and write out of bounds. +#include +#include +#include +#include +using namespace fk; + +struct Selector { + FK_HOST_DEVICE_FUSE uint at(const uint& z) { return z == 0 ? 1u : 2u; } +}; + +int main() { + constexpr int W = 8, H = 4, B = 3; + Stream stream; + + Tensor in(W, H, B, 1, MemType::Device); + Tensor out(W, H, B, 1, MemType::Device); + // canary buffer allocated right after 'out' to detect OOB writes + Tensor canary(W, H, B, 1, MemType::Device); + + std::vector hostIn(W * H * B); + for (int i = 0; i < W * H * B; ++i) hostIn[i] = float(i % 100); + cudaMemcpy(in.ptr().data, hostIn.data(), W * H * B * sizeof(float), cudaMemcpyHostToDevice); + cudaMemset(out.ptr().data, 0, W * H * B * sizeof(float)); + std::vector sentinel(W * H * B, 777.f); + cudaMemcpy(canary.ptr().data, sentinel.data(), W * H * B * sizeof(float), cudaMemcpyHostToDevice); + + const auto seq1 = buildOperationSequence( + TensorRead::build(in), Mul::build(10.f), TensorWrite::build(out)); + const auto seq2 = buildOperationSequence( + TensorRead::build(in), Add::build(1.f), TensorWrite::build(out)); + + Executor>:: + executeOperations(stream, seq1, seq2); + stream.sync(); + + std::vector got(W * H * B), can(W * H * B); + cudaMemcpy(got.data(), out.ptr().data, W * H * B * sizeof(float), cudaMemcpyDeviceToHost); + cudaMemcpy(can.data(), canary.ptr().data, W * H * B * sizeof(float), cudaMemcpyDeviceToHost); + + int wrong = 0, corrupted = 0; + for (int z = 0; z < B; ++z) + for (int i = 0; i < W * H; ++i) { + const float v = got[z * W * H + i]; + const float x = hostIn[z * W * H + i]; + const float expect = z == 0 ? x * 10.f : x + 1.f; + if (v != expect) ++wrong; + } + for (int i = 0; i < W * H * B; ++i) if (can[i] != 777.f) ++corrupted; + + printf("wrong output elements: %d / %d\n", wrong, W * H * B); + printf("canary corrupted elements: %d / %d\n", corrupted, W * H * B); + printf(wrong == 0 && corrupted == 0 ? "PASS (intended semantics)\n" + : "FAIL (sum-of-z launch)\n"); + return wrong == 0 && corrupted == 0 ? 0 : 1; +} diff --git a/tests/test_cpu_backend.py b/tests/test_cpu_backend.py new file mode 100644 index 0000000..02cbd3e --- /dev/null +++ b/tests/test_cpu_backend.py @@ -0,0 +1,77 @@ +"""Tests for the CPU backend (ParArch::CPU): same fused chains, host memory, +numpy in/out, validated against numpy references and against the GPU backend. +Requires numpy + clang++ (no CUDA needed for the CPU path itself).""" +import sys + +try: + import numpy as np +except ImportError: + print("SKIP: numpy not available") + sys.exit(0) + +import fkl +from harness import check_true, run + + +def t_cpu_basic_vf(): + x = np.arange(32, dtype=np.float32).reshape(4, 8) + pipe = fkl.compose(fkl.TensorRead(), fkl.Mul(2.0), fkl.Add(1.0), + fkl.TensorWrite(), target="cpu") + out = pipe(x) + check_true("CPU VF Mul+Add", np.allclose(out, x * 2 + 1)) + + +def t_cpu_uchar3_preproc(): + img = np.random.randint(0, 256, (6, 10, 3), dtype=np.uint8) + pipe = fkl.compose(fkl.TensorRead(), fkl.Cast("float32"), + fkl.Div((255.0, 255.0, 255.0)), fkl.TensorSplit(), + target="cpu") + out = pipe(img) + ref = img.transpose(2, 0, 1).astype(np.float32) / 255.0 + check_true("CPU uchar3 -> CHW float", out.shape == (3, 6, 10) + and np.allclose(out, ref, atol=1e-6)) + + +def t_cpu_crop_resize(): + img = np.random.rand(12, 16).astype(np.float32) * 100 + pipe = fkl.compose(fkl.TensorRead(), fkl.Crop(4, 2, 8, 8), + fkl.Resize(4, 4), fkl.TensorWrite(), target="cpu") + out = pipe(img) + check_true("CPU BVF Crop+Resize shape", out.shape == (4, 4)) + + +def t_cpu_matches_gpu(): + """The exact same chain must produce identical results on both targets.""" + img = np.random.randint(0, 256, (8, 12, 3), dtype=np.uint8) + ops = lambda: (fkl.TensorRead(), fkl.Cast("float32"), + fkl.Mul((2.0, 0.5, 1.0)), fkl.Sub((10.0, 0.0, 5.0)), + fkl.TensorSplit()) + cpu_out = fkl.compose(*ops(), target="cpu")(img) + + try: + gpu_buf = fkl.DeviceBuffer(12, 8, "uint8", channels=3) + except Exception: + check_true("CPU==GPU (skipped: no GPU)", True) + return + gpu_buf.copy_from_host(img.tobytes()) + gpu_out_buf = fkl.compose(*ops())(gpu_buf) + gpu_out = np.frombuffer(gpu_out_buf.copy_to_host(), + dtype=np.float32).reshape(cpu_out.shape) + check_true("CPU == GPU bit-comparable", np.allclose(cpu_out, gpu_out, atol=1e-5)) + + +def t_cpu_values_no_recompile(): + x = np.ones((4, 8), dtype=np.float32) + pipe = fkl.compose(fkl.TensorRead(), fkl.Mul(3.0), fkl.TensorWrite(), + target="cpu") + a = pipe(x) + pipe.ops[1]._v = None # not how values change; use a fresh compose + pipe2 = fkl.compose(fkl.TensorRead(), fkl.Mul(7.0), fkl.TensorWrite(), + target="cpu") + b = pipe2(x) + check_true("CPU values via params", float(a[0, 0]) == 3.0 and float(b[0, 0]) == 7.0) + + +if __name__ == "__main__": + run([t_cpu_basic_vf, t_cpu_uchar3_preproc, t_cpu_crop_resize, + t_cpu_matches_gpu, t_cpu_values_no_recompile], "cpu-backend") diff --git a/tests/test_thread_fusion.py b/tests/test_thread_fusion.py new file mode 100644 index 0000000..5fd04bb --- /dev/null +++ b/tests/test_thread_fusion.py @@ -0,0 +1,60 @@ +"""ThreadFusion (TF::ENABLED) tests: same results, vectorized accesses.""" +import struct +from harness import dev_f32, dev_u8, unf32, check, check_true, run +import fkl + + +def t_tf_correctness_f32(): + W, H = 256, 32 # wide → thread-fusable + src = [float(i % 977) for i in range(W * H)] + ops = (fkl.TensorRead(), fkl.Mul(2.0), fkl.Add(1.0), fkl.TensorWrite()) + base = fkl.compose(*ops)(dev_f32(src, W, H)) + tf = fkl.compose(*ops, thread_fusion=True)(dev_f32(src, W, H)) + check("TF f32 result == baseline", unf32(tf.copy_to_host(), W * H), + unf32(base.copy_to_host(), W * H)) + + +def t_tf_odd_width(): + """Non-divisible width exercises the THREAD_DIVISIBLE=false tail path.""" + W, H = 251, 8 # prime width + src = [float(i) for i in range(W * H)] + ops = (fkl.TensorRead(), fkl.Mul(3.0), fkl.TensorWrite()) + tf = fkl.compose(*ops, thread_fusion=True)(dev_f32(src, W, H)) + check("TF odd width exact", unf32(tf.copy_to_host(), W * H), + [v * 3.0 for v in src]) + + +def t_tf_uchar3(): + W, H = 128, 16 + n = W * H * 3 + vals = [(i * 7) % 256 for i in range(n)] + ops = (fkl.TensorRead(), fkl.Cast("float32"), fkl.Div((255.0,) * 3), + fkl.TensorSplit()) + base = fkl.compose(*ops)(dev_u8(vals, W, H, ch=3)) + tf = fkl.compose(*ops, thread_fusion=True)(dev_u8(vals, W, H, ch=3)) + check("TF uchar3 CHW == baseline", unf32(tf.copy_to_host(), n), + unf32(base.copy_to_host(), n)) + + +def t_tf_distinct_cache_entries(): + ops = (fkl.TensorRead(), fkl.Mul(1.5), fkl.TensorWrite()) + a = fkl.compose(*ops) + b = fkl.compose(*ops, thread_fusion=True) + x = dev_f32([1.0] * 64 * 4, 64, 4) + a(x); b(x) + so_a = list(a._variants.values())[0][3] + so_b = list(b._variants.values())[0][3] + check_true("TF compiles a distinct .so", so_a != so_b) + + +def t_tf_cpu_rejected_gracefully(): + """thread_fusion on CPU target silently disables (no CUDA TF on CPU).""" + p = fkl.compose(fkl.TensorRead(), fkl.Mul(2.0), fkl.TensorWrite(), + target="cpu", thread_fusion=True) + check_true("TF ignored for CPU target", p.thread_fusion is False) + + +if __name__ == "__main__": + run([t_tf_correctness_f32, t_tf_odd_width, t_tf_uchar3, + t_tf_distinct_cache_entries, t_tf_cpu_rejected_gracefully], + "thread-fusion") diff --git a/tests/test_torch_integration.py b/tests/test_torch_integration.py new file mode 100644 index 0000000..1ec1773 --- /dev/null +++ b/tests/test_torch_integration.py @@ -0,0 +1,103 @@ +"""REAL torch integration tests (requires torch+cuda): zero-copy in both +directions, external streams, out= preallocation, CircularTensor->torch. + +Run with a torch-enabled python: + .venv-torch/bin/python tests/test_torch_integration.py +""" +import sys +import fkl + +try: + import torch + assert torch.cuda.is_available() +except Exception as e: + print(f"SKIP: torch+cuda not available ({e})") + sys.exit(0) + +from harness import check, check_true, run, unf32 + + +def t_torch_in_dlpack_out(): + """torch tensor in -> fused kernel -> DeviceBuffer -> torch.from_dlpack.""" + x = torch.full((16, 32), 3.0, device="cuda", dtype=torch.float32) + pipe = fkl.compose(fkl.TensorRead(), fkl.Mul(2.0), fkl.Add(1.0), + fkl.TensorWrite()) + out = pipe(x) # torch in -> torch out + check_true("torch in -> torch out (auto-alloc)", + isinstance(out, torch.Tensor) and out.is_cuda + and torch.allclose(out, torch.full_like(out, 7.0))) + + +def t_dlpack_zero_copy_roundtrip(): + """DeviceBuffer -> torch.from_dlpack shares the SAME memory.""" + buf = fkl.DeviceBuffer(8, 4, "float32") + import struct + buf.copy_from_host(struct.pack("32f", *[float(i) for i in range(32)])) + t = torch.from_dlpack(buf) + check_true("from_dlpack shape+device", tuple(t.shape) == (4, 8) and t.is_cuda) + check("from_dlpack values", t.flatten().tolist(), + [float(i) for i in range(32)]) + # zero-copy proof: mutate via torch, the underlying memory changes + t += 100.0 + t2 = torch.from_dlpack(buf) if not hasattr(buf, "_exported_dlpack") else t + check_true("zero-copy (mutation visible)", float(t.flatten()[0]) == 100.0) + + +def t_torch_stream_async(): + """Launch on a torch stream: async semantics, caller syncs.""" + x = torch.full((64, 64), 1.0, device="cuda") + pipe = fkl.compose(fkl.TensorRead(), fkl.Mul(5.0), fkl.TensorWrite()) + s = torch.cuda.Stream() + with torch.cuda.stream(s): + out = pipe(x, stream=s.cuda_stream) + s.synchronize() + check_true("external torch stream", torch.allclose(out, torch.full_like(out, 5.0))) + + +def t_out_preallocated_torch(): + x = torch.full((8, 8), 2.0, device="cuda") + out = torch.empty(8, 8, device="cuda") + pipe = fkl.compose(fkl.TensorRead(), fkl.Add(3.0), fkl.TensorWrite()) + ret = pipe(x, out=out) + check_true("out= preallocated torch tensor", + ret is out and torch.allclose(out, torch.full_like(out, 5.0))) + + +def t_uchar3_image_preproc_to_torch(): + """The DNN-ingest path: uint8 HWC torch image -> fused preproc -> CHW.""" + img = torch.randint(0, 256, (24, 32, 3), device="cuda", dtype=torch.uint8) + pipe = fkl.compose(fkl.TensorRead(), fkl.Cast("float32"), + fkl.Div((255.0, 255.0, 255.0)), fkl.TensorSplit()) + out = pipe(img) # (3, 24, 32) float + ref = img.permute(2, 0, 1).float() / 255.0 + check_true("uchar3 HWC -> CHW float matches torch reference", + tuple(out.shape) == (3, 24, 32) + and torch.allclose(out, ref, atol=1e-6)) + + +def t_circular_tensor_to_torch(): + """CircularTensor window consumed by torch via DLPack.""" + ct = fkl.CircularTensor(16, 8, batch=3, dtype="float32") + for k in (1.0, 2.0, 3.0): + ct.update(torch.full((8, 16), k, device="cuda")) + win = torch.from_dlpack(ct.snapshot()) # (3, 8, 16) + check_true("CircularTensor -> torch window", + tuple(win.shape) == (3, 8, 16) + and win[0].max().item() == 3.0 and win[1].max().item() == 2.0 + and win[2].max().item() == 1.0) + + +def t_batch_hf_from_torch_list(): + imgs = [torch.full((8, 16), float(k), device="cuda") for k in (1, 2, 3)] + pipe = fkl.compose(fkl.TensorRead(), fkl.Mul(10.0), fkl.TensorWrite()) + out = pipe(imgs) # (3, 8, 16) + check_true("HF batch from list of torch tensors", + tuple(out.shape) == (3, 8, 16) + and out[0].max().item() == 10.0 and out[2].max().item() == 30.0) + + +if __name__ == "__main__": + run([t_torch_in_dlpack_out, t_dlpack_zero_copy_roundtrip, + t_torch_stream_async, t_out_preallocated_torch, + t_uchar3_image_preproc_to_torch, t_circular_tensor_to_torch, + t_batch_hf_from_torch_list], "torch-integration")