PyTorch-style high-level API: Tensor / Image / ImageBatch + fluent fused pipelines

README.md

@@ -56,6 +56,49 @@ out = pipe(x, stream=torch.cuda.current_stream())   # zero-copy, async
 
 No torch? Use the built-in dependency-free `DeviceBuffer` (CUDA driver via ctypes).
 
+## PyTorch-style API: Tensor / Image / pipe
+
+A torch-flavoured layer over the same machinery — for users who think in
+tensors and image batches rather than op chains. Everything is zero-copy and
+every pipeline still lowers to **one** fused kernel:
+
+```python
+import torch, fkl
+
+frames = torch.randint(0, 256, (8, 720, 1280, 3),           # NHWC uint8
+                       device="cuda", dtype=torch.uint8)
+batch = fkl.ImageBatch(frames)          # zero-copy per-plane views
+                                        # (or a list of same-size images)
+
+out = (fkl.pipe(batch)
+       .resize((224, 224))              # BVF: fused into the read
+       .normalize(MEAN, STD)            # Cast+Sub+Div, in registers
+       .to_planar("CHW")                # planar NCHW write
+       .run())                          # ONE kernel for all 8 images
+
+net(out.torch())                        # zero-copy (8, 3, 224, 224) cuda tensor
+```
+
+- `fkl.Tensor(obj)` wraps torch CUDA tensors, `__cuda_array_interface__`
+  objects (cupy/numba) and `DeviceBuffer` without copying; exposes
+  `.shape/.dtype/.device`, `__cuda_array_interface__`, DLPack and
+  `.torch()`. `fkl.Tensor.empty(shape, dtype)` allocates.
+- `fkl.Image` = HWC image semantics (`width/height/channels`, C 1..4);
+  `fkl.ImageBatch` = list of same-shape images OR one batched
+  `(N, H, W[, C])` tensor. Input layout is `'HWC'`; planar `'CHW'` is an
+  OUTPUT layout produced by `.to_planar()` (the existing split ops).
+- `fkl.pipe(x)` methods (`crop`, `resize`, `border`, `cvt_color`,
+  `normalize`, `mul/add/sub/div`, `cast`, `saturate_cast`, `apply(op)`,
+  `to_planar`) each append an existing symbolic op; `.run()` goes through
+  `compose()`'s JIT cache — same rules: types compile once, values travel
+  in `params[]`. Geometry ops come first (they fuse into the read).
+- One-op eager calls: `fkl.F.resize(img, (64, 64))`,
+  `fkl.F.cvt_color(img, "BGR2GRAY")`, ... (each a one-op fused pipeline).
+
+Out of scope by design (use torch itself): autograd, broadcasting,
+>4-channel images (FKL vector pixels are 1..4 channels), multi-input graphs
+(the ABI is single-input/single-output; batches are horizontal fusion).
+
 ## Setup
 
 ### Plug and play (wheel with vendored headers)
@@ -205,11 +248,12 @@ testing the package.
 
 ## Examples
 
-`examples/` — 8 runnable, dependency-free scripts (see `examples/README.md`):
-vertical fusion basics, single-kernel DNN preprocessing, multi-ROI batch
-(HF), multi-camera + Divergent HF, color pipelines, warping + border
-policies, torch/DLPack interop, and a fused-vs-unfused benchmark
-(~5x at 1080p for a 6-op chain).
+`examples/` — 12 runnable scripts, most dependency-free (see
+`examples/README.md`): vertical fusion basics, single-kernel DNN
+preprocessing, multi-ROI batch (HF), multi-camera + Divergent HF, color
+pipelines, warping + border policies, torch/DLPack interop, a
+fused-vs-unfused benchmark (~5x at 1080p for a 6-op chain), temporal video
+windows, YOLO end-to-end, and the torch-style `pipe()` API.
 
 ## Current scope / next steps
 

examples/12_torch_style_api.py

