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

[johnnynunez]

Motivation

compose() is op-chain-first: using it well means knowing the FKL op catalog (TensorRead, SaturateCast, TensorSplit, ...) and the read/write chain rules. DNN-preprocessing users coming from torch/torchvision think in tensors, images and batches. This adds that spelling as a thin layer over the existing machinery — every pipeline still lowers to ONE fused kernel through compose()'s JIT cache, and nothing about codegen, fusion or the C ABI changes.

batch = fkl.ImageBatch(frames)          # list of HWC images, or ONE (N,H,W,C) tensor
out = (fkl.pipe(batch)
       .resize((224, 224))              # BVF: fused into the read
       .normalize(MEAN, STD)            # Cast-if-needed + Sub + Div, in registers
       .to_planar("CHW")                # planar NCHW write (TensorSplit)
       .run())                          # ONE kernel for the whole batch
net(out.torch())                        # zero-copy (N, 3, 224, 224) cuda tensor

What's included

• fkl.Tensor — zero-copy wrapper over torch CUDA tensors, __cuda_array_interface__ objects (cupy/numba) and DeviceBuffer: .shape/.dtype/.device (scalar dtype; channels live in the shape), __cuda_array_interface__, __dlpack__/__dlpack_device__, .torch() (cached zero-copy view), metadata-only .reshape(), Tensor.empty(shape, dtype). A vector dtype spec appends the channel dim (empty((4, 6), 'uint8x3') ≡ empty((4, 6, 3), 'uint8')). Dtypes without an FKL kernel surface (float16/bfloat16/int64 torch tensors) raise a clear TypeError at construction instead of failing later.
• fkl.Image — HWC image semantics (width/height/channels, C 1..4) with eager shape validation; fkl.ImageBatch — list of same-shape images (validated for matching W/H/C/dtype/device) OR one batched (N, H, W[, C]) tensor split into zero-copy per-plane views; both run as horizontal fusion through the existing list-input path (BatchRead, N part of the kernel type).
• fkl.pipe(x) — fluent builder: crop, resize, border, cvt_color, normalize, mul/add/sub/div, cast, saturate_cast, apply(op) (escape hatch for any symbolic op), to_planar('CHW'), then .run(out=None, stream=None) -> fkl.Tensor. Each method appends an existing symbolic op; the same rules as compose() apply (types compile once, values ride in params[]; the kernel handle is cached on the pipeline). run() returns semantic shapes — (C,H,W) / (N,C,H,W) after to_planar, and 2-D sources keep a height-1 dim so (1, W) round-trips — and auto-allocated outputs follow the SOURCE's framework (torch in -> torch out, including batched-tensor plane views and reshaped Tensors). .source() shows the single generated TU for the exact variant .run() launches, including the effective ThreadFusion choice.
• Eager errors instead of nvcc errors or memory corruption: geometry-after-compute (BVF ordering), multi-ROI crop on a batch source, normalize channel mismatches, to_planar on single-channel chains, ops after to_planar, mixed-device batches — and out= is validated against the planned output (element count, dtype, device) BEFORE the launch, since a too-small or wrongly-typed out would otherwise mean an out-of-bounds device write or a silent reinterpretation.
• fkl.F — eager one-op fused pipelines (fkl.F.resize(img, (64, 64)), cvt_color, normalize, to_planar, ...).
• _make_dlpack_capsule gains an optional dims override (validated against the buffer's element count; default behaviour unchanged) so reshaped / planar-split outputs export the semantic shape over the same memory, and exports of one buffer are refcounted: the device memory is freed when the LAST consumer releases its capsule, so several Tensor views handed to torch cannot free memory out from under each other.

Out of scope, documented (module docstring, README, skill): autograd, broadcasting, >4-channel images (FKL vector pixels are 1..4 channels), multi-input graphs (single-input/single-output ABI; batches are horizontal fusion).

Tests / docs

• tests/test_highlevel.py: 54 dependency-free checks (DeviceBuffer-backed) — wrapper semantics (incl. vector-spec empty and shared-buffer DLPack ownership through the real capsule deleters), zero-copy plane views, pipelines vs pure-Python per-op references and vs equivalent compose() chains, batched-tensor == list-of-images equality, single-fused-kernel source check, ThreadFusion source/run, (1, W) round-trip, out=/stream/kernel-reuse, and the eager error paths (ordering, batch/channel mismatches, mixed devices, out= size/dtype validated before launch with the rejected buffer verified untouched). All pass on sm_121 / CUDA 13.0 (nvcc backend).
• tests/test_highlevel_torch.py: torch interop (Tensor over torch tensors, .torch() both ways, torch-backed Image/ImageBatch pipelines, out= into torch tensors incl. rejection cases, fp16/bf16 rejection, torch-allocated outputs for plane-view/reshaped sources); skips cleanly when torch is absent.
• Full pre-existing matrix re-run green (the one failing Saturate case fails identically on unmodified main with this toolkit — pre-existing, unrelated).
• New runnable examples/12_torch_style_api.py (batch of HWC uint8 frames -> resized, normalized, NCHW float in one fused kernel), README section, examples index row, usage-skill section.

🤖 Generated with Claude Code