Compose-time buffer binding: argument-free fused kernels + DeviceBuffer.from_ptr

README.md

@@ -56,6 +56,44 @@ out = pipe(x, stream=torch.cuda.current_stream())   # zero-copy, async
 
 No torch? Use the built-in dependency-free `DeviceBuffer` (CUDA driver via ctypes).
 
+## Compose-time buffer binding (argument-free kernels)
+
+The read/write ops optionally take a concrete buffer (cuda torch tensor, any
+`__cuda_array_interface__` object, or `DeviceBuffer`) at construction. A
+chain whose read is bound knows its dtype/shape immediately, so `compose()`
+compiles **eagerly** — the kernel is ready-to-run before the first call:
+
+```python
+x   = torch.empty(1080, 1920, 3, device="cuda", dtype=torch.uint8)
+out = torch.empty(3, 1080, 1920, device="cuda")
+
+k = fkl.compose(
+    fkl.TensorRead(x),          # input bound -> eager compile here
+    fkl.Cast("float32"),
+    fkl.Div((255.0,) * 3),
+    fkl.TensorSplit(out),       # output bound
+)
+for frame in camera:
+    x.copy_(frame)              # refill the bound buffer in place
+    k()                         # no arguments; k(stream=s) for async
+```
+
+Bound buffers are call **defaults**: `k(y)` / `k(y, out=z)` override them and
+are validated against the compiled signature (a different dtype/shape raises
+— compose an unbound chain for that). Binding only the read auto-allocates
+the output per call; binding only the write keeps the usual lazy `pipe(x)`
+path with the bound buffer as default `out=`. Unbound chains behave exactly
+as before.
+
+Raw pointers from C-style integrations bind the same way via a non-owning
+wrapper (the caller keeps ownership of the allocation — the wrapper never
+frees it):
+
+```python
+buf = fkl.DeviceBuffer.from_ptr(ptr, (1080, 1920, 3), "uint8")  # zero-copy
+k = fkl.compose(fkl.TensorRead(buf), ..., fkl.TensorWrite(dst))
+```
+
 ## Setup
 
 ### Plug and play (wheel with vendored headers)
@@ -105,11 +143,15 @@ backends (clang with auto-shims, nvcc):
 | test_circular_tensor             | 21    | PASS  | PASS |
 | test_thread_fusion               | 5     | PASS  | PASS |
 | test_e2e (timing/cache)          | 1     | PASS  | PASS |
+| test_bound_compose (binding)     | 41    | PASS  | PASS |
+
+(test_bound_compose verified on NVIDIA GB10, sm_121, CUDA 13.0, both
+backends.)
 
-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).
+Plus, separately: test_torch_integration (9 checks verified on real torch
+2.12+cu130: zero-copy both ways, external streams, CircularTensor->torch;
++5 bound-compose checks pending a torch-enabled run) and test_cpu_backend
+(5 checks, ParArch::CPU with numpy in/out, CPU==GPU cross-validated).
 
 ## CPU backend (no CUDA required)
 