@@ -0,0 +1,65 @@
+"""Example 12 — PyTorch-style API: Tensor / Image / ImageBatch + fkl.pipe().
+
+The classic DNN-ingest chain, written the way a torch user expects:
+
+    batch of HWC uint8 frames -> resize -> normalize -> NCHW float planes
+
+...and it STILL lowers to ONE fused FKL kernel: the fluent methods only
+append the same symbolic ops compose() takes; .run() JIT-compiles once per
+signature (disk cache) and then costs a single launch. resize fuses into the
+read (BVF), normalize runs in registers (VF), the batch is horizontal fusion
+(HF), to_planar is the planar TensorSplit write.
+
+Out of scope by design: autograd, broadcasting, >4-channel images,
+multi-input graphs — this is a preprocessing API layer, not a framework.
+"""
+import fkl
+from _util import gpu_image_u8, to_floats, synthetic_rgb
+
+SRC_W, SRC_H = 320, 240          # camera frames
+NET_W, NET_H = 64, 64            # model input
+MEAN = (123.675, 116.28, 103.53)     # ImageNet mean (RGB)
+STD = (58.395, 57.12, 57.375)
+
+# three same-size frames (deterministic, shifted patterns)
+frames = []
+for k in range(3):
+    pat = synthetic_rgb(SRC_W, SRC_H)
+    frames.append(gpu_image_u8([(v + 40 * k) % 256 for v in pat],
+                               SRC_W, SRC_H, channels=3))
+
+batch = fkl.ImageBatch(frames)       # also takes ONE (N, H, W, C) tensor
+print(f"OK  {batch!r}")
+
+pre = (fkl.pipe(batch)
+       .resize((NET_W, NET_H))       # bilinear, fused into the read (BVF)
+       .normalize(MEAN, STD)         # (x - mean) / std, in registers (VF)
+       .to_planar("CHW"))            # planar NCHW write (TensorSplit)
+
+out = pre.run()                      # ONE kernel launch for all 3 images
+print(f"OK  3x {SRC_W}x{SRC_H} HWC uint8 -> {out!r}")
+assert out.shape == (3, 3, NET_H, NET_W)
+
+plane = to_floats(out.base, 16)
+print(f"    img0 R-plane[:4] = {[round(v, 3) for v in plane[:4]]}")
+
+# Zero-copy hand-off to a DNN runtime: out.torch() gives a (3, 3, 64, 64)
+# cuda tensor over the SAME memory when torch is installed (DLPack).
+
+# Reusing the Pipeline object skips codegen entirely (the kernel handle is
+# cached on it); a NEW pipeline with different VALUES (sizes, mean/std,
+# rects keep the same types) still reuses the compiled .so from disk.
+pre.run()
+print("OK  second run: cached kernel, single ctypes call")
+
+# Single images work the same, and one-op eager calls live in fkl.F:
+roi = (fkl.pipe(fkl.Image(frames[0]))
+       .crop(40, 30, 240, 180)
+       .resize((NET_W, NET_H))
+       .normalize(MEAN, STD)
+       .to_planar("CHW")
+       .run())
+print(f"OK  single image crop->resize->normalize->CHW: {roi.shape}")
+
+gray = fkl.F.cvt_color(fkl.Image(frames[0]), "RGB2GRAY")   # eager one-op
+print(f"OK  fkl.F.cvt_color: {gray.shape} {gray.dtype.base}")

examples/README.md

@@ -34,6 +34,7 @@ disk cache in `~/.cache/fkl/` and start instantly.
 | `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` |
+| `12_torch_style_api.py` | PyTorch-style layer: `Image`/`ImageBatch` + `fkl.pipe(batch).resize().normalize().to_planar().run()` -> ONE fused kernel; eager `fkl.F.*` one-op calls |
 
 ## The one-page mental model
 
@@ -48,3 +49,6 @@ out  = pipe(x2)           # later calls: single ctypes call, ~70 us
   recompile.
 - Lists = Horizontal Fusion: `Crop([r1, r2, ...])` or `pipe([img1, img2])`.
 - `compose_divergent(plane_map, chain1, chain2)` = Divergent HF.
