From f5b2ced1c4b080facc15557773a357106f7c9e50 Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Sat, 18 Jul 2026 18:40:17 +0200 Subject: [PATCH] compose-time binding: bind IO buffers on read/write ops + from_ptr MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Read/write descriptors optionally take a concrete buffer at construction (TensorRead(x), TensorWrite(out); also TensorSplit/SplitWrite/TensorTSplit dest and TensorPack source). Binding is VALUE-level, like params[]: it never touches token()/codegen — only which pointer the launch uses. A chain with a bound read knows its dtype/shape at compose() time, so FusedKernel compiles EAGERLY there and runs argument-free: k(), k(stream=s). Bound buffers are call defaults — k(y) / k(y, out=z) override them, validated against the compiled signature (ValueError on dtype/shape mismatch; a bound kernel is one signature by design). Read-only binding auto-allocates the output per call; write-only binding keeps the lazy pipe(x) path with the bound buffer as default out=. compose_divergent rejects bindings explicitly (its batch arrives at call time). Unbound chains are untouched (generate_cu output verified byte-identical vs upstream/main). DeviceBuffer.from_ptr(ptr, shape, dtype, stream=None, device=0) wraps an external raw device pointer NON-OWNING (C-style integrations): __del__ and the DLPack deleter never free it — the caller keeps ownership/lifetime. The optional producer stream is advertised via __cuda_array_interface__ v3. Guards on a raw-pointer API: NULL/None ptr raises ValueError; passing a vector dtype AND the trailing channel dim (shape (H, W, 3) + uint8x3, which would silently size the wrapper 3x the allocation) raises TypeError; a producer stream handle of 0 is advertised as 1 (CAI v3 spells the legacy default stream as 1 — the raw 0 is disallowed by the spec). tests/test_bound_compose.py: 41 checks (argless call, eager compile + cache-hit second compose, overrides, from_ptr round-trip incl. raw-pointer output, from_ptr guards, error cases) — 41/41 on GB10 (sm_121, CUDA 13.0) on BOTH backends (clang, nvcc). Full core suite re-run on both backends: all green except the pre-existing Saturate failures (also fail on pristine upstream/main on this machine, unrelated). torch additions in test_torch_integration.py skip cleanly without torch. New example: examples/12_bound_pipeline.py. Co-Authored-By: Claude Fable 5 --- README.md | 66 +++++-- examples/12_bound_pipeline.py | 76 ++++++++ examples/README.md | 3 + fkl/jit.py | 112 ++++++++++- fkl/operations.py | 49 ++++- fkl/tensor.py | 91 ++++++++- skills/fkl-python-testing/SKILL.md | 1 + skills/fkl-python-usage/SKILL.md | 28 +++ tests/test_bound_compose.py | 300 +++++++++++++++++++++++++++++ tests/test_torch_integration.py | 33 +++- 10 files changed, 733 insertions(+), 26 deletions(-) create mode 100644 examples/12_bound_pipeline.py create mode 100644 tests/test_bound_compose.py diff --git a/README.md b/README.md index 34005d8..3c5b480 100644 --- a/README.md +++ b/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 diff --git a/examples/12_bound_pipeline.py b/examples/12_bound_pipeline.py new file mode 100644 index 0000000..34dd1c1 --- /dev/null +++ b/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), , 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.") diff --git a/examples/README.md b/examples/README.md index e4da7a1..52cf13e 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_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. diff --git a/fkl/jit.py b/fkl/jit.py index 38779a3..558e072 100644 --- a/fkl/jit.py +++ b/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 = {} diff --git a/fkl/operations.py b/fkl/operations.py index f3f4b72..5abf051 100644 --- a/fkl/operations.py +++ b/fkl/operations.py @@ -374,8 +374,13 @@ class _ReadOp(Op): if/elif ladder in codegen) means a Read IOp is self-contained — exactly how the C++ side carries the read in its IOp type. emit_read is the read counterpart to cpp() for COMPUTE ops. + + `_bound` optionally holds a concrete device buffer bound at construction + (compose-time binding). Like params[], it is a VALUE-level concern: it + never appears in token()/codegen, only in which pointer the launch uses. """ role = READ + _bound = None # optional compose-time-bound input buffer def emit_read(self, state, mem, n_inputs): """-> (in_decl, read_expr). in_decl declares the host buffer object @@ -385,8 +390,9 @@ def emit_read(self, state, mem, n_inputs): class _WriteOp(Op): """A chain-tail memory Write IOp. emit_write() is the write counterpart - to _ReadOp.emit_read().""" + to _ReadOp.emit_read(). `_bound` mirrors _ReadOp._bound for outputs.""" role = WRITE + _bound = None # optional compose-time-bound output buffer def emit_write(self, state, mem, pbase): """-> (out_decl, write_expr). out_decl declares the host output @@ -397,9 +403,18 @@ def emit_write(self, state, mem, pbase): class TensorRead(_ReadOp): """Read input. 2D Ptr2D -> PerThreadRead<_2D>; 3D Tensor -> TensorRead; batch of B same-size 2D images -> std::array -> PerThreadRead - (BatchRead under the hood, the HF read).""" + (BatchRead under the hood, the HF read). + + TensorRead(x) BINDS a concrete buffer (cuda torch tensor, any + __cuda_array_interface__ object, or DeviceBuffer — incl. from_ptr) at + compose() time: the kernel takes dtype/shape from it, compiles eagerly, + and can then be called with no arguments. Passing an argument at call + time still overrides the binding (same dtype/shape required).""" name = "TensorRead" + def __init__(self, source=None): + self._bound = source + def emit_read(self, state, mem, n_inputs): t = state.dtype.ctype if n_inputs > 1: @@ -424,8 +439,14 @@ def token(self, state): class TensorWrite(_WriteOp): + """Write output. TensorWrite(out) BINDS a concrete destination buffer at + compose() time (see TensorRead): the bound kernel writes there when no + out= is passed. out= at call time still overrides the binding.""" name = "TensorWrite" + def __init__(self, dest=None): + self._bound = dest + def emit_write(self, state, mem, pbase): t = state.dtype.ctype if state.planes > 1: @@ -443,9 +464,13 @@ def token(self, state): class TensorSplit(_WriteOp): """Write vector pixels as separate planes (planar layout for DNNs). uchar3 HxW -> 3 planes of HxW uchar. planes = batch (thread.z); - color_planes = channels (split offsets).""" + color_planes = channels (split offsets). TensorSplit(out) binds a + concrete destination at compose() time (see TensorWrite).""" name = "TensorSplit" + def __init__(self, dest=None): + self._bound = dest + def out_shape(self, shape): return shape # planes encoded by channels at alloc time @@ -596,9 +621,13 @@ def token(self, state): class SplitWrite(_WriteOp): """Write a vector-pixel stream into per-channel 2D planes laid out - contiguously (like TensorSplit but via SplitWrite params).""" + contiguously (like TensorSplit but via SplitWrite params). + SplitWrite(out) binds a concrete destination (see TensorWrite).""" name = "SplitWrite" + def __init__(self, dest=None): + self._bound = dest + def emit_write(self, state, mem, pbase): base_t = state.dtype.base_ctype ch = state.dtype.channels @@ -687,13 +716,15 @@ def token(self, state): class TensorPack(_ReadOp): """Read PLANAR (CHW) data as packed vector pixels: the inverse of - TensorSplit. Input buffer must be planar with `channels` planes.""" + TensorSplit. Input buffer must be planar with `channels` planes. + TensorPack(ch, source) binds a concrete input (see TensorRead).""" name = "TensorPack" - def __init__(self, channels: int): + def __init__(self, channels: int, source=None): if channels not in (2, 3, 4): raise ValueError("TensorPack needs 2..4 channels") self._ch = channels + self._bound = source def out_dtype(self, dt): return dt.with_channels(self._ch) @@ -712,9 +743,13 @@ def token(self, state): class TensorTSplit(_WriteOp): """Write to T3D transposed-planes layout (color_planes outermost): - layout [C][planes][H][W] instead of TensorSplit's [planes][C][H][W].""" + layout [C][planes][H][W] instead of TensorSplit's [planes][C][H][W]. + TensorTSplit(out) binds a concrete destination (see TensorWrite).""" name = "TensorTSplit" + def __init__(self, dest=None): + self._bound = dest + def emit_write(self, state, mem, pbase): base_t = state.dtype.base_ctype ch = state.dtype.channels diff --git a/fkl/tensor.py b/fkl/tensor.py index e028029..0a24ba1 100644 --- a/fkl/tensor.py +++ b/fkl/tensor.py @@ -3,6 +3,10 @@ Supports torch.Tensor(cuda), cupy.ndarray, numba, and the built-in DeviceBuffer via __cuda_array_interface__. Trailing dims of size 2..4 map to CUDA vector pixel types (uchar3, float4, ...) per fkl.types.from_cai. + +Raw device pointers (C-style integrations) enter through +DeviceBuffer.from_ptr(ptr, shape, dtype): a NON-OWNING wrap that gives the +pointer the full interop surface (CAI, DLPack, compose-time binding). """ from __future__ import annotations from dataclasses import dataclass @@ -108,6 +112,7 @@ class DeviceBuffer: allocation happens with that device's primary context current).""" _cuda = None _ctxs = {} + _owns = True # False for from_ptr views (never free external memory) def __init__(self, width: int, height: int = 1, dtype="float32", channels: int = 1, planes: int = 1, device: int = 0): @@ -128,6 +133,72 @@ def from_state(cls, st, device: int = 0): return cls(st.width, st.height, st.dtype.base, channels=st.dtype.channels, planes=st.planes, device=device) + @classmethod + def from_ptr(cls, ptr: int, shape, dtype="float32", stream=None, + device: int = 0) -> "DeviceBuffer": + """NON-OWNING wrap of an external raw device pointer. + + For C-style integrations: hand a device pointer straight into the + fkl world (compose-time binding, kernel I/O, DLPack export) with no + copies and no ownership transfer. + + buf = fkl.DeviceBuffer.from_ptr(ptr, (H, W, 3), "uint8") + k = fkl.compose(fkl.TensorRead(buf), ..., fkl.TensorWrite(dst)) + + ptr : raw CUDA device pointer (int, non-null), C-contiguous data. + shape : array dims with the library-wide folding rules of + __cuda_array_interface__ — (W,), (H, W), (planes, H, W); + a trailing dim of 2..4 folds into a vector pixel type. + A vector dtype spec ('uint8x3') may be passed instead of + the trailing dim — never together with it (that would + double-count the channels; it raises TypeError). + stream : optional producer stream (int handle / torch / cupy + stream). Recorded in __cuda_array_interface__ (v3) so + stream-aware consumers synchronize with it; a handle of + 0 (the legacy default stream) is advertised as 1, the + spelling CAI v3 requires. fkl kernels still take their + launch stream per call. + device : CUDA device ordinal the pointer lives on. + + LIFETIME: the wrapper never frees `ptr` (not on __del__, not via a + DLPack consumer's deleter). The CALLER owns the allocation and must + keep it alive for as long as the wrapper — or any kernel bound to + it — is in use.""" + if ptr is None or int(ptr) == 0: + raise ValueError("from_ptr requires a non-null device pointer") + cls._load_driver() + from .types import dtype as _d, from_shape + dt0 = _d(dtype) # accepts 'float32', 'uint8x3', ('uint8', 3), DType + if dt0.channels > 1: + dims = list(shape) + if len(dims) >= 2 and dims[-1] == dt0.channels: + # (H, W, 3) + 'uint8x3' would double-count the channels and + # silently size the wrapper channels x the real allocation + raise TypeError( + f"ambiguous spec: trailing dim of shape {tuple(shape)} " + f"matches the channels of vector dtype {dt0}; drop the " + f"trailing dim or pass the base dtype {dt0.base!r}") + if len(dims) == 1: + w, h, p = dims[0], 1, 1 + elif len(dims) == 2: + (h, w), p = dims, 1 + elif len(dims) == 3: + p, h, w = dims + else: + raise TypeError(f"unsupported shape {tuple(shape)} for {dt0}") + dt = dt0 + else: + dt, w, h, p = from_shape(tuple(shape), dt0.base) + self = cls.__new__(cls) + self.device = int(device) + self.dtype = dt + self.width, self.height, self.planes = int(w), int(h), int(p) + self._nbytes = self.width * self.height * self.planes * dt.itemsize + self._dptr = ctypes.c_void_p(int(ptr)) + self._owns = False + self._stream = stream + return self + @classmethod def _activate(cls, device: int): """Make `device`'s primary context current for this thread.""" @@ -190,8 +261,17 @@ def __cuda_array_interface__(self): dims.append(self.width) if self.dtype.channels > 1: dims.append(self.dtype.channels) - return {"shape": tuple(dims), "typestr": ts, - "data": (self.ptr, False), "version": 3} + cai = {"shape": tuple(dims), "typestr": ts, + "data": (self.ptr, False), "version": 3} + # producer stream recorded by from_ptr(stream=...) — CAI v3 allows + # advertising it so stream-aware consumers synchronize before use. + # The spec disallows the raw 0: the legacy default stream is spelled + # 1, so a caller's stream=0 keeps its synchronization contract. + s = getattr(self, "_stream", None) + if s is not None: + h = stream_handle(s) + cai["stream"] = h if h else 1 + return cai # ---- DLPack export (zero-copy into torch/cupy/jax) ---------------------- # torch.from_dlpack(buf) / cupy.from_dlpack(buf) reuse OUR device memory: @@ -206,6 +286,8 @@ def __dlpack__(self, stream=None): def __del__(self): try: if getattr(self, "_dptr", None) and self._dptr.value: + if not getattr(self, "_owns", True): + return # from_ptr view: the caller owns the memory if getattr(self, "_exported_dlpack", False): return # ownership moved to the DLPack consumer self._cuda.cuMemFree_v2(self._dptr) @@ -289,7 +371,10 @@ def _deleter(mt_ptr): if entry is not None: b = entry[2] try: - if getattr(b, "_dptr", None) and b._dptr.value: + # from_ptr views never own the memory: the DLPack consumer + # gets a zero-copy view, the external owner keeps the bytes + if (getattr(b, "_owns", True) + and getattr(b, "_dptr", None) and b._dptr.value): b._cuda.cuMemFree_v2(b._dptr) b._dptr = ctypes.c_void_p() except Exception: diff --git a/skills/fkl-python-testing/SKILL.md b/skills/fkl-python-testing/SKILL.md index a6b9697..4ab9fe8 100644 --- a/skills/fkl-python-testing/SKILL.md +++ b/skills/fkl-python-testing/SKILL.md @@ -18,6 +18,7 @@ python3 test_horizontal_fusion.py # HF: batch crops -> planes python3 test_operations.py # every op vs CPU reference python3 test_matrix.py # legacy broad matrix python3 test_e2e.py # timing + cache behaviour +python3 test_bound_compose.py # compose-time binding + from_ptr ``` Each suite exits 0/1 and prints [PASS]/[FAIL] per case. First run compiles diff --git a/skills/fkl-python-usage/SKILL.md b/skills/fkl-python-usage/SKILL.md index cc18d9b..e6ff9b7 100644 --- a/skills/fkl-python-usage/SKILL.md +++ b/skills/fkl-python-usage/SKILL.md @@ -52,6 +52,34 @@ 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. +## Compose-time buffer binding (argument-free kernels) + +Read/write ops optionally take a concrete buffer at construction. A chain +with a bound read compiles EAGERLY inside compose() (dtype/shape are known +then) and runs with no arguments: + +```python +k = fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), fkl.TensorWrite(out)) +k() # reads x, writes out — ready before the first call +k(stream=s) # argless + async on a caller-owned stream +k(y); k(y, out=z) # overrides: validated against the compiled signature +``` + +Rules: +- Binding is by POINTER: refill x in place and just call k() again. +- Bound buffers are call defaults. Overrides with a different dtype/shape + raise ValueError (a bound kernel is compiled for ONE signature; compose + an unbound chain for multi-signature use). +- Read-only binding auto-allocates the output per call; write-only binding + keeps lazy pipe(x) with the bound buffer as default out=. +- GPU-only (target="cpu" rejects bindings); lists (batch HF) cannot be + bound — pass them at call time. +- TensorSplit/SplitWrite/TensorTSplit take dest the same way; + TensorPack(ch, source) binds a planar input. +- Raw device pointers: `fkl.DeviceBuffer.from_ptr(ptr, shape, dtype)` is a + NON-OWNING zero-copy wrap (caller keeps ownership/lifetime; shape folds a + trailing 2..4 dim into a vector pixel type, or pass 'uint8x3' + (H, W)). + ## Operation catalog Memory: TensorRead(), TensorWrite(), TensorSplit() [packed->planar], diff --git a/tests/test_bound_compose.py b/tests/test_bound_compose.py new file mode 100644 index 0000000..7b0757e --- /dev/null +++ b/tests/test_bound_compose.py @@ -0,0 +1,300 @@ +"""COMPOSE-TIME BUFFER BINDING tests: TensorRead(x)/TensorWrite(out) bind +concrete buffers at compose(); the kernel compiles EAGERLY and runs with no +arguments (k() / k(stream=s)). Explicit args still override the binding, +validated against the compiled signature. Plus DeviceBuffer.from_ptr +(non-owning raw-pointer wrap) bound straight into a chain. +""" +import ctypes + +from harness import dev_f32, dev_u8, unf32, unu8, f32, check, check_true, run +import fkl + + +def t_bound_argless(): + """read+write bound at compose -> k() runs with no arguments.""" + W, H = 33, 17 # non-pow2 to catch pitch errors + src = [float((y * W + x) % 97) for y in range(H) for x in range(W)] + x = dev_f32(src, W, H) + out = fkl.DeviceBuffer(W, H, "float32") + k = fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), fkl.Add(1.0), + fkl.TensorWrite(out)) + ret = k() + check_true("bound: k() returns the bound output", ret is out) + check("bound: argless call computes", unf32(out.copy_to_host(), W * H), + [v * 2 + 1 for v in src]) + # the binding is by POINTER: update the input in place, rerun argless + src2 = [v + 1 for v in src] + x.copy_from_host(f32(src2)) + k() + check("bound: rerun sees updated input data", + unf32(out.copy_to_host(), W * H), [v * 2 + 1 for v in src2]) + + +def t_bound_stream(): + """k(stream=s) on a bound kernel: async launch on a caller-owned + driver-API stream, caller syncs.""" + W = 64 + src = [float(i % 53) for i in range(W)] + x, out = dev_f32(src, W), fkl.DeviceBuffer(W, 1, "float32") + k = fkl.compose(fkl.TensorRead(x), fkl.Mul(4.0), fkl.TensorWrite(out)) + + cu = fkl.DeviceBuffer._cuda + cu.cuStreamCreate.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_uint] + cu.cuStreamSynchronize.argtypes = [ctypes.c_void_p] + cu.cuStreamDestroy_v2.argtypes = [ctypes.c_void_p] + s = ctypes.c_void_p() + check_true("bound: cuStreamCreate", cu.cuStreamCreate(ctypes.byref(s), 0) == 0) + k(stream=int(s.value)) # async: caller owns the stream + cu.cuStreamSynchronize(s) + check("bound: k(stream=s) computes", unf32(out.copy_to_host(), W), + [v * 4 for v in src]) + cu.cuStreamDestroy_v2(s) + + +def t_eager_compile(): + """Binding compiles at compose() time; a second compose of the same + chain signature is a pure cache hit (same .so).""" + W = 32 + src = [float(i) for i in range(W)] + x, out = dev_f32(src, W), fkl.DeviceBuffer(W, 1, "float32") + k = fkl.compose(fkl.TensorRead(x), fkl.Mul(3.0), fkl.TensorWrite(out)) + check_true("eager: compiled inside compose() (before any call)", + len(k._variants) == 1) + so1 = list(k._variants.values())[0][3] + import time + t0 = time.perf_counter() + k2 = fkl.compose(fkl.TensorRead(x), fkl.Mul(7.0), fkl.TensorWrite(out)) + ms = (time.perf_counter() - t0) * 1e3 + so2 = list(k2._variants.values())[0][3] + check_true("eager: second compose reuses the cached .so", + so1 == so2, f"compose took {ms:.1f} ms") + k2() + check("eager: cached kernel computes with its own values", + unf32(out.copy_to_host(), W), [v * 7 for v in src]) + + +def t_read_only_bound(): + """Binding only the read: k() auto-allocates the output per call.""" + W = 16 + src = [float(i) for i in range(W)] + k = fkl.compose(fkl.TensorRead(dev_f32(src, W)), fkl.Mul(2.0), + fkl.TensorWrite()) + out = k() + check("bound: read-only binding auto-allocates output", + unf32(out.copy_to_host(), W), [v * 2 for v in src]) + + +def t_write_only_bound(): + """Binding only the write: pipe(x) uses the bound buffer as default out + (lazy compile as usual — input signature arrives with x).""" + W = 16 + src = [float(i) for i in range(W)] + out = fkl.DeviceBuffer(W, 1, "float32") + k = fkl.compose(fkl.TensorRead(), fkl.Sub(1.0), fkl.TensorWrite(out)) + check_true("bound: write-only binding stays lazy", len(k._variants) == 0) + ret = k(dev_f32(src, W)) + check_true("bound: write-only binding is the default out", ret is out) + check("bound: write-only values", unf32(out.copy_to_host(), W), + [v - 1 for v in src]) + + +def t_override_at_call(): + """Explicit args override the binding (same compiled signature).""" + W = 24 + a = [float(i % 19) for i in range(W)] + b = [float(i * 3 % 23) for i in range(W)] + x, out = dev_f32(a, W), fkl.DeviceBuffer(W, 1, "float32") + k = fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), fkl.TensorWrite(out)) + y = dev_f32(b, W) + r1 = k(y) # input override -> still writes bound out + check_true("override: k(y) writes the bound output", r1 is out) + check("override: k(y) values", unf32(out.copy_to_host(), W), + [v * 2 for v in b]) + z = fkl.DeviceBuffer(W, 1, "float32") + r2 = k(y, out=z) # both overridden + check_true("override: k(y, out=z) returns z", r2 is z) + check("override: k(y, out=z) values", unf32(z.copy_to_host(), W), + [v * 2 for v in b]) + k() # bindings intact after overrides + check("override: k() still uses the bindings", + unf32(out.copy_to_host(), W), [v * 2 for v in a]) + + +def t_bound_uchar3_split(): + """Bound DNN-ingest chain: uchar3 -> normalize -> planar CHW, argless. + TensorSplit's destination layout is (C, H, W) of the base dtype.""" + W, H = 8, 4 + n = W * H * 3 + src = [(i * 7) % 256 for i in range(n)] + x = dev_u8(src, W, H, ch=3) + out = fkl.DeviceBuffer(W, H, "float32", planes=3) # 3 planes of HxW + k = fkl.compose(fkl.TensorRead(x), fkl.Cast("float32"), + fkl.Div((255.0,) * 3), fkl.TensorSplit(out)) + k() + got = unf32(out.copy_to_host(), n) + exp = [src[p * 3 + c] / 255.0 + for c in range(3) for p in range(W * H)] + check("bound: uchar3 -> planar CHW argless", got, exp) + + +def t_from_ptr_roundtrip(): + """from_ptr wraps a DeviceBuffer's own pointer: same bytes, non-owning, + usable as a bound chain input.""" + W, H = 8, 4 + src = [float((y * W + x) % 31) for y in range(H) for x in range(W)] + owner = dev_f32(src, W, H) + view = fkl.DeviceBuffer.from_ptr(owner.ptr, (H, W), "float32") + check_true("from_ptr: same pointer, non-owning", + view.ptr == owner.ptr and not view._owns) + check("from_ptr: reads the owner's bytes", + unf32(view.copy_to_host(), W * H), src) + out = fkl.DeviceBuffer(W, H, "float32") + fkl.compose(fkl.TensorRead(view), fkl.Mul(2.0), fkl.TensorWrite(out))() + check("from_ptr: kernel reads through the wrapper", + unf32(out.copy_to_host(), W * H), [v * 2 for v in src]) + del view # non-owning: must NOT free the owner's memory + check("from_ptr: owner memory survives wrapper deletion", + unf32(owner.copy_to_host(), W * H), src) + + +def t_from_ptr_output_and_vector(): + """from_ptr as the bound OUTPUT (raw destination pointer) and with a + vector dtype spec ('uint8x3' == shape (H, W, 3)).""" + W, H = 4, 4 + n = W * H * 3 + src = [(i * 5) % 256 for i in range(n)] + x = dev_u8(src, W, H, ch=3) + dst = fkl.DeviceBuffer(W, H, "uint8", channels=3) + wrap = fkl.DeviceBuffer.from_ptr(dst.ptr, (H, W), "uint8x3") + check_true("from_ptr: vector dtype spec folds channels", + wrap.dtype.ctype == "uchar3" and wrap.width == W + and wrap.height == H) + fkl.compose(fkl.TensorRead(x), fkl.VectorReorder(2, 1, 0), + fkl.TensorWrite(wrap))() + got = [float(v) for v in unu8(dst.copy_to_host(), n)] + exp = [float(src[p * 3 + (2 - c)]) for p in range(W * H) for c in range(3)] + check("from_ptr: raw-pointer output receives results", got, exp) + + +def t_from_ptr_guards(): + """from_ptr argument validation: null pointers, ambiguous vector-dtype + + trailing-channel-dim specs, and CAI v3 stream spelling.""" + W, H = 8, 4 + owner = fkl.DeviceBuffer(W, H, "uint8", channels=3) + + # NULL is never a valid device pointer — fail at the call site, not + # later in .ptr / copy_to_host + for bad in (0, None): + try: + fkl.DeviceBuffer.from_ptr(bad, (H, W), "uint8x3") + check_true(f"from_ptr: null ptr ({bad!r}) rejected", False) + except ValueError as e: + check_true(f"from_ptr: null ptr ({bad!r}) rejected", + "non-null" in str(e), str(e)[:60]) + + # vector dtype AND the trailing channel dim would double-count the + # channels ((H, W, 3) x uchar3 = 3x the real allocation) — reject + try: + fkl.DeviceBuffer.from_ptr(owner.ptr, (H, W, 3), "uint8x3") + check_true("from_ptr: vector dtype + trailing dim rejected", False) + except TypeError as e: + check_true("from_ptr: vector dtype + trailing dim rejected", + "ambiguous" in str(e), str(e)[:60]) + + # CAI v3 stream key: 0 (legacy default stream) must be spelled 1, + # None means "no synchronization needed" (key omitted) + cai = fkl.DeviceBuffer.from_ptr(owner.ptr, (H, W), "uint8x3", + stream=0).__cuda_array_interface__ + check_true("from_ptr: stream=0 advertised as legacy default (1)", + cai.get("stream") == 1) + cai = fkl.DeviceBuffer.from_ptr(owner.ptr, (H, W), "uint8x3", + stream=None).__cuda_array_interface__ + check_true("from_ptr: stream=None omits the CAI stream key", + "stream" not in cai) + cai = fkl.DeviceBuffer.from_ptr(owner.ptr, (H, W), "uint8x3", + stream=1234).__cuda_array_interface__ + check_true("from_ptr: explicit stream handle advertised as-is", + cai.get("stream") == 1234) + + +def t_errors(): + W = 16 + src = [float(i) for i in range(W)] + x, out = dev_f32(src, W), fkl.DeviceBuffer(W, 1, "float32") + + # argless call without a bound input + k = fkl.compose(fkl.TensorRead(), fkl.Mul(2.0), fkl.TensorWrite()) + try: + k() + check_true("err: argless call without binding raises", False) + except TypeError as e: + check_true("err: argless call without binding raises", + "bound" in str(e), str(e)[:60]) + + kb = fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), fkl.TensorWrite(out)) + # dtype mismatch on input override + try: + kb(fkl.DeviceBuffer(W, 1, "int32")) + check_true("err: input override dtype mismatch", False) + except ValueError as e: + check_true("err: input override dtype mismatch", + "mismatch" in str(e), str(e)[:60]) + # shape mismatch on input override + try: + kb(fkl.DeviceBuffer(W * 2, 1, "float32")) + check_true("err: input override shape mismatch", False) + except ValueError: + check_true("err: input override shape mismatch", True) + # out= override mismatch on a bound kernel + try: + kb(out=fkl.DeviceBuffer(W // 2, 1, "float32")) + check_true("err: out= override size mismatch", False) + except ValueError: + check_true("err: out= override size mismatch", True) + # batch (list) override of an eagerly-compiled bound kernel + try: + kb([x, x]) + check_true("err: batch override of bound kernel", False) + except ValueError: + check_true("err: batch override of bound kernel", True) + # bound output validated at compose() time + try: + fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), + fkl.TensorWrite(fkl.DeviceBuffer(W // 2, 1, "float32"))) + check_true("err: bound output size mismatch at compose", False) + except ValueError: + check_true("err: bound output size mismatch at compose", True) + # binding is GPU-only (CPU chains take numpy arrays per call) + try: + fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), fkl.TensorWrite(), + target="cpu") + check_true("err: cpu target rejects binding", False) + except ValueError: + check_true("err: cpu target rejects binding", True) + # a batch (list) cannot be bound — pass it at call time instead + try: + fkl.compose(fkl.TensorRead([x, x]), fkl.Mul(2.0), fkl.TensorWrite()) + check_true("err: list binding rejected at compose", False) + except TypeError: + check_true("err: list binding rejected at compose", True) + # divergent chains do not take bindings (batch arrives at call time) + try: + fkl.compose_divergent( + [1, 1], + [fkl.TensorRead(x), fkl.Mul(2.0), fkl.TensorWrite()]) + check_true("err: compose_divergent rejects binding", False) + except ValueError: + check_true("err: compose_divergent rejects binding", True) + # unbound behavior unchanged: lazy compile, explicit args required + check_true("unbound: compose stays lazy", len(k._variants) == 0) + o = k(dev_f32(src, W)) + check("unbound: pipe(x) path intact", unf32(o.copy_to_host(), W), + [v * 2 for v in src]) + + +if __name__ == "__main__": + run([t_bound_argless, t_bound_stream, t_eager_compile, + t_read_only_bound, t_write_only_bound, t_override_at_call, + t_bound_uchar3_split, t_from_ptr_roundtrip, + t_from_ptr_output_and_vector, t_from_ptr_guards, t_errors], + "bound-compose") diff --git a/tests/test_torch_integration.py b/tests/test_torch_integration.py index 1ec1773..a71b04a 100644 --- a/tests/test_torch_integration.py +++ b/tests/test_torch_integration.py @@ -96,8 +96,39 @@ def t_batch_hf_from_torch_list(): and out[0].max().item() == 10.0 and out[2].max().item() == 30.0) +def t_bound_compose_torch(): + """Compose-time binding with torch tensors: eager compile, argless call, + argless call on an external torch stream.""" + x = torch.full((8, 16), 3.0, device="cuda", dtype=torch.float32) + out = torch.empty_like(x) + k = fkl.compose(fkl.TensorRead(x), fkl.Mul(2.0), fkl.Add(1.0), + fkl.TensorWrite(out)) + check_true("bound compose: eager compile at compose()", + len(k._variants) == 1) + ret = k() # no arguments at all + check_true("bound compose: torch argless call", + ret is out and torch.allclose(out, torch.full_like(out, 7.0))) + x.fill_(5.0) # binding is by pointer + s = torch.cuda.Stream() + with torch.cuda.stream(s): + k(stream=s.cuda_stream) # argless + async + s.synchronize() + check_true("bound compose: argless on external torch stream", + torch.allclose(out, torch.full_like(out, 11.0))) + y = torch.full((8, 16), 1.0, device="cuda") + k(y) # override the binding + check_true("bound compose: override writes bound out", + torch.allclose(out, torch.full_like(out, 3.0))) + try: + k(torch.full((4, 4), 1.0, device="cuda")) + check_true("bound compose: shape override mismatch raises", False) + except ValueError: + check_true("bound compose: shape override mismatch raises", True) + + if __name__ == "__main__": run([t_torch_in_dlpack_out, t_dlpack_zero_copy_roundtrip, t_torch_stream_async, t_out_preallocated_torch, t_uchar3_image_preproc_to_torch, t_circular_tensor_to_torch, - t_batch_hf_from_torch_list], "torch-integration") + t_batch_hf_from_torch_list, t_bound_compose_torch], + "torch-integration")