@@ -153,8 +195,9 @@ true multi-GPU runs still need a 2+ GPU machine.
   [#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.
+163 checks per backend (+ 14 torch + 5 cpu separately). Steady-state hot launch ~70 µs/call; cache-hit
+compose ~0 ms (lazy; bound compose eagerly compiles — cache-hit ~1 ms);
+cold compile ~1-3 s once per chain signature.
 
 ## Temporal video: CircularTensor
 
@@ -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/` — runnable, mostly 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, a
+fused-vs-unfused benchmark (~5x at 1080p for a 6-op chain), and
+compose-time bound pipelines (argument-free `k()`).
 
 ## Current scope / next steps
 

examples/12_bound_pipeline.py

@@ -0,0 +1,76 @@
+"""Example 12 — Compose-time buffer binding (argument-free kernels).
+
+Bind the input/output buffers directly on the read/write ops:
+
+    k = fkl.compose(fkl.TensorRead(x), <ops...>, fkl.TensorWrite(out))
+    k()          # no arguments: reads x, writes out
+
+Key ideas demonstrated:
+  * Binding makes compose() EAGER: the kernel compiles inside compose()
+    (dtype/shape come from the bound buffers), so it is ready-to-run
+    before the first call — no first-call compile hiccup in the hot loop.
+  * Bound buffers are call DEFAULTS: k(y) / k(y, out=z) override them,
+    validated against the compiled signature (clear errors on mismatch).
+  * DeviceBuffer.from_ptr(ptr, shape, dtype) wraps a RAW device pointer
+    (non-owning) so C-style integrations bind pointers into the chain.
+"""
+import struct
+import fkl
+from _util import gpu_image_f32, to_floats
+
+W, H = 64, 32
+src = [float(i % 100) for i in range(W * H)]
+
+# ---- 1. bind input AND output: a fully wired, argument-free kernel --------
+x = gpu_image_f32(src, W, H)
+out = fkl.DeviceBuffer(W, H, "float32")
+
+k = fkl.compose(
+    fkl.TensorRead(x),        # <- input bound here
+    fkl.Mul(2.0),
+    fkl.Add(10.0),
+    fkl.TensorWrite(out),     # <- output bound here
+)
+# compose() already compiled the kernel (eager because the read is bound):
+assert len(k._variants) == 1, "expected eager compilation at compose() time"
+print("OK  compose() compiled eagerly: kernel ready before the first call")
+
+k()                           # the whole pipeline, no arguments
+got = to_floats(out, W * H)
+assert all(abs(g - (v * 2 + 10)) < 1e-5 for g, v in zip(got, src))
+print(f"OK  k() ran argument-free: out[:4] = {got[:4]}")
+
+# The binding is by POINTER: update the input in place and just call again.
+x.copy_from_host(struct.pack(f"{W * H}f", *[v + 1 for v in src]))
+k()
+got = to_floats(out, W * H)
+assert abs(got[0] - ((src[0] + 1) * 2 + 10)) < 1e-5
+print("OK  in-place input update + k(): same kernel, new data")
+
+# ---- 2. overrides still work (validated against the compiled signature) ---
+y = gpu_image_f32([v * 3 for v in src], W, H)
+k(y)                          # same dtype/shape: reuses the compiled kernel
+assert abs(to_floats(out, 1)[0] - (src[0] * 3 * 2 + 10)) < 1e-4
+print("OK  k(y) overrode the bound input, still wrote the bound output")
+
+try:
+    k(gpu_image_f32([0.0] * 16, 4, 4))   # wrong shape -> clear error
+    raise SystemExit("expected a ValueError!")
+except ValueError as e:
+    print(f"OK  mismatching override rejected: {str(e)[:56]}...")
+
+# ---- 3. raw pointers: DeviceBuffer.from_ptr (non-owning) -------------------
+# Wrap an EXTERNAL device pointer (here: another buffer's own pointer, the
+# way a C library would hand you cudaMalloc'd memory). The wrapper never
+# frees it — the caller keeps ownership and must keep it alive.
+dst = fkl.DeviceBuffer(W, H, "float32")
+dst_view = fkl.DeviceBuffer.from_ptr(dst.ptr, (H, W), "float32")
+
+k2 = fkl.compose(fkl.TensorRead(x), fkl.Div(2.0), fkl.TensorWrite(dst_view))
+k2()
+assert abs(to_floats(dst, 1)[0] - (src[0] + 1) / 2) < 1e-5
+del dst_view                  # non-owning: dst's memory is untouched
+assert abs(to_floats(dst, 1)[0] - (src[0] + 1) / 2) < 1e-5
+print("OK  from_ptr: raw device pointer bound as the kernel's destination")
+
+print("\nAll bound-pipeline demos passed.")

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_bound_pipeline.py` | compose-time buffer binding: `TensorRead(x)`/`TensorWrite(out)` -> eager compile, argument-free `k()`, call-time overrides, `DeviceBuffer.from_ptr` raw pointers |
 
 ## The one-page mental model
 
@@ -48,3 +49,5 @@ 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.
+- BINDING: `compose(fkl.TensorRead(x), ..., fkl.TensorWrite(out))` compiles
+  eagerly and runs as `k()` — buffers are call defaults, args override.

fkl/jit.py

@@ -15,6 +15,15 @@
 (chain, input dtype/layout) compiles once; every other call -- with any crop
 rect, any mul factor, any size values that keep the same types -- reuses the
 cached .so. Values travel through params[].
+
+COMPOSE-TIME BINDING: TensorRead(x) / TensorWrite(out) accept concrete
+buffers. A chain with a bound read compiles EAGERLY inside compose() (the
+dtype/shape are known then) and runs argument-free:
+
+    k = fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), fkl.TensorWrite(out))
+    k()             # ready-to-run: no arguments, writes into `out`
+    k(stream=s)     # same, async on an external stream
+    k(y)            # override the binding (same dtype/shape enforced)
 """
 from __future__ import annotations
 import ctypes
@@ -38,6 +47,14 @@ class FusedKernel:
     """A composed chain. Compiles lazily on first call (input dtype/layout
     is only known then) and caches per signature.
 