+- Torch-style spelling of the same thing:
+  `fkl.pipe(batch).resize((w, h)).normalize(m, s).to_planar("CHW").run()`
+  -> still ONE fused kernel (see `12_torch_style_api.py`).

fkl/__init__.py

@@ -27,6 +27,8 @@
 from .backend import CompilerBackend, set_backend, get_backend, clear_cache
 from .tensor import DeviceBuffer
 from .types import DType, dtype
+from .highlevel import Tensor, Image, ImageBatch, Pipeline, pipe
+from . import functional as F
 
 __all__ = [
     "TensorRead", "TensorWrite", "TensorSplit", "SplitWrite",
@@ -43,6 +45,7 @@
     "ScoreMod", "ALiBi", "SoftCap", "SlidingWindow", "BlockMask",
     "CompilerBackend", "set_backend", "get_backend", "clear_cache",
     "DeviceBuffer", "DType", "dtype",
+    "Tensor", "Image", "ImageBatch", "Pipeline", "pipe", "F",
 ]
 
 __version__ = "0.6.0"

fkl/functional.py

@@ -0,0 +1,69 @@
+"""fkl.F — eager, torch-style convenience functions.
+
+Each function is a ONE-op fused pipeline (fkl.pipe(x).<op>().run()): first
+call per signature JIT-compiles (disk-cached), then it costs a single kernel
+launch. For multi-step preprocessing prefer fkl.pipe(...) so the WHOLE chain
+fuses into one kernel instead of one kernel per call.
+
+    import fkl
+    y = fkl.F.resize(img, (64, 64))
+    g = fkl.F.cvt_color(img, "RGB2GRAY")
+    n = fkl.F.normalize(img, MEAN, STD)
+    p = fkl.F.to_planar(img)                 # HWC -> CHW planar
+"""
+from .highlevel import pipe
+
+
+def resize(x, size, interp="linear", aspect_ratio="ignore", background=None,
+           out=None, stream=None):
+    """Bilinear/nearest resize to size=(w, h) (int = square)."""
+    return pipe(x).resize(size, interp=interp, aspect_ratio=aspect_ratio,
+                          background=background).run(out=out, stream=stream)
+
+
+def crop(x, rect, out=None, stream=None):
+    """rect=(x, y, w, h), or a list of rects for multi-ROI (batch planes)."""
+    return pipe(x).crop(rect).run(out=out, stream=stream)
+
+
+def cvt_color(x, code, out=None, stream=None):
+    """OpenCV-style color conversion, e.g. 'RGB2BGR', 'BGR2GRAY'."""
+    return pipe(x).cvt_color(code).run(out=out, stream=stream)
+
+
+def normalize(x, mean, std, out=None, stream=None):
+    """(x - mean) / std per channel (auto-casts non-float chains to float32)."""
+    return pipe(x).normalize(mean, std).run(out=out, stream=stream)
+
+
+def cast(x, dtype, out=None, stream=None):
+    return pipe(x).cast(dtype).run(out=out, stream=stream)
+
+
+def saturate_cast(x, dtype, out=None, stream=None):
+    return pipe(x).saturate_cast(dtype).run(out=out, stream=stream)
+
+
+def to_planar(x, layout="CHW", out=None, stream=None):
+    """Packed HWC pixels -> planar 'CHW' ('NCHW' for batches)."""
+    return pipe(x).to_planar(layout).run(out=out, stream=stream)
+
+
+def mul(x, value, out=None, stream=None):
+    return pipe(x).mul(value).run(out=out, stream=stream)
+
+
+def add(x, value, out=None, stream=None):
+    return pipe(x).add(value).run(out=out, stream=stream)
+
+
+def sub(x, value, out=None, stream=None):
+    return pipe(x).sub(value).run(out=out, stream=stream)
+
+
+def div(x, value, out=None, stream=None):
+    return pipe(x).div(value).run(out=out, stream=stream)
+
+
+__all__ = ["resize", "crop", "cvt_color", "normalize", "cast",
+           "saturate_cast", "to_planar", "mul", "add", "sub", "div"]