diff --git a/README.md b/README.md index 34005d8..dcc22bc 100644 --- a/README.md +++ b/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 diff --git a/examples/12_torch_style_api.py b/examples/12_torch_style_api.py new file mode 100644 index 0000000..c5633d0 --- /dev/null +++ b/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}") diff --git a/examples/README.md b/examples/README.md index e4da7a1..e0d0d6f 100644 --- a/examples/README.md +++ b/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`). diff --git a/fkl/__init__.py b/fkl/__init__.py index 13bb6c7..b707ea7 100644 --- a/fkl/__init__.py +++ b/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" diff --git a/fkl/functional.py b/fkl/functional.py new file mode 100644 index 0000000..92008c1 --- /dev/null +++ b/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).().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"] diff --git a/fkl/highlevel.py b/fkl/highlevel.py new file mode 100644 index 0000000..7eaad3c --- /dev/null +++ b/fkl/highlevel.py @@ -0,0 +1,681 @@ +"""PyTorch-flavoured high-level API: Tensor / Image / ImageBatch + pipe(). + +An API LAYER over the existing machinery, not a framework: + + fkl.Tensor thin zero-copy wrapper over torch CUDA tensors, + __cuda_array_interface__ objects and DeviceBuffer, with + .torch() / DLPack export and .reshape() (metadata only). + fkl.Image image semantics on a Tensor: width/height/channels, + 'HWC' input layout (planar 'CHW' OUTPUT comes from + .to_planar(), which lowers to the existing split ops). + fkl.ImageBatch batch of same-shape images: a list of Images OR a + batched (N, H, W[, C]) tensor (split into zero-copy + per-plane views; runs as horizontal fusion). + fkl.pipe(x) fluent builder. Every method appends an existing + symbolic op; .run() lowers the WHOLE chain to ONE fused + FKL kernel through compose()'s machinery (JIT-compiled + once per signature, disk-cached, then a single launch). + + out = (fkl.pipe(batch) + .resize((64, 64)) # BVF: fused into the read + .cvt_color("RGB2BGR") # VF: registers + .normalize(MEAN, STD) # Cast+Sub+Div composition + .to_planar("CHW") # planar NCHW write + .run()) # ONE kernel launch + out.torch() # zero-copy cuda tensor + +Explicitly OUT of scope (use torch itself for these): + - autograd (fkl kernels are preprocessing, not differentiable graph nodes) + - broadcasting (values broadcast per-channel only, like the raw ops) + - >4-channel images (FKL vector pixels are 1..4 channels) + - multi-input graphs (the ABI is single-input/single-output; batches are + horizontal fusion over one logical input) +""" +from __future__ import annotations +from typing import List, Optional, Tuple + +from .jit import FusedKernel +from .codegen import plan, generate_cu +from .operations import (Op, COMPUTE, TensorRead, TensorWrite, TensorSplit, + Mul, Add, Sub, Div, Cast, SaturateCast, Crop, Resize, + ColorConversion, Warping, BorderReader, Deinterlace) +from .tensor import as_device_view, DeviceBuffer, _make_dlpack_capsule +from .types import DType, dtype as _dt, from_shape, _BASES + + +def _prod(dims) -> int: + n = 1 + for d in dims: + n *= int(d) + return n + + +def _natural_shape(obj) -> Tuple[int, ...]: + s = getattr(obj, "shape", None) # torch / cupy / numba + if s is not None: + return tuple(int(d) for d in s) + cai = getattr(obj, "__cuda_array_interface__", None) + if cai is not None: + return tuple(int(d) for d in cai["shape"]) + raise TypeError(f"cannot infer shape of {type(obj).__name__}") + + +class Tensor: + """Zero-copy wrapper over any CUDA array (torch tensor, cupy/numba array, + DeviceBuffer, ... anything with __cuda_array_interface__). + + Never copies, never owns new memory (except Tensor.empty, which allocates + a DeviceBuffer). Exposes .shape/.dtype/.device, __cuda_array_interface__, + __dlpack__/__dlpack_device__ and .torch() (zero-copy view when torch is + installed). C-contiguous inputs only, like the rest of fkl.""" + __slots__ = ("_base", "_view", "_base_shape", "_shape", "_torch_view") + + def __init__(self, data): + if isinstance(data, Tensor): + self._base, self._view = data._base, data._view + self._base_shape, self._shape = data._base_shape, data._shape + self._torch_view = data._torch_view + return + self._view = as_device_view(data) # validates CUDA + C-contiguity + if self._view.dtype.base not in _BASES: + # e.g. float16/bfloat16/int64 torch tensors: as_device_view maps + # them, but there is no FKL kernel/interop surface for them yet. + raise TypeError( + f"fkl.Tensor does not support {self._view.dtype.base} yet; " + "cast to a supported dtype first (e.g. .to(torch.float32))") + self._base = data + self._base_shape = _natural_shape(data) + self._shape = self._base_shape + self._torch_view = None + + @classmethod + def empty(cls, shape, dtype="float32", device: int = 0) -> "Tensor": + """Allocate an uninitialized DeviceBuffer-backed Tensor. + + dtype is a scalar spec ('float32', 'uint8', ...; channels live in + shape, torch-style) or a vector spec like 'uint8x3', which APPENDS + the channel dim to shape: Tensor.empty((4, 6), 'uint8x3') has shape + (4, 6, 3) and scalar dtype uint8 — exactly like + Tensor.empty((4, 6, 3), 'uint8').""" + shape = tuple(int(s) for s in ((shape,) if isinstance(shape, int) + else shape)) + dt = _dt(dtype) + if dt.channels > 1: + if len(shape) == 1: + w, h, p = shape[0], 1, 1 + elif len(shape) == 2: + (h, w), p = shape, 1 + elif len(shape) == 3: + p, h, w = shape + else: + raise ValueError(f"unsupported shape {shape} for vector dtype") + ch = dt.channels + shape = shape + (ch,) # channels live in .shape (torch-style) + else: + d2, w, h, p = from_shape(shape, dt.base) + ch = d2.channels + buf = DeviceBuffer(w, h, dt.base, channels=ch, planes=p, device=device) + t = cls(buf) + return t if t._shape == shape else t.reshape(shape) + + # ---- torch-style metadata ------------------------------------------ + @property + def base(self): + """The wrapped object (torch tensor / DeviceBuffer / ...).""" + return self._base + + @property + def shape(self) -> Tuple[int, ...]: + return self._shape + + @property + def ndim(self) -> int: + return len(self._shape) + + @property + def dtype(self) -> DType: + """Scalar element type (channels live in .shape, torch-style).""" + return DType(self._view.dtype.base, 1) + + @property + def device(self) -> int: + return self._view.device + + @property + def ptr(self) -> int: + return self._view.ptr + + @property + def nbytes(self) -> int: + return _prod(self._shape) * _BASES[self._view.dtype.base][2] + + def reshape(self, *shape) -> "Tensor": + """Metadata-only reshape (zero-copy; memory is C-contiguous).""" + if len(shape) == 1 and isinstance(shape[0], (tuple, list)): + shape = tuple(shape[0]) + shape = tuple(int(s) for s in shape) + if _prod(shape) != _prod(self._shape): + raise ValueError(f"cannot reshape {self._shape} -> {shape} " + "(element counts differ)") + t = Tensor.__new__(Tensor) + t._base, t._view = self._base, self._view + t._base_shape, t._shape = self._base_shape, shape + t._torch_view = None + return t + + # ---- interop -------------------------------------------------------- + @property + def __cuda_array_interface__(self): + return {"shape": self._shape, "typestr": self._view.dtype.typestr, + "data": (self._view.ptr, False), "version": 3} + + def __dlpack_device__(self): + return (2, self._view.device) # (kDLCUDA, device_id) + + def __dlpack__(self, stream=None): + if isinstance(self._base, DeviceBuffer): + # our own allocation: export under THIS Tensor's (possibly + # reshaped) dims. Ownership moves to the consumers: the device + # memory is freed when the LAST exported capsule is released + # (several Tensor views may export the same buffer). + return _make_dlpack_capsule(self._base, dims=self._shape) + if self._shape == self._base_shape and hasattr(self._base, "__dlpack__"): + if stream is None: + return self._base.__dlpack__() + return self._base.__dlpack__(stream=stream) + raise BufferError( + "cannot export DLPack: the base object has no __dlpack__ or the " + "Tensor was reshaped; use .torch() or export .base directly") + + def torch(self): + """Zero-copy torch view of this Tensor (torch must be installed). + + DeviceBuffer-backed Tensors hand memory ownership to torch (DLPack + deleter): the device memory is freed when the last exported view is + released, so several Tensor views over one buffer are safe. The + view is cached so repeated calls return the same tensor.""" + if self._torch_view is not None: + return self._torch_view + try: + import torch as _torch + except ImportError as e: # pragma: no cover - torch-less machines + raise RuntimeError("fkl.Tensor.torch() requires torch") from e + if isinstance(self._base, _torch.Tensor): + t = self._base + if tuple(t.shape) != self._shape: + t = t.view(self._shape) + else: + try: + t = _torch.from_dlpack(self) + except (BufferError, TypeError, RuntimeError): + # CAI-only objects without DLPack: torch.as_tensor consumes + # __cuda_array_interface__ zero-copy on matching device. + t = _torch.as_tensor(self._base, + device=f"cuda:{self.device}") + if tuple(t.shape) != self._shape: + t = t.reshape(self._shape) + self._torch_view = t + return t + + # kernels take the raw base object (so auto-allocated outputs match the + # input's framework); a reshaped Tensor must pass itself, since only its + # own __cuda_array_interface__ carries the override shape. + @property + def _kernel_arg(self): + return self._base if self._shape == self._base_shape else self + + def __repr__(self): + return (f"fkl.Tensor(shape={self._shape}, dtype={self.dtype.base}, " + f"device=cuda:{self.device})") + + +class Image(Tensor): + """A single 'HWC' image: shape (H, W) or (H, W, C) with C in 2..4. + + Input layout is 'HWC' (packed vector pixels — what cameras, decoders and + OpenCV produce, and what FKL's ops consume). Planar 'CHW' is an OUTPUT + layout: produce it with pipe(...).to_planar('CHW') / fkl.F.to_planar, + which lower to the existing TensorSplit write.""" + __slots__ = () + + def __init__(self, data, layout: str = "HWC"): + if str(layout).upper() != "HWC": + raise ValueError( + f"unsupported input layout {layout!r}: Image inputs are " + "'HWC'; planar 'CHW' output comes from .to_planar()") + super().__init__(data) + s = self._shape + if len(s) == 3 and not (2 <= s[-1] <= 4): + raise ValueError( + f"Image got shape {s}: trailing dim {s[-1]} is not a channel " + "count (2..4). >4-channel images are out of scope; squeeze a " + "trailing 1; for a batch of 2-D images use fkl.ImageBatch") + if len(s) not in (2, 3): + raise ValueError( + f"Image needs (H, W) or (H, W, C<=4), got {s}; for a batch " + "use fkl.ImageBatch") + + @classmethod + def from_tensor(cls, t, layout: str = "HWC") -> "Image": + """From an fkl.Tensor or any CUDA array object (zero-copy).""" + return cls(t, layout=layout) + + @property + def layout(self) -> str: + return "HWC" + + @property + def height(self) -> int: + return self._shape[0] + + @property + def width(self) -> int: + return self._shape[1] + + @property + def channels(self) -> int: + return self._shape[2] if len(self._shape) == 3 else 1 + + def __repr__(self): + return (f"fkl.Image({self.width}x{self.height}x{self.channels} HWC, " + f"dtype={self.dtype.base}, device=cuda:{self.device})") + + +class _PlaneView: + """Zero-copy view of one plane of a batched (N, ...) CUDA array. Exposes + just enough for as_device_view (__cuda_array_interface__ + .device) and + keeps the owner alive.""" + __slots__ = ("_cai", "device", "_owner") + + def __init__(self, ptr, shape, typestr, device, owner): + self._cai = {"shape": tuple(int(d) for d in shape), "typestr": typestr, + "data": (int(ptr), False), "version": 3} + self.device = int(device) + self._owner = owner + + @property + def __cuda_array_interface__(self): + return self._cai + + +class ImageBatch: + """A batch of same-shape 'HWC' images. + + Build it from a list of Images/arrays, or from ONE batched tensor of + shape (N, H, W) or (N, H, W, C<=4) — the tensor is split into zero-copy + per-plane views. Pipelines over an ImageBatch run as HORIZONTAL FUSION: + one kernel, N thread-planes (batch reads), N part of the kernel type.""" + + def __init__(self, images, layout: str = "HWC"): + if str(layout).upper() != "HWC": + raise ValueError("ImageBatch input layout must be 'HWC'") + if isinstance(images, (list, tuple)): + if not images: + raise ValueError("ImageBatch needs at least one image") + imgs = [im if isinstance(im, Image) else Image(im) + for im in images] + key0 = (imgs[0].width, imgs[0].height, imgs[0].channels, + imgs[0].dtype.base, imgs[0].device) + for im in imgs[1:]: + if (im.width, im.height, im.channels, im.dtype.base, + im.device) != key0: + raise ValueError("all images in a batch must share " + "(W, H, C, dtype, device)") + self._images = imgs + self._tensor = None + return + t = images if isinstance(images, Tensor) else Tensor(images) + s = t.shape + if len(s) == 4 and not (2 <= s[-1] <= 4): + raise ValueError(f"batched tensor got shape {s}: trailing dim " + "must be a channel count (2..4)") + if len(s) not in (3, 4): + raise ValueError( + f"batched tensor must be (N, H, W) or (N, H, W, C<=4), got {s}") + stride = _prod(s[1:]) * _BASES[t.dtype.base][2] + ts = t.__cuda_array_interface__["typestr"] + self._images = [Image(_PlaneView(t.ptr + i * stride, s[1:], ts, + t.device, t)) + for i in range(s[0])] + self._tensor = t + + def __len__(self) -> int: + return len(self._images) + + def __getitem__(self, i) -> Image: + return self._images[i] + + def __iter__(self): + return iter(self._images) + + @property + def images(self) -> Tuple[Image, ...]: + return tuple(self._images) + + @property + def tensor(self) -> Optional[Tensor]: + """The batched source tensor, if this batch was built from one.""" + return self._tensor + + @property + def width(self) -> int: + return self._images[0].width + + @property + def height(self) -> int: + return self._images[0].height + + @property + def channels(self) -> int: + return self._images[0].channels + + @property + def dtype(self) -> DType: + return self._images[0].dtype + + @property + def device(self) -> int: + return self._images[0].device + + @property + def layout(self) -> str: + return "HWC" + + def __repr__(self): + return (f"fkl.ImageBatch({len(self)} x {self.width}x{self.height}" + f"x{self.channels} HWC, dtype={self.dtype.base})") + + +# ===================== fluent pipeline builder ============================= + +# ReadBack-family ops: fused backwards into the read (BVF) -> must be the +# LEADING stages of a chain, exactly like in compose(). +_GEOMETRY_OPS = (Crop, Resize, Warping, BorderReader, Deinterlace) + + +class Pipeline: + """Fluent, source-bound pipeline that lowers to ONE fused kernel. + + Every method appends an existing symbolic op (the same descriptors + compose() takes) and returns self. .run() materializes: the whole chain + becomes a single fused FKL kernel — geometry ops fuse backwards into the + read (BVF), compute ops fuse vertically in registers (VF), an ImageBatch + source adds horizontal fusion (HF) — JIT-compiled once per signature and + disk-cached, exactly like any compose() chain. Values (rects, sizes, + means) travel in params[]: an identical chain built with different + values reuses the cached .so. + + Ordering rule (mirrors compose()): crop/resize/warp/border are ReadBack + ops and must come before compute ops (cvt_color/normalize/...). + Reuse the Pipeline object for hot loops: the kernel handle is cached on + it after the first .run(). + + Note: value arithmetic (mul/add/sub/div) on integer VECTOR-pixel chains + (e.g. uchar3) does not compile in FKL (make_ promotes small-int vectors) + — .cast('float32') first; .normalize() and .resize() already move the + chain to float. + """ + + def __init__(self, src, thread_fusion: bool = False): + if isinstance(src, (list, tuple)): + src = ImageBatch(list(src)) + if isinstance(src, ImageBatch): + self._batch, self._single = src, None + im0 = src[0] + self._in_dt = DType(im0.dtype.base, im0.channels) + self._in_shape = (im0.width, im0.height, 1) + self._n_inputs = len(src) + self._keep_h = True # images are 2-D by definition + else: + t = src if isinstance(src, Tensor) else Tensor(src) + self._batch, self._single = None, t + dt, w, h, p = from_shape(t.shape, t.dtype.base) + self._in_dt = dt + self._in_shape = (w, h, p) + self._n_inputs = 1 + # 2-D-natured sources keep their height dim in the output even + # when it is 1 ((1, W) round-trips); 1-D vectors stay 1-D. Same + # channel folding rule as from_shape. + dims = t.shape + if len(dims) >= 2 and 2 <= dims[-1] <= 4: + dims = dims[:-1] + self._keep_h = len(dims) >= 2 + self._ops: List[Op] = [] + self._write: Optional[Op] = None + self._closed = False + self._geometry_open = True + self._thread_fusion = bool(thread_fusion) + self._kernel: Optional[FusedKernel] = None + + # ---- builder plumbing ------------------------------------------------- + def _append(self, op: Op) -> "Pipeline": + if self._closed: + raise RuntimeError( + ".to_planar() is the pipeline's write stage; no ops after it") + if isinstance(op, _GEOMETRY_OPS): + if not self._geometry_open: + raise ValueError( + f"{type(op).__name__} must come before compute ops: " + "geometry ops fuse backwards into the read (BVF), so " + "put crop/resize/warp/border first, like in compose()") + else: + self._geometry_open = False + self._ops.append(op) + self._kernel = None + return self + + def _chain_dtype(self) -> DType: + dt = self._in_dt + for op in self._ops: + dt = op.out_dtype(dt) + return dt + + @property + def ops(self) -> Tuple[Op, ...]: + """The symbolic compute ops appended so far (read/write excluded).""" + return tuple(self._ops) + + # ---- geometry (ReadBack / BVF) ----------------------------------------- + def crop(self, *args) -> "Pipeline": + """crop(x, y, w, h) / crop((x, y, w, h)) — single ROI. + crop([(x, y, w, h), ...]) — multi-ROI from ONE image (each rect + becomes a batch plane: horizontal fusion).""" + if len(args) == 4: + return self._append(Crop(*args)) + if len(args) == 1: + a = args[0] + if len(a) == 4 and all(isinstance(v, (int, float)) for v in a): + return self._append(Crop(*a)) + if self._n_inputs > 1: + raise ValueError( + "multi-ROI crop needs a single source image (the rect " + "list itself produces the batch planes)") + return self._append(Crop(list(a))) + raise ValueError( + "crop(x, y, w, h), crop((x, y, w, h)) or crop([(x, y, w, h), ...])") + + def resize(self, size, interp: str = "linear", + aspect_ratio: str = "ignore", background=None) -> "Pipeline": + """size = (w, h) target size, or an int for a square. Output becomes + float32 (interpolation), like fkl.Resize.""" + if isinstance(size, int): + w = h = size + else: + w, h = size + return self._append(Resize(int(w), int(h), interp=interp, + aspect_ratio=aspect_ratio, + background=background)) + + def border(self, mode: str = "replicate", value=0.0) -> "Pipeline": + """Out-of-bounds read policy (place before crop/resize).""" + return self._append(BorderReader(mode, value=value)) + + # ---- compute (VF) -------------------------------------------------------- + def cvt_color(self, code: str) -> "Pipeline": + """OpenCV-style color conversion code, e.g. 'RGB2BGR', 'BGR2GRAY'.""" + return self._append(ColorConversion(code)) + + def normalize(self, mean, std) -> "Pipeline": + """(x - mean) / std, per channel. Composed from the existing ops: + Cast('float32') (only when the chain is not float yet) + Sub + Div — + all fused in registers. mean/std are given in the chain's value + range at this point (e.g. 0..255 for uint8 sources).""" + dt = self._chain_dtype() + for name, v in (("mean", mean), ("std", std)): + if isinstance(v, (list, tuple)) and len(v) != dt.channels: + raise ValueError( + f"normalize {name} has {len(v)} values but the chain " + f"carries {dt.channels} channel(s) here") + if dt.base not in ("float32", "float64"): + self._append(Cast("float32")) + self._append(Sub(mean)) + return self._append(Div(std)) + + def mul(self, value) -> "Pipeline": + return self._append(Mul(value)) + + def add(self, value) -> "Pipeline": + return self._append(Add(value)) + + def sub(self, value) -> "Pipeline": + return self._append(Sub(value)) + + def div(self, value) -> "Pipeline": + return self._append(Div(value)) + + def cast(self, to) -> "Pipeline": + return self._append(Cast(to)) + + def saturate_cast(self, to) -> "Pipeline": + return self._append(SaturateCast(to)) + + def apply(self, op: Op) -> "Pipeline": + """Escape hatch: append any fkl symbolic op (e.g. fkl.VectorReorder, + fkl.MxVFloat3) to the fused chain.""" + if getattr(op, "role", None) != COMPUTE: + raise ValueError("apply() takes compute ops; the read/write " + "stages are managed by the pipeline") + return self._append(op) + + # ---- write stage ----------------------------------------------------- + def to_planar(self, layout: str = "CHW") -> "Pipeline": + """Terminal layout change: packed HWC pixels -> planar 'CHW' (or + 'NCHW' for batches) via the existing TensorSplit write. Must be the + last stage; single-channel chains are already planar.""" + if str(layout).upper() not in ("CHW", "NCHW"): + raise ValueError("to_planar supports 'CHW' / 'NCHW' output only") + if self._closed: + raise RuntimeError("to_planar() already applied") + dt = self._chain_dtype() + if dt.channels < 2: + raise ValueError("to_planar needs a multi-channel chain " + "(single-channel data is already planar)") + self._write = TensorSplit() + self._closed = True + self._kernel = None + return self + + # ---- materialization --------------------------------------------------- + def _full_ops(self) -> List[Op]: + write = self._write if self._write is not None else TensorWrite() + return [TensorRead(), *self._ops, write] + + def _ensure_kernel(self) -> FusedKernel: + if self._kernel is None: + self._kernel = FusedKernel(self._full_ops(), target="gpu", + thread_fusion=self._thread_fusion) + return self._kernel + + def _check_out(self, out_arg, out_st, sem: Tuple[int, ...]): + """Eager validation of a preallocated out= BEFORE the launch: a + too-small buffer would be an out-of-bounds device write, a wrong + dtype a silent reinterpretation.""" + vout = as_device_view(out_arg) + want = (out_st.width * out_st.height * out_st.planes + * out_st.dtype.channels) + got = vout.width * vout.height * vout.planes * vout.dtype.channels + if vout.dtype.base != out_st.dtype.base or got != want: + raise ValueError( + f"out= mismatch: the pipeline produces {sem} " + f"{out_st.dtype.base} ({want} elements), out= provides " + f"{_natural_shape(out_arg)} {vout.dtype.base} " + f"({got} elements)") + in_dev = (self._batch.device if self._batch is not None + else self._single.device) + if vout.device != in_dev: + raise ValueError(f"out= lives on cuda:{vout.device} but the " + f"pipeline input is on cuda:{in_dev}") + + def run(self, out=None, stream=None) -> Tensor: + """Execute the fused chain: ONE kernel launch. First call per + signature JIT-compiles (disk-cached forever); afterwards it is a + single ctypes call. stream: torch stream / cupy stream / raw handle + (async, caller syncs); None = internal stream, synced before return. + out: preallocated output (fkl.Tensor / torch tensor / DeviceBuffer), + validated against the planned output before the launch. + Returns an fkl.Tensor (use .torch() for a zero-copy torch view).""" + k = self._ensure_kernel() + _, out_st = plan(k.ops, self._in_dt, self._in_shape, self._n_inputs) + sem = _semantic_shape(out_st, split=isinstance(k.ops[-1], TensorSplit), + keep_h=self._keep_h) + src = self._batch.tensor if self._batch is not None else self._single + if out is not None: + out_arg = out._kernel_arg if isinstance(out, Tensor) else out + self._check_out(out_arg, out_st, sem) + else: + # allocate in the SOURCE's framework (torch in -> torch out), + # even when the kernel args are fkl-internal views (reshaped + # Tensors, batched-tensor plane views). A batch built from a + # LIST allocates from its first image inside the kernel call. + base = None if src is None else src.base + if isinstance(base, _PlaneView): + base = base._owner.base # plane of a batched source tensor + out_arg = (None if base is None + else k._alloc_out(as_device_view(base), out_st)) + if self._batch is not None: + args = [im._kernel_arg for im in self._batch.images] + raw = k(args, out=out_arg, stream=stream) + else: + raw = k(self._single._kernel_arg, out=out_arg, stream=stream) + return Tensor(raw).reshape(sem) + + def source(self) -> str: + """Debug: the generated C++ for the exact variant .run() would + launch (including the effective ThreadFusion choice for this + source's dtype/width) — one TU, one executeOperations call.""" + k = self._ensure_kernel() + return generate_cu(k.ops, self._in_dt, self._in_shape, self._n_inputs, + thread_fusion=k._tf_effective(self._in_dt, + self._in_shape)) + + def __repr__(self): + names = [type(op).__name__ for op in self._ops] + if self._write is not None: + names.append("to_planar") + src = (f"batch[{self._n_inputs}]" if self._batch is not None + else "image") + return f"fkl.Pipeline({src}: {' -> '.join(names) or ''})" + + +def _semantic_shape(out_st, split: bool, keep_h: bool = False) -> Tuple[int, ...]: + """Logical output shape for a chain output state: planar (C, H, W) / + (N, C, H, W) for split writes, (H, W[, C]) / (N, H, W[, C]) otherwise. + keep_h preserves a height-1 dim (2-D sources round-trip (1, W)); + without it height-1 collapses to (W,), matching 1-D vector sources.""" + ch, p = out_st.dtype.channels, out_st.planes + H, W = out_st.height, out_st.width + if split: + core = (ch, H, W) + else: + core = (H, W) if (H > 1 or keep_h) else (W,) + if ch > 1: + core = core + (ch,) + return (p,) + core if p > 1 else core + + +def pipe(src, thread_fusion: bool = False) -> Pipeline: + """Start a fluent pipeline bound to `src`: an Image / ImageBatch / + Tensor, a raw CUDA array, or a list of same-shape images (treated as an + ImageBatch -> horizontal fusion). The chain fuses into ONE kernel at + .run(). GPU only (for the CPU executor use fkl.compose(target='cpu')).""" + return Pipeline(src, thread_fusion=thread_fusion) diff --git a/fkl/tensor.py b/fkl/tensor.py index e028029..ee257ac 100644 --- a/fkl/tensor.py +++ b/fkl/tensor.py @@ -255,20 +255,38 @@ class _DLManagedTensor(ctypes.Structure): _dlpack_keepalive = {} # capsule id -> (managed, shape_arr, buffer) -def _make_dlpack_capsule(buf: "DeviceBuffer"): +def _make_dlpack_capsule(buf: "DeviceBuffer", dims=None): + """dims=None exports the buffer's natural shape. An explicit dims + reinterprets the SAME C-contiguous memory under another shape (used by + fkl.Tensor.reshape / pipeline outputs, e.g. planar (C, H, W) views over + a split write); it must cover exactly the buffer's base elements. + + A buffer may be exported several times (e.g. two reshaped Tensor views + each handed to torch): exports are refcounted, and the device memory is + freed when the LAST consumer releases its capsule.""" code_bits = _DL_CODES.get(buf.dtype.base) if code_bits is None: raise TypeError(f"DLPack: unsupported base dtype {buf.dtype.base}") code, bits = code_bits - dims = [] - if buf.planes > 1: - dims.append(buf.planes) - if buf.height > 1 or buf.planes > 1: - dims.append(buf.height) - dims.append(buf.width) - if buf.dtype.channels > 1: - dims.append(buf.dtype.channels) + if dims is None: + dims = [] + if buf.planes > 1: + dims.append(buf.planes) + if buf.height > 1 or buf.planes > 1: + dims.append(buf.height) + dims.append(buf.width) + if buf.dtype.channels > 1: + dims.append(buf.dtype.channels) + else: + dims = [int(d) for d in dims] + n = 1 + for d in dims: + n *= d + n_base = buf._nbytes // _BASES[buf.dtype.base][2] + if n != n_base: + raise ValueError( + f"DLPack dims {tuple(dims)} cover {n} elements; buffer has {n_base}") ndim = len(dims) shape_arr = (ctypes.c_int64 * ndim)(*dims) @@ -289,7 +307,10 @@ def _deleter(mt_ptr): if entry is not None: b = entry[2] try: - if getattr(b, "_dptr", None) and b._dptr.value: + # several capsules may share the buffer: last one frees + b._dlpack_exports -= 1 + if (b._dlpack_exports == 0 + and getattr(b, "_dptr", None) and b._dptr.value): b._cuda.cuMemFree_v2(b._dptr) b._dptr = ctypes.c_void_p() except Exception: @@ -298,6 +319,7 @@ def _deleter(mt_ptr): managed.deleter = _deleter # keep the struct + shape array + deleter + buffer alive until consumed _dlpack_keepalive[key] = (managed, shape_arr, buf, _deleter) + buf._dlpack_exports = getattr(buf, "_dlpack_exports", 0) + 1 buf._exported_dlpack = True # our __del__ must not double-free pycapsule_new = ctypes.pythonapi.PyCapsule_New diff --git a/skills/fkl-python-usage/SKILL.md b/skills/fkl-python-usage/SKILL.md index cc18d9b..003befc 100644 --- a/skills/fkl-python-usage/SKILL.md +++ b/skills/fkl-python-usage/SKILL.md @@ -52,6 +52,38 @@ out = pipe(x, stream=torch.cuda.current_stream()) # async on torch stream With an external stream the call does NOT sync (caller owns the stream). Without one, an internal static stream is used and synced before return. +## PyTorch-style layer (Tensor / Image / ImageBatch / pipe / F) + +The same machinery with torch-flavoured spelling (fkl/highlevel.py, +fkl/functional.py). Everything zero-copy; every pipeline lowers to ONE +fused kernel through compose()'s JIT cache: + +```python +t = fkl.Tensor(x) # torch cuda tensor / CAI obj / DeviceBuffer +t = fkl.Tensor.empty((H, W, 3), "uint8") +img = fkl.Image(x) # HWC, C in 1..4, planes == 1 +batch = fkl.ImageBatch([i0, i1, i2]) # or ONE (N, H, W[, C]) tensor + # (split into zero-copy plane views -> HF) +out = (fkl.pipe(batch) + .crop(x, y, w, h) # geometry FIRST (BVF into the read) + .resize((w2, h2)) + .normalize(MEAN, STD) # Cast('float32') if needed + Sub + Div + .to_planar("CHW") # terminal: TensorSplit planar write + .run(out=None, stream=None)) # -> fkl.Tensor; .torch() = zero-copy view +one = fkl.F.resize(img, (64, 64)) # eager one-op fused pipelines +``` + +Semantics to remember: +- run() returns fkl.Tensor with the SEMANTIC shape ((C,H,W) / (N,C,H,W) + after to_planar); .torch()/.reshape() are metadata-only over the same + C-contiguous memory. +- pipe methods append the SAME symbolic ops compose() takes (apply(op) for + anything else); types compile once, values ride in params[]. +- Geometry ops (crop/resize/border/warp) must precede compute ops — the + builder raises eagerly if you reorder. +- OUT of scope: autograd, broadcasting, >4-channel images, multi-input + graphs. Input layout is 'HWC' only; 'CHW' is an output layout. + ## Operation catalog Memory: TensorRead(), TensorWrite(), TensorSplit() [packed->planar], diff --git a/tests/test_highlevel.py b/tests/test_highlevel.py new file mode 100644 index 0000000..57ed270 --- /dev/null +++ b/tests/test_highlevel.py @@ -0,0 +1,428 @@ +"""High-level torch-style API (Tensor / Image / ImageBatch / pipe / F). + +Dependency-free (DeviceBuffer-backed). Pipelines are validated against +pure-Python per-op references, or against the equivalent hand-built +compose() chain (whose ops are themselves reference-tested in +test_operations.py). Torch interop lives in test_highlevel_torch.py. +""" +from harness import (dev_f32, dev_u8, f32, unf32, unu8, u8, check, + check_true, run) +import fkl + + +# ---- Tensor --------------------------------------------------------------- + +def t_tensor_wrap(): + W, H = 8, 4 + src = [float(y * W + x) for y in range(H) for x in range(W)] + buf = dev_f32(src, W, H) + t = fkl.Tensor(buf) + check_true("Tensor wrap: shape/dtype/device", + t.shape == (H, W) and t.dtype.base == "float32" + and t.device == 0 and t.ndim == 2) + check_true("Tensor wrap: zero-copy", + t.base is buf and t.ptr == buf.ptr + and t.__cuda_array_interface__["data"][0] == buf.ptr) + + +def t_tensor_empty(): + t = fkl.Tensor.empty((4, 6, 3), "uint8") + check_true("Tensor.empty: shape/dtype/nbytes", + t.shape == (4, 6, 3) and t.dtype.base == "uint8" + and t.nbytes == 72) + vals = [(i * 5) % 256 for i in range(72)] + t.base.copy_from_host(u8(vals)) + check("Tensor.empty: DeviceBuffer roundtrip", + unu8(t.base.copy_to_host(), 72), vals) + + +def t_tensor_empty_vector(): + t = fkl.Tensor.empty((4, 6), "uint8x3") # channel dim appended to shape + check_true("Tensor.empty vector spec: (4, 6) x uint8x3 -> (4, 6, 3)", + t.shape == (4, 6, 3) and t.dtype.base == "uint8" + and t.nbytes == 72) + t1 = fkl.Tensor.empty((5,), "float32x2") # 1-D: (W,) -> (W, C) + t3 = fkl.Tensor.empty((2, 3, 4), "uint8x4") # 3-D: (P, H, W) -> +C + check_true("Tensor.empty vector spec: 1-D and 3-D shapes", + t1.shape == (5, 2) and t3.shape == (2, 3, 4, 4)) + eq = fkl.Tensor.empty((4, 6, 3), "uint8") # scalar-spec equivalent + check_true("Tensor.empty: vector spec == scalar spec with channel dim", + eq.shape == t.shape and eq.dtype == t.dtype + and eq.nbytes == t.nbytes) + + +def t_tensor_reshape(): + t = fkl.Tensor.empty((2, 8), "float32") + r = t.reshape(4, 4) + check_true("reshape: zero-copy metadata", + r.shape == (4, 4) and r.ptr == t.ptr and r.base is t.base) + try: + t.reshape(5, 5) + ok = False + except ValueError: + ok = True + check_true("reshape: validates element count", ok) + + +def t_tensor_dlpack(): + src = [float(i) for i in range(16)] + t = fkl.Tensor(dev_f32(src, 8, 2)) + check_true("__dlpack_device__ = (kDLCUDA, 0)", + t.__dlpack_device__() == (2, 0)) + cap = t.reshape(16).__dlpack__() # reshaped export over DeviceBuffer + check_true("__dlpack__ capsule (reshaped DeviceBuffer base)", + cap is not None) + + +def t_dlpack_shared_buffer_ownership(): + """Two Tensor views exporting ONE DeviceBuffer: the memory must survive + until the LAST consumer releases its capsule (no use-after-free).""" + import ctypes + from fkl.tensor import _DLManagedTensor + t = fkl.Tensor.empty((4, 4), "float32") + buf = t.base + caps = [t.__dlpack__(), t.reshape(16).__dlpack__()] + get = ctypes.pythonapi.PyCapsule_GetPointer + get.restype = ctypes.c_void_p + get.argtypes = [ctypes.py_object, ctypes.c_char_p] + + def release(cap): # what a DLPack consumer does when it is dropped + mt = ctypes.cast(get(cap, b"dltensor"), + ctypes.POINTER(_DLManagedTensor)) + mt.contents.deleter(mt) + + release(caps[0]) + alive_after_first = bool(buf._dptr.value) + release(caps[1]) + check_true("DLPack: shared buffer freed by the LAST exported view", + alive_after_first and not buf._dptr.value) + + +# ---- Image / ImageBatch ---------------------------------------------------- + +def t_image_semantics(): + W, H = 6, 4 + img = fkl.Image(dev_u8([(i * 3) % 256 for i in range(W * H * 3)], W, H, ch=3)) + check_true("Image: w/h/c/layout", + img.width == W and img.height == H and img.channels == 3 + and img.layout == "HWC" and img.shape == (H, W, 3)) + img2 = fkl.Image.from_tensor(fkl.Tensor(dev_u8([0] * (W * H), W, H))) + check_true("Image.from_tensor (single channel)", + isinstance(img2, fkl.Image) and img2.channels == 1) + + +def t_image_rejects(): + bad = 0 + try: # 5 "channels" is neither a vector pixel nor a valid image + fkl.Image(fkl.DeviceBuffer(5, 4, "uint8", channels=5)) + except ValueError: + bad += 1 + try: # batch-shaped (N, H, W) input + fkl.Image(fkl.DeviceBuffer(16, 8, "float32", planes=3)) + except ValueError: + bad += 1 + try: # CHW input layout is not supported (output-only, via to_planar) + fkl.Image(fkl.DeviceBuffer(8, 8, "uint8", channels=3), layout="CHW") + except ValueError: + bad += 1 + check_true("Image rejects >4ch / batch-shaped / CHW-input", bad == 3) + + +def t_image_batch_list(): + W, H = 8, 6 + imgs = [dev_u8([(i + k) % 256 for i in range(W * H * 3)], W, H, ch=3) + for k in range(3)] + b = fkl.ImageBatch(imgs) + check_true("ImageBatch(list): len/w/h/c", + len(b) == 3 and b.width == W and b.height == H + and b.channels == 3 and isinstance(b[0], fkl.Image)) + + +def t_image_batch_tensor(): + W, H, N = 4, 3, 2 + base = fkl.DeviceBuffer(W, H, "uint8", channels=3, planes=N) + base.copy_from_host(u8([(i * 7) % 256 for i in range(N * H * W * 3)])) + b = fkl.ImageBatch(base) + per_plane = H * W * 3 + check_true("ImageBatch(tensor): zero-copy plane views", + len(b) == N and b[0].ptr == base.ptr + and b[1].ptr == base.ptr + per_plane + and b[0].shape == (H, W, 3) and b.tensor is not None) + + +def t_image_batch_device_mismatch(): + from fkl.highlevel import _PlaneView + W, H = 4, 3 + a = dev_u8([0] * (W * H * 3), W, H, ch=3) + b = dev_u8([0] * (W * H * 3), W, H, ch=3) + ghost = _PlaneView(b.ptr, (H, W, 3), "|u1", 1, b) # claims cuda:1 + try: + fkl.ImageBatch([a, ghost]) + ok = False + except ValueError: + ok = True + check_true("ImageBatch(list): mixed devices rejected eagerly", ok) + + +# ---- pipelines vs references ------------------------------------------------ + +def t_pipe_elementwise(): + W, H = 16, 8 + src = [float((y * W + x) % 97) for y in range(H) for x in range(W)] + out = fkl.pipe(fkl.Image(dev_f32(src, W, H))).mul(2.0).add(1.0).run() + check_true("pipe: returns fkl.Tensor with image shape", + isinstance(out, fkl.Tensor) and out.shape == (H, W)) + check("pipe: mul+add vs python ref", + unf32(out.base.copy_to_host(), W * H), [v * 2 + 1 for v in src]) + + +def t_pipe_crop(): + W, H = 33, 17 # non-pow2 to catch pitch errors + src = [(y * W + x) % 251 for y in range(H) for x in range(W)] + out = fkl.pipe(fkl.Image(dev_u8(src, W, H))).crop(5, 3, 8, 6).run() + exp = [float(((y + 3) * W + (x + 5)) % 251) + for y in range(6) for x in range(8)] + check_true("pipe: crop output shape", out.shape == (6, 8)) + check("pipe: crop vs position-encoded ref", + [float(v) for v in unu8(out.base.copy_to_host(), 48)], exp) + + +def t_pipe_resize(): + W, H = 16, 8 + src = [float((y * W + x) % 53) for y in range(H) for x in range(W)] + out = fkl.pipe(fkl.Image(dev_f32(src, W, H))).resize((7, 5)).run() + ref = fkl.compose(fkl.TensorRead(), fkl.Resize(7, 5), + fkl.TensorWrite())(dev_f32(src, W, H)) + check_true("pipe: resize output shape+dtype", + out.shape == (5, 7) and out.dtype.base == "float32") + check("pipe: resize == compose(Resize) chain", + unf32(out.base.copy_to_host(), 35), + unf32(ref.copy_to_host(), 35), tol=1e-5) + + +def t_pipe_cvt_color(): + W, H = 8, 2 + src = [(i * 17) % 256 for i in range(W * H * 3)] + out = fkl.pipe(fkl.Image(dev_u8(src, W, H, ch=3))).cvt_color("BGR2GRAY").run() + exp = [] + for p in range(W * H): + b, g, r = src[p * 3:(p + 1) * 3] + exp.append(float(int(0.299 * r + 0.587 * g + 0.114 * b + 0.5))) + check("pipe: cvt_color BGR2GRAY vs BT.601 ref", + [float(v) for v in unu8(out.base.copy_to_host(), W * H)], exp, + tol=1.0) + + +def t_pipe_normalize_autocast(): + W, H = 6, 4 + MEAN, STD = (110.0, 120.0, 130.0), (50.0, 55.0, 60.0) + src = [(i * 13) % 256 for i in range(W * H * 3)] + out = fkl.pipe(fkl.Image(dev_u8(src, W, H, ch=3))).normalize(MEAN, STD).run() + exp = [(float(v) - MEAN[i % 3]) / STD[i % 3] for i, v in enumerate(src)] + check_true("pipe: normalize auto-casts uint8 -> float32", + out.dtype.base == "float32" and out.shape == (H, W, 3)) + check("pipe: normalize vs (x-mean)/std ref", + unf32(out.base.copy_to_host(), W * H * 3), exp) + + +def t_pipe_to_planar(): + W, H = 4, 3 + src = [(i * 7) % 256 for i in range(W * H * 3)] + p = fkl.pipe(fkl.Image(dev_u8(src, W, H, ch=3))).to_planar("CHW") + out = p.run() + exp = [float(src[px * 3 + c]) for c in range(3) for px in range(W * H)] + check_true("pipe: to_planar semantic shape (C, H, W)", + out.shape == (3, H, W)) + check("pipe: to_planar HWC->CHW vs ref", + [float(v) for v in unu8(out.base.copy_to_host(), W * H * 3)], exp) + try: + p.mul(2.0) + ok = False + except RuntimeError: + ok = True + check_true("pipe: no ops after to_planar", ok) + + +def t_pipe_dnn_single_kernel(): + """Flagship chain: crop -> resize -> normalize -> CHW, ONE fused kernel.""" + W, H = 32, 24 + MEAN = (123.675, 116.28, 103.53) + STD = (58.395, 57.12, 57.375) + src = [(i * 3 + (i % 3) * 7) % 256 for i in range(W * H * 3)] + p = (fkl.pipe(fkl.Image(dev_u8(src, W, H, ch=3))) + .crop(4, 2, 24, 20) + .resize((8, 8)) + .normalize(MEAN, STD) + .to_planar("CHW")) + out = p.run() + ref = fkl.compose(fkl.TensorRead(), fkl.Crop(4, 2, 24, 20), + fkl.Resize(8, 8), fkl.Sub(MEAN), fkl.Div(STD), + fkl.TensorSplit())(dev_u8(src, W, H, ch=3)) + check_true("pipe: DNN chain output (3, 8, 8)", out.shape == (3, 8, 8)) + check("pipe: DNN chain == compose reference chain", + unf32(out.base.copy_to_host(), 3 * 64), + unf32(ref.copy_to_host(), 3 * 64)) + # ONE fused kernel: a single TU with a single entry; every stage appears + # inside the one fused launch (executeOperations shows up twice only + # because the template has an async and a sync stream branch). + src = p.source() + check_true("pipe: lowers to ONE fused kernel", + src.count("void fkl_entry") == 1 + and src.count("executeOperations") == 2 + and all(tok in src for tok in + ("Crop", "Resize", "Sub", "Div", "TensorSplit"))) + + +def t_pipe_batch_list(): + """HF: batch pipeline plane b == single-image pipeline on image b.""" + W, H = 12, 10 + MEAN, STD = (100.0, 110.0, 120.0), (50.0, 51.0, 52.0) + srcs = [[(i * (k + 3)) % 256 for i in range(W * H * 3)] for k in range(3)] + imgs = [dev_u8(s, W, H, ch=3) for s in srcs] + out = (fkl.pipe(fkl.ImageBatch(imgs)) + .resize((6, 5)).normalize(MEAN, STD).to_planar("CHW").run()) + check_true("pipe(batch): semantic shape (N, C, H, W)", + out.shape == (3, 3, 5, 6)) + got = unf32(out.base.copy_to_host(), 3 * 3 * 5 * 6) + per = 3 * 5 * 6 + for b in range(3): + single = (fkl.pipe(fkl.Image(imgs[b])) + .resize((6, 5)).normalize(MEAN, STD).to_planar("CHW").run()) + check(f"pipe(batch): plane {b} == single-image pipeline", + got[b * per:(b + 1) * per], + unf32(single.base.copy_to_host(), per)) + + +def t_pipe_batch_tensor_equals_list(): + W, H, N = 8, 6, 3 + flat = [(i * 11) % 256 for i in range(N * H * W * 3)] + per = H * W * 3 + batched = fkl.DeviceBuffer(W, H, "uint8", channels=3, planes=N) + batched.copy_from_host(u8(flat)) + imgs = [dev_u8(flat[k * per:(k + 1) * per], W, H, ch=3) for k in range(N)] + ops = lambda p: p.resize((4, 3)).mul(2.0).to_planar("CHW").run() + a = ops(fkl.pipe(fkl.ImageBatch(batched))) + b = ops(fkl.pipe(fkl.ImageBatch(imgs))) + n = N * 3 * 3 * 4 + check_true("pipe: batched tensor shape == list shape", a.shape == b.shape) + check("pipe: ImageBatch(tensor) == ImageBatch(list)", + unf32(a.base.copy_to_host(), n), unf32(b.base.copy_to_host(), n)) + + +def t_functional_eager(): + W, H = 10, 6 + src = [(i * 19) % 256 for i in range(W * H * 3)] + img = fkl.Image(dev_u8(src, W, H, ch=3)) + r = fkl.F.resize(img, (5, 3)) + ref = fkl.compose(fkl.TensorRead(), fkl.Resize(5, 3), + fkl.TensorWrite())(dev_u8(src, W, H, ch=3)) + check("F.resize == compose(Resize)", + unf32(r.base.copy_to_host(), 45), unf32(ref.copy_to_host(), 45), + tol=1e-5) + m = fkl.F.mul(fkl.Image(dev_f32([float(i) for i in range(24)], 6, 4)), 3.0) + check("F.mul vs python ref", unf32(m.base.copy_to_host(), 24), + [i * 3.0 for i in range(24)]) + pl = fkl.F.to_planar(img) + exp = [float(src[px * 3 + c]) for c in range(3) for px in range(W * H)] + check("F.to_planar vs ref", + [float(v) for v in unu8(pl.base.copy_to_host(), W * H * 3)], exp) + + +def t_pipe_out_and_reuse(): + W, H = 8, 4 + src = [float(i % 31) for i in range(W * H)] + p = fkl.pipe(fkl.Image(dev_f32(src, W, H))).mul(3.0) + p.run() + pre = fkl.Tensor.empty((H, W), "float32") + out = p.run(out=pre) + check_true("pipe: out= writes into the provided Tensor", + out.ptr == pre.ptr) + check("pipe: out= values", unf32(pre.base.copy_to_host(), W * H), + [v * 3 for v in src]) + check_true("pipe: kernel cached across runs (one variant)", + p._kernel is not None and len(p._kernel._variants) == 1) + + +def t_pipe_out_validation(): + W, H = 8, 4 + src = [float(i % 29) for i in range(W * H)] + p = fkl.pipe(fkl.Image(dev_f32(src, W, H))).mul(2.0) + small = fkl.Tensor.empty((2, 2), "float32") + small.base.copy_from_host(f32([-7.0] * 4)) + bad = 0 + try: # too small: would be an out-of-bounds device write + p.run(out=small) + except ValueError: + bad += 1 + try: # same element count, wrong dtype: silent reinterpretation + p.run(out=fkl.Tensor.empty((H, W), "int32")) + except ValueError: + bad += 1 + check_true("pipe: out= size/dtype validated BEFORE the launch", bad == 2) + check("pipe: rejected out= buffer left untouched", + unf32(small.base.copy_to_host(), 4), [-7.0] * 4) + + +def t_pipe_source_thread_fusion(): + src = [float(i % 19) for i in range(8 * 4)] + p = fkl.pipe(fkl.Image(dev_f32(src, 8, 4)), thread_fusion=True).mul(2.0) + check_true("pipe: source() shows the TF variant that will run", + "TF::ENABLED" in p.source()) + out = p.run() + check("pipe: thread_fusion run vs python ref", + unf32(out.base.copy_to_host(), 32), [v * 2 for v in src]) + q = fkl.pipe(fkl.Image(dev_f32([0.0] * 24, 6, 4)), + thread_fusion=True).mul(2.0) + check_true("pipe: source() shows scalar fallback (row not 16B-aligned)", + "TF::ENABLED" not in q.source()) + + +def t_pipe_height1_roundtrip(): + W = 8 + src = [float(i) for i in range(W)] + img = fkl.Image(fkl.Tensor(dev_f32(src, W)).reshape(1, W)) + out = fkl.pipe(img).mul(2.0).run() + check_true("pipe: (1, W) image keeps its rank", out.shape == (1, W)) + check("pipe: (1, W) values", unf32(out.base.copy_to_host(), W), + [v * 2 for v in src]) + vec = fkl.pipe(fkl.Tensor(dev_f32(src, W))).mul(2.0).run() + check_true("pipe: 1-D source stays 1-D", vec.shape == (W,)) + + +def t_pipe_errors(): + W, H = 8, 8 + img = fkl.Image(dev_u8([0] * (W * H * 3), W, H, ch=3)) + bad = 0 + try: # geometry after compute: BVF ordering violated + fkl.pipe(img).mul(2.0).resize((4, 4)) + except ValueError: + bad += 1 + try: # multi-ROI crop over a batch source + imgs = [dev_u8([0] * (W * H * 3), W, H, ch=3) for _ in range(2)] + fkl.pipe(fkl.ImageBatch(imgs)).crop([(0, 0, 4, 4), (1, 1, 4, 4)]) + except ValueError: + bad += 1 + try: # normalize channel mismatch + fkl.pipe(img).normalize((1.0, 2.0), (1.0, 2.0)) + except ValueError: + bad += 1 + try: # to_planar on single-channel chain + fkl.pipe(fkl.Image(dev_u8([0] * (W * H), W, H))).to_planar() + except ValueError: + bad += 1 + check_true("pipe: ordering/batch/channel/planar errors are eager", + bad == 4) + + +if __name__ == "__main__": + run([t_tensor_wrap, t_tensor_empty, t_tensor_empty_vector, + t_tensor_reshape, t_tensor_dlpack, t_dlpack_shared_buffer_ownership, + t_image_semantics, t_image_rejects, t_image_batch_list, + t_image_batch_tensor, t_image_batch_device_mismatch, + t_pipe_elementwise, t_pipe_crop, + t_pipe_resize, t_pipe_cvt_color, t_pipe_normalize_autocast, + t_pipe_to_planar, t_pipe_dnn_single_kernel, t_pipe_batch_list, + t_pipe_batch_tensor_equals_list, t_functional_eager, + t_pipe_out_and_reuse, t_pipe_out_validation, + t_pipe_source_thread_fusion, t_pipe_height1_roundtrip, + t_pipe_errors], "highlevel-api") diff --git a/tests/test_highlevel_torch.py b/tests/test_highlevel_torch.py new file mode 100644 index 0000000..ac0314b --- /dev/null +++ b/tests/test_highlevel_torch.py @@ -0,0 +1,128 @@ +"""Torch interop for the high-level API (requires torch+cuda; skips cleanly +otherwise): fkl.Tensor over torch tensors, .torch() zero-copy views both +ways, torch-backed Image/ImageBatch pipelines, out= into torch tensors. + +Run with a torch-enabled python: + .venv-torch/bin/python tests/test_highlevel_torch.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 + + +def t_tensor_wraps_torch(): + x = torch.arange(32, device="cuda", dtype=torch.float32).reshape(4, 8) + t = fkl.Tensor(x) + check_true("Tensor(torch): shape/dtype/device/ptr", + t.shape == (4, 8) and t.dtype.base == "float32" + and t.device == 0 and t.ptr == x.data_ptr()) + check_true("Tensor(torch).torch() is the SAME tensor", t.torch() is x) + v = t.reshape(8, 4).torch() + check_true("reshaped .torch() is a zero-copy view", + tuple(v.shape) == (8, 4) and v.data_ptr() == x.data_ptr()) + + +def t_empty_torch_view(): + import struct + t = fkl.Tensor.empty((4, 4), "float32") + v = t.torch() # DLPack: ownership moves to torch + v.fill_(3.5) + check_true("Tensor.empty().torch(): cached zero-copy view", + t.torch() is v and tuple(v.shape) == (4, 4) and v.is_cuda) + check("Tensor.empty(): DeviceBuffer sees torch writes (zero-copy)", + list(struct.unpack("16f", t.base.copy_to_host())), [3.5] * 16) + + +def t_torch_image_pipeline(): + H, W = 24, 32 + MEAN, STD = (123.675, 116.28, 103.53), (58.395, 57.12, 57.375) + x = (torch.arange(H * W * 3, device="cuda", dtype=torch.float32) + .mul(7).remainder(256).to(torch.uint8).reshape(H, W, 3).contiguous()) + out = (fkl.pipe(fkl.Image(x)).normalize(MEAN, STD).to_planar("CHW") + .run()) + v = out.torch() + ref = ((x.float() - torch.tensor(MEAN, device="cuda")) + / torch.tensor(STD, device="cuda")).permute(2, 0, 1).contiguous() + check_true("torch pipeline: (C, H, W) float out", + tuple(v.shape) == (3, H, W) and v.dtype == torch.float32) + check_true("torch pipeline: values match torch reference", + torch.allclose(v, ref, atol=1e-3)) + + +def t_torch_batch_pipeline(): + N, H, W = 3, 10, 12 + x = (torch.arange(N * H * W * 3, device="cuda", dtype=torch.float32) + .mul(13).remainder(256).to(torch.uint8) + .reshape(N, H, W, 3).contiguous()) + b = fkl.ImageBatch(x) + check_true("ImageBatch(torch NHWC): plane views", + len(b) == N and b[0].ptr == x.data_ptr()) + out = fkl.pipe(b).resize((6, 5)).to_planar("CHW").run() + v = out.torch() + check_true("torch batch pipeline: (N, C, H, W) out", + tuple(v.shape) == (N, 3, 5, 6)) + single = fkl.pipe(fkl.Image(x[0])).resize((6, 5)).to_planar("CHW").run() + check_true("torch batch plane 0 == single-image pipeline", + torch.allclose(v[0], single.torch(), atol=1e-4)) + + +def t_torch_out_prealloc(): + x = torch.full((8, 8), 2.0, device="cuda") + pre = torch.empty(8, 8, device="cuda") + out = fkl.pipe(x).add(3.0).run(out=pre) + check_true("pipe: out= torch tensor filled in place", + torch.allclose(pre, torch.full_like(pre, 5.0)) + and out.ptr == pre.data_ptr()) + bad = 0 + try: # too small: must be rejected BEFORE the launch + fkl.pipe(x).add(3.0).run(out=torch.empty(2, 2, device="cuda")) + except ValueError: + bad += 1 + try: # matching element count, wrong dtype + fkl.pipe(x).add(3.0).run( + out=torch.empty(8, 8, device="cuda", dtype=torch.int32)) + except ValueError: + bad += 1 + check_true("pipe: out= torch tensor validated (size + dtype)", bad == 2) + + +def t_torch_half_rejected(): + for dt in (torch.float16, torch.bfloat16): + try: + fkl.Tensor(torch.zeros(4, 4, device="cuda", dtype=dt)) + ok = False + except TypeError: + ok = True + check_true(f"Tensor(torch {dt}): clear TypeError (no FKL surface)", + ok) + + +def t_torch_out_allocation_follows_source(): + # batched torch tensor: the kernel args are fkl-internal plane views, + # but the auto-allocated output must still be torch (torch in -> + # torch out); same for a reshaped Tensor over a torch base. + x = (torch.arange(2 * 4 * 6 * 3, device="cuda", dtype=torch.float32) + .remainder(256).to(torch.uint8).reshape(2, 4, 6, 3).contiguous()) + out = fkl.pipe(fkl.ImageBatch(x)).cast("float32").to_planar("CHW").run() + check_true("ImageBatch(torch tensor): auto-allocated out is torch", + type(out.base).__module__.split(".")[0] == "torch" + and tuple(out.shape) == (2, 3, 4, 6)) + y = torch.arange(16, device="cuda", dtype=torch.float32) + r = fkl.pipe(fkl.Tensor(y).reshape(4, 4)).mul(2.0).run() + check_true("reshaped torch source: auto-allocated out is torch", + type(r.base).__module__.split(".")[0] == "torch" + and torch.allclose(r.torch(), y.reshape(4, 4) * 2)) + + +if __name__ == "__main__": + run([t_tensor_wraps_torch, t_empty_torch_view, t_torch_image_pipeline, + t_torch_batch_pipeline, t_torch_out_prealloc, t_torch_half_rejected, + t_torch_out_allocation_follows_source], "highlevel-torch")