+    COMPOSE-TIME BINDING: if the chain's read/write ops carry bound buffers
+    (fkl.TensorRead(x) / fkl.TensorWrite(out)), the input dtype/shape are
+    known immediately, so the kernel compiles EAGERLY here and can be called
+    with no arguments — k() / k(stream=s). Bound buffers are call defaults:
+    k(y) / k(y, out=z) override them, validated against the compiled
+    signature. Binding a read without a write auto-allocates the output per
+    call, exactly like pipe(x) does today.
+
     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)."""
@@ -54,6 +71,28 @@ def __init__(self, ops: List[Op], target: str = "gpu",
         self.target = target
         self.thread_fusion = bool(thread_fusion) and target == "gpu"
         self._variants = {}  # signature -> (entry fn, params buffer, out_state)
+        # ---- compose-time buffer binding (optional) -----------------------
+        self._bound_in = getattr(ops[0], "_bound", None)
+        self._bound_out = getattr(ops[-1], "_bound", None)
+        self._bound_st = None  # (in DType, (w, h, planes)) once eagerly compiled
+        self._is_bound = self._bound_in is not None or self._bound_out is not None
+        if self._is_bound and target != "gpu":
+            raise ValueError("compose-time buffer binding requires "
+                             "target='gpu' (CPU chains take numpy arrays "
+                             "per call)")
+        if self._bound_in is not None:
+            if isinstance(self._bound_in, (list, tuple)):
+                raise TypeError("cannot bind a batch (list of images) at "
+                                "compose time; pass the list at call time")
+            vin: DeviceView = as_device_view(self._bound_in)
+            in_shape = (vin.width, vin.height, vin.planes)
+            # EAGER compile: dtype/shape are known now, so the kernel is
+            # ready-to-run before the first call (disk-cache hit if this
+            # signature was ever compiled before).
+            _, _, out_st, _ = self._get_variant(vin.dtype, in_shape)
+            self._bound_st = (vin.dtype, in_shape)
+            if self._bound_out is not None:
+                self._check_out(as_device_view(self._bound_out), out_st)
 
     # ---- compile path (cold, once per signature) --------------------------
     def _tf_effective(self, dt: DType, shape) -> bool:
@@ -93,13 +132,57 @@ def _get_variant(self, dt: DType, shape, n_inputs: int = 1):
         self._variants[sig] = variant
         return variant
 
+    # ---- bound-buffer validation (VALUE-level, like params[]) -------------
+    @staticmethod
+    def _check_out(vout: DeviceView, out_st):
+        """Validate an output buffer against the compiled output state.
+        Layouts may legitimately differ from out_st's (w, h, planes) view of
+        the same bytes (e.g. TensorSplit destinations are (C*planes, H, W)),
+        so check what the kernel actually assumes about the destination
+        pointer: base dtype + total byte size."""
+        need = (out_st.width * out_st.height * out_st.planes
+                * out_st.dtype.itemsize)
+        got = vout.width * vout.height * vout.planes * vout.dtype.itemsize
+        if vout.dtype.base != out_st.dtype.base or got != need:
+            raise ValueError(
+                f"output buffer mismatch: kernel writes {need} bytes of "
+                f"{out_st.dtype.base} (w={out_st.width}, h={out_st.height}, "
+                f"planes={out_st.planes}, channels={out_st.dtype.channels}); "
+                f"buffer has {got} bytes of {vout.dtype.base}")
+
+    def _check_in(self, vin: DeviceView):
+        """A bound kernel was compiled eagerly for ONE input signature; an
+        explicit-argument override must match it (a different dtype/shape is
+        a different kernel — compose an unbound chain for that)."""
+        b_dt, b_shape = self._bound_st
+        if vin.dtype != b_dt or (vin.width, vin.height, vin.planes) != b_shape:
+            raise ValueError(
+                f"input override mismatch: kernel was compiled at compose "
+                f"time for {b_dt} (w={b_shape[0]}, h={b_shape[1]}, "
+                f"planes={b_shape[2]}); got {vin.dtype} (w={vin.width}, "
+                f"h={vin.height}, planes={vin.planes})")
+
     # ---- HOT PATH ----------------------------------------------------------
-    def __call__(self, x, out=None, stream=None):
+    def __call__(self, x=None, out=None, stream=None):
+        if x is None:
+            if self._bound_in is None:
+                raise TypeError(
+                    "this kernel has no bound input buffer; call pipe(x) or "
+                    "bind one at compose time with fkl.TensorRead(x)")
+            x = self._bound_in
+        if out is None:
+            out = self._bound_out  # stays None when unbound: auto-alloc below
         if self.target == "cpu":
             return self._call_cpu(x, out)
         if isinstance(x, (list, tuple)):
+            if self._bound_st is not None:
+                raise ValueError(
+                    "kernel was compiled at compose time for a single bound "
+                    "input; a batch (list) call needs an unbound compose()")
             return self._call_batch(x, out, stream)
         vin: DeviceView = as_device_view(x)
+        if self._bound_st is not None:
+            self._check_in(vin)
         if vin.device != 0:
             # make the input's device current for the launch + output alloc
             from .tensor import DeviceBuffer
@@ -110,7 +193,11 @@ def __call__(self, x, out=None, stream=None):
 
         if out is None:
             out = self._alloc_out(vin, out_st)
-        vout: DeviceView = as_device_view(out)
+            vout: DeviceView = as_device_view(out)
+        else:
+            vout: DeviceView = as_device_view(out)
+            if self._is_bound:
+                self._check_out(vout, out_st)
 
         dims = _FklDims(vin.width, vin.height, vin.planes,
                         out_st.width, out_st.height, out_st.planes)
@@ -139,7 +226,11 @@ def _call_batch(self, xs, out=None, stream=None):
 
         if out is None:
             out = self._alloc_out(v0, out_st)
-        vout: DeviceView = as_device_view(out)
+            vout: DeviceView = as_device_view(out)
+        else:
+            vout: DeviceView = as_device_view(out)
+            if self._is_bound:
+                self._check_out(vout, out_st)
 
         # device-pointer array (host-side; FKL copies Ptr2D params by value
         # into kernel arguments at launch)
@@ -226,7 +317,15 @@ def source_for(self, dtype_spec, shape):
 
 def compose(*ops: Op, target: str = "gpu",
             thread_fusion: bool = False) -> FusedKernel:
-    """thread_fusion=True opts into FKL's ThreadFusion (TF::ENABLED):
+    """Build a FusedKernel from an op chain (TensorRead ... TensorWrite).
+
+    Unbound (default): lazy — the first call with a concrete input compiles
+    the kernel for that signature. Bound: pass concrete buffers to the read/
+    write ops (fkl.TensorRead(x), fkl.TensorWrite(out)) and the kernel
+    compiles EAGERLY here, ready to be called with no arguments (see
+    FusedKernel).
+
+    thread_fusion=True opts into FKL's ThreadFusion (TF::ENABLED):
     each thread processes multiple elements with vectorized accesses.
     DEFAULT IS FALSE, matching FKL's own TransformDPP<> default
     (TF::DISABLED): per upstream guidance it only pays off in a small set
@@ -264,6 +363,11 @@ def __init__(self, plane_map, chains):
         for ch in chains:
             if ch[0].role != READ or ch[-1].role != WRITE:
                 raise ValueError("every chain needs TensorRead() ... TensorWrite()")
+            if (getattr(ch[0], "_bound", None) is not None
+                    or getattr(ch[-1], "_bound", None) is not None):
+                raise ValueError(
+                    "compose_divergent does not support compose-time buffer "
+                    "binding; pass the batch (and out=) at call time")
         self.plane_map = list(plane_map)
         self.chains = [list(c) for c in chains]
         self._variants = {}