From ec0869e0a714f175c53f7da450588c8ad2550004 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Mon, 12 Oct 2026 16:20:51 -0400 Subject: [PATCH] Conversion routes say what happens; the observer belongs to its context StaysInPlace ("stays in place") is the one statement of "same library, same device: nothing is converted". planned returns it, the run goes through it, and move_to_device, to_numpy and from_numpy ask it instead of restating it. DeviceMove applies only when the device differs. SharedHostBuffer ("shared host buffer") names a host exchange in which the result is the source's own buffer. Libraries declare it by inheriting NumpyBufferSharing (NumPy, torch, JAX; TensorFlow copies both ways, measured). Whether one array was shared is read after the exchange from where the two arrays begin; a copy is recorded as "host copy" with the declined route named. to_numpy and from_numpy are conversions through ConversionRoute.run, so the record is written in one place. ConversionObserver.current is a ContextVar read through the class: each thread and asyncio task has its own observer, and a thread run in a copied context sees the spawner's. Co-Authored-By: Claude Opus 5.5 --- ADDING_NEW_FRAMEWORKS.md | 1 + CHANGELOG.md | 106 ++++++++++- README.md | 24 ++- docs/source/api_reference.rst | 31 +++- docs/source/converters.rst | 8 +- src/arraybridge/conversion.py | 190 ++++++++++++++++--- src/arraybridge/libraries.py | 109 ++++++----- stubs/jax/__init__.pyi | 1 + stubs/numpy/__init__.pyi | 4 + stubs/torch/__init__.pyi | 1 + tests/fakes.py | 6 +- tests/test_conversion.py | 340 +++++++++++++++++++++++++++++++++- tests/test_toy_library.py | 2 +- 13 files changed, 716 insertions(+), 107 deletions(-) diff --git a/ADDING_NEW_FRAMEWORKS.md b/ADDING_NEW_FRAMEWORKS.md index 3911db1..7d88b74 100644 --- a/ADDING_NEW_FRAMEWORKS.md +++ b/ADDING_NEW_FRAMEWORKS.md @@ -176,6 +176,7 @@ class must define. | Parent | The library... | Abstract methods | |---|---|---| | `HostDevices` | keeps arrays in main memory. | `_wrap_host_array(data)`, `_move_to_host(data)` | +| `NumpyBufferSharing` | is a `HostDevices` library whose main-memory array and a NumPy array can be one buffer (measure it before declaring it). | those of `HostDevices`, plus `_first_element_address(data)` | | `AcceleratorDevices` | can keep arrays on numbered accelerators. | `_accelerator_count()`, `_accelerator_scope(device)`, `_move_to_accelerator(data, device)`, `_upload(data, device)`, `_device_if_accepted(value)` | | `CudaLibrary` | is an `AcceleratorDevices` library with its own CUDA streams and a memory cache to release. | those of `AcceleratorDevices`, plus `create_stream()`, `stream_scope(stream)`, `current_device()`, `_free_cached_memory()` | | `DLPackImport` | can wrap another library's buffer through DLPack. | `_import_dlpack(payload)` | diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c6639a..5c3c1a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -77,12 +77,112 @@ two different libraries. `MemoryType`'s metaclass overrides are annotated with `EnumMeta`'s signatures; its `__call__` override is gone (`Enum` already asks `_missing_`). +### Conversion + +The mechanism names a consumer can now see, in the order the routes are tried: + +| Mechanism | Route class | Meaning | +|---|---|---| +| `"stays in place"` | `StaysInPlace` (new) | Same library, same device. Nothing is converted; the input object is returned. No record. | +| `"device move"` | `DeviceMove` | Same library, another device. | +| `"dlpack"` | `DLPackExchange` | The target wraps the source's buffer through DLPack. | +| `"shared host buffer"` | `SharedHostBuffer` (new) | Both arrays are in main memory and the result begins at the same address as the source: no bytes were copied. | +| `"host copy"` | `HostCopy` | The values were copied through a NumPy array in main memory. | + +**The plan and the run agree that nothing is converted.** +`ConversionRoute.planned(library, device, library, device)` returns +`StaysInPlace`. In 0.5.0 it returned `DeviceMove`, while running that +conversion returned the input and recorded nothing, so a caller had to restate +the rule. `DeviceMove` now applies only when the device differs. The rule is +stated once, in `StaysInPlace.applies`; `convert_to`, `to_numpy`, `from_numpy` +and `move_to_device` all ask it. Every route has `writes_a_record`; it is +`False` for `StaysInPlace` only. A record means data was moved, copied or +wrapped. + +**A host exchange that shares the buffer is no longer called a copy.** On the +host, torch and NumPy share one buffer in both directions, and JAX hands NumPy +its own buffer; 0.5.0 recorded these as `"host copy"`. They are now recorded +as `"shared host buffer"`. A library declares that it can share by inheriting +`NumpyBufferSharing` (NumPy, torch and JAX do; TensorFlow copies in both +directions and does not). Whether one array was shared is decided after the +exchange by comparing where the two arrays begin (`host_address_of`). + +What was measured (Python 3.12, NumPy 2.5.4, torch 2.14.1 CPU, JAX 0.11.2, +TensorFlow 2.21.0 CPU; pointers, `np.shares_memory` and writing to one array +and reading the other): + +| Exchange on the host | Result | +|---|---| +| NumPy to torch | shared, for every dtype tried, Fortran order, stepped and offset views, read-only arrays. A view with a negative step is copied (arraybridge makes it contiguous first). | +| torch to NumPy | shared | +| JAX to NumPy | shared (the NumPy array is read-only) | +| NumPy to JAX | copied, unless the buffer is C-contiguous, starts on a 64-byte boundary and keeps its dtype; then shared | +| NumPy to TensorFlow, TensorFlow to NumPy | copied | +| an array with no elements | nothing to share: `"host copy"` of 0 bytes | + +**What `planned` returns when sharing depends on the array.** A plan sees +only declarations and devices, never an array. It returns the route that is +tried first: `SharedHostBuffer` for NumPy, torch and JAX with NumPy on the +host. The record names what was done. Where the exchange copied (NumPy to JAX +for an ordinary NumPy allocation, a reversed view to torch) the record's +mechanism is `"host copy"` and `declined_mechanisms` is +`("shared host buffer",)`, the way a declined DLPack exchange is already +reported. The exchange is done once either way. + +What a consumer that keyed on the old answers sees: + +- `planned(...) is DeviceMove` is no longer true for the same library and + device; test `is StaysInPlace` (or `route.writes_a_record`). +- Rows that were `"host copy"` for NumPy and torch on the host, and for JAX + to NumPy on the host, are now `"shared host buffer"`. `"host copy"` rows + that follow a declined `SharedHostBuffer` carry it in `declined_mechanisms`, + which `ConversionCounters` does not key on. +- `planned` for NumPy with torch or JAX on the host (either direction) is + `SharedHostBuffer`. Between torch and JAX on the host it is still `"dlpack"`. + Anything to or from an accelerator, pyclesperanto or TensorFlow is planned + as before. +- The `nbytes` of a `from_numpy` record is the size of the array that was + made (it was the size of the NumPy input; they differ when JAX narrows + float64 to float32). + +`ConversionRoute.taken(data, result, source, target)` is new: the route that +names what `convert` did, the route itself unless it says otherwise. +`ArrayLibrary.shares_numpy_buffer_on(device)` and +`ArrayLibrary.host_address_of(data)` are new, as is +`NumpyLibrary.registered()`. `HostCopy` and `SharedHostBuffer` share the new +parent `HostExchange`, which has no mechanism of its own. +`arraybridge.libraries` no longer imports `ConversionObserver`, +`ConversionRecord` and `HostCopy` (they were never exported from it; they are +in `arraybridge.conversion`). + +**`ConversionObserver.current` belongs to the context that reads it.** It was +one attribute for the whole process, and `installed()` put back whatever it +had seen on entry: two threads that entered and left in the same order left +the second thread's observer installed for the rest of the process. The hook +is now a `contextvars.ContextVar`. + +- `ConversionObserver.current` reads the observer of the running thread or + asyncio task. Assigning to it sets the observer of that context only. +- `observer.install()` makes it the observer of this context and returns the + one it replaced. `with observer.installed():` does so for a block. +- A new thread starts with the default observer, which ignores conversions. + **An observer installed in the main thread no longer sees conversions made + in worker threads.** To observe a thread you start, run it in a copy of your + context: `Thread(target=contextvars.copy_context().run, args=(work,))`, or + `executor.submit(contextvars.copy_context().run, work)`. +- An asyncio task starts with the observer of the code that created it, and + what it installs stays in that task. +- `current` is read on the class (`ConversionObserver.current`). It is no + longer an attribute of an observer instance. + ### Not verified on a GPU 0.5.1 was tested with NumPy on Python 3.10, 3.12 and 3.13 and with the torch -CPU wheel on 3.12. CuPy, TensorFlow, JAX and pyclesperanto, and torch on CUDA, -were exercised only through the fakes in `tests/fakes.py`. `TESTING.md` has -the commands for a GPU machine. +CPU wheel on 3.12. The host exchanges in the Conversion section were also run +against the JAX and TensorFlow CPU builds named there. CuPy and pyclesperanto, +every accelerator path of torch, TensorFlow and JAX, and everything else +TensorFlow and JAX do, were exercised only through the fakes in +`tests/fakes.py`. `TESTING.md` has the commands for a GPU machine. ## 0.5.0 diff --git a/README.md b/README.md index 308a430..0b0d859 100644 --- a/README.md +++ b/README.md @@ -99,10 +99,13 @@ library does not have. ## Which route a conversion takes -A conversion tries three mechanisms in order: `"device move"` (same library, -other device), `"dlpack"` (the target wraps the source's buffer) and -`"host copy"` (download to NumPy, then upload). `ConversionRoute.planned` -tells which one will be tried first. It reads only the library classes and the +A conversion tries five mechanisms in order: `"stays in place"` (same library +and device: nothing is converted), `"device move"` (same library, other +device), `"dlpack"` (the target wraps the source's buffer), +`"shared host buffer"` (both arrays are in main memory and the result is the +source's own buffer) and `"host copy"` (the values are copied through NumPy). +`ConversionRoute.planned` tells which one will be tried first; the record of +a conversion names the one that was used. It reads only the library classes and the device values, so it works without the libraries or the hardware: ```python @@ -115,11 +118,12 @@ gpu = AcceleratorDevice(0) assert ConversionRoute.planned(torch, gpu, cupy, gpu).mechanism == "dlpack" assert ConversionRoute.planned(torch, HOST, cupy, gpu).mechanism == "host copy" assert ConversionRoute.planned(torch, HOST, torch, gpu).mechanism == "device move" +assert ConversionRoute.planned(torch, gpu, torch, gpu).mechanism == "stays in place" ``` To see what conversions actually ran, install a `ConversionCounters` observer. It keeps a count, bytes and seconds per source, target and mechanism. This -example needs JAX: +example needs TensorFlow: ```python import numpy as np @@ -128,16 +132,18 @@ from arraybridge import HOST, ConversionCounters, convert_memory counters = ConversionCounters() with counters.installed(): - convert_memory(np.zeros((4, 4), dtype=np.float32), "numpy", "jax", HOST) + convert_memory(np.zeros((4, 4), dtype=np.float32), "numpy", "tensorflow", HOST) row = counters.rows()[0] -assert (row["source_library"], row["target_library"]) == ("numpy", "jax") +assert (row["source_library"], row["target_library"]) == ("numpy", "tensorflow") assert (row["mechanism"], row["count"], row["nbytes"]) == ("host copy", 1, 64) print(counters.format_table()) ``` -A call that returns its input unchanged is not counted. With no observer -installed nothing is timed. +A conversion that stays in place is not counted. With no observer installed +nothing is timed. An observer belongs to the thread or asyncio task that +installed it; a thread you start is observed when you run it in a copy of your +context (`contextvars.copy_context().run`). ## Array operations in the array's own library diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst index f5a3d53..0e1120c 100644 --- a/docs/source/api_reference.rst +++ b/docs/source/api_reference.rst @@ -138,6 +138,10 @@ Conversion reporting ``ConversionRoute`` has one subclass per mechanism, tried in this order: +``StaysInPlace`` (``"stays in place"``) + Same library and same device. The input is returned and no record is + written (``writes_a_record`` is ``False`` for this route only). + ``DeviceMove`` (``"device move"``) Source and target are the same library; the array changes device only. @@ -145,15 +149,24 @@ Conversion reporting The target wraps the source's buffer. Applies when the target imports DLPack from the source's device and the source exports from it. +``SharedHostBuffer`` (``"shared host buffer"``) + Both arrays are in main memory and the result begins at the address of the + source: no bytes are copied. Applies when both libraries are + ``NumpyBufferSharing`` (NumPy, torch, JAX) and both devices are the host. + Whether one array is shared is known after the exchange; one that was + copied is recorded as ``"host copy"`` with ``"shared host buffer"`` under + ``declined_mechanisms``. + ``HostCopy`` (``"host copy"``) - Download to a NumPy array, then upload to the target. Always applies. + The values are copied through a NumPy array. Always applies. ``ConversionRoute.planned(source_library, source_device, target_library, target_device)`` returns the route class that a conversion of that pair tries first. It reads only the library classes and the device values: no library is imported and no device is touched. The conversion selects its route with the same test. A route that applies but cannot carry a particular array is -skipped and the next one is tried. +skipped and the next one is tried, so the record of a conversion, not the +plan, says what was done to that array. .. code-block:: python @@ -166,19 +179,21 @@ skipped and the next one is tried. assert ConversionRoute.planned(torch, gpu, cupy, gpu).mechanism == "dlpack" assert ConversionRoute.planned(torch, HOST, cupy, gpu).mechanism == "host copy" - assert ConversionRoute.planned(numpy_library, HOST, torch, HOST).mechanism == "host copy" + assert ConversionRoute.planned(numpy_library, HOST, torch, HOST).mechanism == "shared host buffer" + assert ConversionRoute.planned(torch, gpu, torch, gpu).mechanism == "stays in place" assert ConversionRoute.planned(torch, HOST, torch, gpu).mechanism == "device move" ``ConversionObserver`` Receives one ``ConversionRecord`` per conversion while it is installed. Subclass it and define ``record(record)``. ``observer.install()`` makes it - the process's observer and returns the one it replaced; - ``with observer.installed():`` does so for a block. - ``ConversionObserver.current`` is the installed observer. The default + the observer of the running thread or asyncio task and returns the one it + replaced; ``with observer.installed():`` does so for a block. + ``ConversionObserver.current`` is the observer of the context that reads + it (a ``contextvars.ContextVar``). A new thread starts with the default; + run it with ``contextvars.copy_context().run`` to give it yours. The default ignores conversions, and then nothing is timed. ``library.convert_to`` (which ``convert_memory`` calls), ``library.to_numpy`` and - ``library.from_numpy`` report. A call that returns its input unchanged does - not. + ``library.from_numpy`` report. A conversion that stays in place does not. ``ConversionRecord`` Fields: ``source_library``, ``source_device``, ``target_library``, diff --git a/docs/source/converters.rst b/docs/source/converters.rst index f6211d6..6040682 100644 --- a/docs/source/converters.rst +++ b/docs/source/converters.rst @@ -45,9 +45,11 @@ answers about them. Routes ------ -A conversion tries, in order: a device move when source and target are the -same library; a DLPack exchange when the target can wrap the source's buffer -on its device; a copy through a NumPy array in main memory. +A conversion tries, in order: nothing at all when source and target are the +same library on the same device; a device move when they are the same library; +a DLPack exchange when the target can wrap the source's buffer on its device; +an exchange through NumPy in main memory, which shares the buffer where both +libraries can (NumPy, torch, JAX) and copies otherwise. ``ConversionRoute.planned`` tells which is tried first for a pair, and a ``ConversionObserver`` receives what actually ran; both are described in :doc:`api_reference`. No pair is promised to be zero-copy. diff --git a/src/arraybridge/conversion.py b/src/arraybridge/conversion.py index 5f4bc75..f810f32 100644 --- a/src/arraybridge/conversion.py +++ b/src/arraybridge/conversion.py @@ -2,17 +2,26 @@ ``ConversionRoute`` leaves are tried in the order they are declared in this module. ``applies`` reads only the two library declarations and the two device -descriptions, so the same test answers "which mechanism will be used" on a -machine without the device and selects the route when the conversion runs. +descriptions, so the same test answers "which mechanism will be tried first" on +a machine without the device and selects the route when the conversion runs. + +A plan sees no array. Where a mechanism depends on the particular array (a +DLPack capsule the target cannot read, a NumPy buffer the target cannot wrap), +the record of the conversion names the mechanism that was used and lists the +planned one under ``declined_mechanisms``. + +A ``ConversionRecord`` is written when data was moved, copied or wrapped. The +route that leaves the array where it is (``StaysInPlace``) writes none. """ from __future__ import annotations import logging import time -from abc import ABC, abstractmethod +from abc import ABC, ABCMeta, abstractmethod from collections.abc import Iterator, Mapping from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import asdict, dataclass, field from types import NotImplementedType from typing import TYPE_CHECKING, ClassVar, TypeGuard, TypeVar @@ -53,6 +62,8 @@ class ConversionRoute(ABC, metaclass=AutoRegisterMeta): __registry_key__ = "mechanism" __registry__: ClassVar[Mapping[str, type[ConversionRoute]]] = {} mechanism: ClassVar[str | None] = None + writes_a_record: ClassVar[bool] = True + """Whether a conversion by this route is reported to the observer.""" @classmethod @abstractmethod @@ -76,6 +87,23 @@ def convert( ) -> T: """Carry ``data`` to ``target`` on ``target_device`` or raise ``RouteDeclinedError``.""" + @classmethod + def taken( + cls, + data: S, + result: T, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + ) -> type[ConversionRoute]: + """The route that names what ``convert`` did to make ``result``: this one. + + A route whose mechanism can only be known once the conversion is done + answers with another route here, and the record lists this one as declined. + """ + + del data, result, source, target + return cls + @classmethod def candidates( cls, @@ -100,7 +128,11 @@ def planned( target: ArrayLibrary[T], target_device: Device, ) -> type[ConversionRoute]: - """The route a conversion of this pair will try first.""" + """The route a conversion of this pair will try first. + + It is the route that runs unless the particular array cannot be carried + by it, which no plan can know; the record then names the one that did. + """ return cls.candidates(source, source_device, target, target_device)[0] @@ -125,14 +157,17 @@ def run( logger.debug("%s declined %s to %s: %s", route.mechanism, source, target, reason) declined.append(str(route.mechanism)) continue - if observer.wants_records and result is not data: + if observer.wants_records and route.writes_a_record: + taken = route.taken(data, result, source, target) + if taken is not route: + declined.append(str(route.mechanism)) observer.record( ConversionRecord( source_library=source.name, source_device=source_device.label, target_library=target.name, target_device=target_device.label, - mechanism=str(route.mechanism), + mechanism=str(taken.mechanism), nbytes=target.nbytes_of(result), seconds=time.perf_counter() - started, declined_mechanisms=tuple(declined), @@ -142,8 +177,42 @@ def run( raise LookupError(f"No conversion route from {source} to {target}") +class StaysInPlace(ConversionRoute): + """Same library, same device: nothing is converted and the input is returned. + + This is the one statement of that rule. ``ConversionRoute.planned`` answers + with this route, the conversion runs through it, and + ``ArrayLibrary.move_to_device`` asks ``applies``. No record is written. + """ + + mechanism = "stays in place" + writes_a_record = False + + @classmethod + def applies( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], + target_device: Device, + ) -> bool: + return source is target and source_device == target_device + + @classmethod + def convert( + cls, + data: S, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + target_device: Device, + ) -> T: + if not _is_array_of_both(data, source, target): + raise TypeError(f"{source} and {target} are different libraries") + return data + + class DeviceMove(ConversionRoute): - """Same library: the array stays native and only changes device.""" + """Same library, another device: the array stays native and changes device.""" mechanism = "device move" @@ -155,7 +224,9 @@ def applies( target: ArrayLibrary[T], target_device: Device, ) -> bool: - return source is target + return source is target and not StaysInPlace.applies( + source, source_device, target, target_device + ) @classmethod def convert( @@ -218,10 +289,29 @@ def convert( raise RouteDeclinedError(str(error)) from error -class HostCopy(ConversionRoute): - """Download to a NumPy array in main memory, then upload to the target.""" +class HostExchange(ConversionRoute): + """Through NumPy: the source hands over a NumPy array and the target takes it.""" - mechanism = "host copy" + @classmethod + def convert( + cls, + data: S, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + target_device: Device, + ) -> T: + return target.unrecorded_from_numpy(source.unrecorded_to_numpy(data), target_device) + + +class SharedHostBuffer(HostExchange): + """Both arrays are in main memory and the result is the source's own buffer. + + Declared by the two libraries (``shares_numpy_buffer_on``). Whether one + particular array is shared is known only afterwards, by comparing where the + two arrays begin; an exchange that copied is recorded as a ``HostCopy``. + """ + + mechanism = "shared host buffer" @classmethod def applies( @@ -231,22 +321,46 @@ def applies( target: ArrayLibrary[T], target_device: Device, ) -> bool: - return True + return source.shares_numpy_buffer_on(source_device) and target.shares_numpy_buffer_on( + target_device + ) @classmethod - def convert( + def taken( cls, data: S, + result: T, source: ArrayLibrary[S], target: ArrayLibrary[T], + ) -> type[ConversionRoute]: + address = source.host_address_of(data) + shared = address is not None and address == target.host_address_of(result) + return cls if shared else HostCopy + + +class HostCopy(HostExchange): + """The values are copied: downloaded to NumPy in main memory, then uploaded.""" + + mechanism = "host copy" + + @classmethod + def applies( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], target_device: Device, - ) -> T: - return target.unrecorded_from_numpy(source.unrecorded_to_numpy(data), target_device) + ) -> bool: + return True @dataclass(frozen=True, slots=True) class ConversionRecord: - """One conversion that produced a new array. Plain fields: pickles and dumps to JSON.""" + """One conversion that moved, copied or wrapped data. Plain fields: pickles, dumps to JSON. + + ``mechanism`` names what was done. ``declined_mechanisms`` lists the routes + tried before it that could not carry this array. + """ source_library: str source_device: str @@ -276,14 +390,33 @@ def from_json(cls, row: Mapping[str, str | int | float | list[str]]) -> Conversi ) -class ConversionObserver(ABC): +class ObserverOfThisContext(ABCMeta): + """Gives ``ConversionObserver.current``: read and assigned per context.""" + + @property + def current(cls) -> ConversionObserver: + """The observer of the running thread or asyncio task.""" + + return _observer.get() + + @current.setter + def current(cls, observer: ConversionObserver) -> None: + _observer.set(observer) + + +class ConversionObserver(ABC, metaclass=ObserverOfThisContext): """Receives one ``ConversionRecord`` per conversion while it is installed. - ``ConversionObserver.current`` is the one hook. Its default ignores - conversions and asks for no records, so nothing is timed or sized. + ``ConversionObserver.current`` is the one hook, and it belongs to the + context (``contextvars``) that reads it. Every thread starts with the + default, which ignores conversions and asks for no records, so nothing is + timed or sized. An asyncio task starts with the observer of the code that + created it. Installing an observer in one thread or task changes nothing + in another; to observe the conversions of a thread you start, run it in a + copy of your context: ``Thread(target=contextvars.copy_context().run, + args=(work,))``. """ - current: ClassVar[ConversionObserver] wants_records: ClassVar[bool] = True @abstractmethod @@ -292,20 +425,19 @@ def record(self, record: ConversionRecord) -> None: @contextmanager def installed(self) -> Iterator[ConversionObserver]: - """Make this the process's observer for the duration of a block.""" + """Make this the observer of this context for the duration of a block.""" - previous = ConversionObserver.current - ConversionObserver.current = self + previous = self.install() try: yield self finally: - ConversionObserver.current = previous + previous.install() def install(self) -> ConversionObserver: - """Make this the process's observer and return the one it replaced.""" + """Make this the observer of this context and return the one it replaced.""" - previous = ConversionObserver.current - ConversionObserver.current = self + previous = _observer.get() + _observer.set(self) return previous @@ -318,7 +450,9 @@ def record(self, record: ConversionRecord) -> None: del record -ConversionObserver.current = IgnoreConversions() +_observer: ContextVar[ConversionObserver] = ContextVar( + "arraybridge conversion observer", default=IgnoreConversions() +) ConversionKey = tuple[str, str, str, str, str] diff --git a/src/arraybridge/libraries.py b/src/arraybridge/libraries.py index 432edff..17285db 100644 --- a/src/arraybridge/libraries.py +++ b/src/arraybridge/libraries.py @@ -43,7 +43,6 @@ import os import sys import threading -import time from abc import ABCMeta, abstractmethod from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import AbstractContextManager, contextmanager, nullcontext @@ -56,7 +55,7 @@ from arraybridge.array_payload import ArrayPayload from arraybridge.arrays import Array -from arraybridge.conversion import ConversionObserver, ConversionRecord, ConversionRoute, HostCopy +from arraybridge.conversion import ConversionRoute, StaysInPlace from arraybridge.devices import HOST, AcceleratorDevice, Device, HostDevice, NamedDevice if TYPE_CHECKING: @@ -465,7 +464,7 @@ def move_to_device(self, data: NativeT, device: Device) -> NativeT: """Return ``data`` on ``device``; the same object when it is already there.""" self.require_device(device) - if self.device_of(data) == device: + if StaysInPlace.applies(self, self.device_of(data), self, device): return data with self._scope(device): return self._move(data, device) @@ -600,50 +599,27 @@ def unrecorded_from_numpy(self, data: np.ndarray, device: Device) -> NativeT: with self.device_scope(device): return self._place(data, device) + def shares_numpy_buffer_on(self, device: Device) -> bool: + """Whether an array of this library on ``device`` and a NumPy array can be one buffer.""" + + del device + return False + + def host_address_of(self, data: NativeT) -> int | None: + """Where ``data`` begins in main memory; ``None`` when no bytes of it are there.""" + + del data + return None + def to_numpy(self, data: NativeT) -> np.ndarray: - """Download ``data`` to a NumPy array in main memory.""" - - observer = ConversionObserver.current - if not observer.wants_records: - return self._to_host(data) - started = time.perf_counter() - source_device = self.device_of(data) - result = self._to_host(data) - if result is not data: - observer.record( - ConversionRecord( - source_library=str(self.name), - source_device=source_device.label, - target_library=str(NumpyLibrary.name), - target_device=HOST.label, - mechanism=str(HostCopy.mechanism), - nbytes=int(result.nbytes), - seconds=time.perf_counter() - started, - ) - ) - return result + """``data`` as a NumPy array in main memory: the conversion to NumPy on the host.""" + + return ConversionRoute.run(data, self, NumpyLibrary.registered(), HOST) def from_numpy(self, data: np.ndarray, device: Device) -> NativeT: - """Create an array of this library on ``device`` from a NumPy array.""" - - observer = ConversionObserver.current - if not observer.wants_records: - return self.unrecorded_from_numpy(data, device) - started = time.perf_counter() - result = self.unrecorded_from_numpy(data, device) - if result is not data: - observer.record( - ConversionRecord( - source_library=str(NumpyLibrary.name), - source_device=HOST.label, - target_library=str(self.name), - target_device=device.label, - mechanism=str(HostCopy.mechanism), - nbytes=int(data.nbytes), - seconds=time.perf_counter() - started, - ) - ) - return result + """An array of this library on ``device``: the conversion from NumPy on the host.""" + + return ConversionRoute.run(data, NumpyLibrary.registered(), self, device) def convert_to(self, data: NativeT, target: ArrayLibrary[TargetT], device: Device) -> TargetT: """Convert ``data`` to ``target`` on ``device`` by the first route that carries it.""" @@ -847,6 +823,27 @@ def _move_to_host(self, data: NativeT) -> NativeT: """The same native array in main memory.""" +class NumpyBufferSharing(HostDevices[NativeT]): + """Libraries whose main-memory arrays can be the very buffer of a NumPy array. + + Declared after measuring it: the library's exchange with NumPy on the host + returns an array that begins at the address of the one it was given, for + some arrays at least. ``host_address_of`` is what decides it for one array. + """ + + def shares_numpy_buffer_on(self, device: Device) -> bool: + return isinstance(device, HostDevice) + + def host_address_of(self, data: NativeT) -> int | None: + if not isinstance(self.device_of(data), HostDevice) or self.nbytes_of(data) == 0: + return None + return self._first_element_address(data) + + @abstractmethod + def _first_element_address(self, data: NativeT) -> int: + """Address of the first element of an array held in main memory.""" + + class AcceleratorDevices(ArrayLibrary[NativeT]): """Libraries that can keep arrays on an accelerator (every one but NumPy).""" @@ -1243,7 +1240,7 @@ def _rescales_into(self, target_dtype: DType) -> bool: # --------------------------------------------------------------------------- -class NumpyLibrary(HostDevices["np.ndarray"], MutableArrays["np.ndarray"]): +class NumpyLibrary(NumpyBufferSharing["np.ndarray"], MutableArrays["np.ndarray"]): name = "numpy" display_name = "NumPy" oom_messages = ("cannot allocate memory", "memory exhausted") @@ -1255,10 +1252,22 @@ class NumpyLibrary(HostDevices["np.ndarray"], MutableArrays["np.ndarray"]): def array_module(self) -> MutableNumpyFunctions[np.ndarray]: return self.module + @classmethod + def registered(cls) -> NumpyLibrary: + """The NumPy library every conversion to or from NumPy names as its other end.""" + + library = ArrayLibrary.named(cls.name) + if not isinstance(library, NumpyLibrary): + raise TypeError(f"{library!r} is registered as '{cls.name}' and is not NumPy") + return library + def device_of(self, data: np.ndarray) -> Device: del data return HOST + def _first_element_address(self, data: np.ndarray) -> int: + return data.ctypes.data + def accepts(self, value: object) -> bool: return isinstance(value, (self.module.ndarray, ArrayPayload)) @@ -1354,7 +1363,7 @@ def _rescales_into(self, target_dtype: DType) -> bool: class TorchLibrary( - HostDevices["torch.Tensor"], + NumpyBufferSharing["torch.Tensor"], CudaLibrary["torch.Tensor"], DLPackImport["torch.Tensor"], DLPackProtocolExport["torch.Tensor"], @@ -1403,6 +1412,9 @@ def _wrap_host_array(self, data: np.ndarray) -> torch.Tensor: def _move_to_host(self, data: torch.Tensor) -> torch.Tensor: return data.cpu() + def _first_element_address(self, data: torch.Tensor) -> int: + return data.data_ptr() + def current_device(self) -> AcceleratorDevice | None: return ( AcceleratorDevice(int(self.module.cuda.current_device())) @@ -1683,7 +1695,7 @@ def _ones_like( class JaxLibrary( - HostDevices["jax.Array"], + NumpyBufferSharing["jax.Array"], AcceleratorDevices["jax.Array"], DLPackImport["jax.Array"], DLPackProtocolExport["jax.Array"], @@ -1743,6 +1755,9 @@ def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> jax.Array: def _wrap_host_array(self, data: np.ndarray) -> jax.Array: return self.module.device_put(data, self.module.devices("cpu")[0]) + def _first_element_address(self, data: jax.Array) -> int: + return data.unsafe_buffer_pointer() + def _reads_dlpack(self, payload: DLPackPayload) -> bool: return _has_dlpack_protocol(payload.source) diff --git a/stubs/jax/__init__.pyi b/stubs/jax/__init__.pyi index 48adf99..941bda4 100644 --- a/stubs/jax/__init__.pyi +++ b/stubs/jax/__init__.pyi @@ -25,6 +25,7 @@ class Array: # A property in current JAX, a method in older releases. device: Device | Callable[[], Device] def astype(self, dtype: object, *, copy: bool = ...) -> Self: ... + def unsafe_buffer_pointer(self) -> int: ... def reshape(self, shape: tuple[int, ...]) -> Self: ... def __sub__(self, other: Self | float) -> Self: ... def __mul__(self, other: float) -> Self: ... diff --git a/stubs/numpy/__init__.pyi b/stubs/numpy/__init__.pyi index c8db68b..5c62abd 100644 --- a/stubs/numpy/__init__.pyi +++ b/stubs/numpy/__init__.pyi @@ -27,8 +27,12 @@ class int64(integer): ... class float32(floating): ... class float64(floating): ... +class _CtypesView: + data: int + class ndarray: shape: tuple[int, ...] + ctypes: _CtypesView dtype: dtype nbytes: int strides: tuple[int, ...] diff --git a/stubs/torch/__init__.pyi b/stubs/torch/__init__.pyi index 53d02af..612f80f 100644 --- a/stubs/torch/__init__.pyi +++ b/stubs/torch/__init__.pyi @@ -31,6 +31,7 @@ class Tensor: def float(self) -> Self: ... def numpy(self) -> numpy.ndarray: ... def element_size(self) -> int: ... + def data_ptr(self) -> int: ... def nelement(self) -> int: ... def min(self) -> Self: ... def max(self) -> Self: ... diff --git a/tests/fakes.py b/tests/fakes.py index 553ea33..25ae295 100644 --- a/tests/fakes.py +++ b/tests/fakes.py @@ -152,7 +152,11 @@ def to(self, device): return type(self)(self.values, int(device.removeprefix("cuda:"))) def numpy(self): - return self.values.copy() + # As torch does in main memory: the NumPy array is the tensor's buffer. + return self.values + + def data_ptr(self): + return self.values.ctypes.data def element_size(self): return self.values.dtype.itemsize diff --git a/tests/test_conversion.py b/tests/test_conversion.py index 1187e4d..6c23a5e 100644 --- a/tests/test_conversion.py +++ b/tests/test_conversion.py @@ -1,8 +1,11 @@ """Conversion routes: which mechanism is planned, which one ran, and what was recorded.""" +import asyncio +import contextvars import json import pickle import sys +import threading from types import SimpleNamespace import numpy as np @@ -17,6 +20,8 @@ DeviceMove, DLPackExchange, HostCopy, + SharedHostBuffer, + StaysInPlace, ) from arraybridge.devices import HOST, AcceleratorDevice from arraybridge.libraries import ( @@ -54,11 +59,24 @@ def planned(source, source_device, target, target_device): class TestPlanning: def test_routes_are_tried_in_declaration_order(self): assert list(ConversionRoute.__registry__.items()) == [ + ("stays in place", StaysInPlace), ("device move", DeviceMove), ("dlpack", DLPackExchange), + ("shared host buffer", SharedHostBuffer), ("host copy", HostCopy), ] + @pytest.mark.parametrize(("library", "device"), placements(), ids=str) + def test_same_library_and_device_stays_in_place(self, library, device): + assert ConversionRoute.planned(library, device, library, device) is StaysInPlace + assert not DeviceMove.applies(library, device, library, device) + assert StaysInPlace.writes_a_record is False + assert all( + route.writes_a_record + for route in ConversionRoute.__registry__.values() + if route is not StaysInPlace + ) + @pytest.mark.parametrize(("source", "source_device"), placements(), ids=str) def test_every_pair_has_a_route_ending_in_a_host_copy(self, source, source_device): for target, target_device in placements(): @@ -68,11 +86,14 @@ def test_every_pair_has_a_route_ending_in_a_host_copy(self, source, source_devic assert candidates[0].mechanism == mechanism assert candidates[-1] is HostCopy if source is target: - assert mechanism == "device move" - elif source is NUMPY or target in (NUMPY, PYCLESPERANTO): - assert mechanism == "host copy" - else: + same_device = source_device == target_device + assert mechanism == ("stays in place" if same_device else "device move") + elif PYCLESPERANTO in (source, target) or GPU in (source_device, target_device): assert mechanism in ("dlpack", "host copy") + elif NUMPY in (source, target): + assert mechanism in ("shared host buffer", "host copy") + else: + assert mechanism in ("dlpack", "shared host buffer", "host copy") @pytest.mark.parametrize( ("source", "source_device", "target", "target_device", "mechanism"), @@ -92,6 +113,20 @@ def test_every_pair_has_a_route_ending_in_a_host_copy(self, source, source_devic (PYCLESPERANTO, GPU, TORCH, GPU, "host copy"), (TORCH, GPU, PYCLESPERANTO, GPU, "host copy"), (TORCH, HOST, TORCH, GPU, "device move"), + (TORCH, GPU, TORCH, HOST, "device move"), + (TORCH, GPU, TORCH, GPU, "stays in place"), + (NUMPY, HOST, NUMPY, HOST, "stays in place"), + (NUMPY, HOST, TORCH, HOST, "shared host buffer"), + (TORCH, HOST, NUMPY, HOST, "shared host buffer"), + (NUMPY, HOST, JAX, HOST, "shared host buffer"), + (JAX, HOST, NUMPY, HOST, "shared host buffer"), + (NUMPY, HOST, TENSORFLOW, HOST, "host copy"), + (TENSORFLOW, HOST, NUMPY, HOST, "host copy"), + (TENSORFLOW, HOST, JAX, HOST, "host copy"), + (NUMPY, HOST, TORCH, GPU, "host copy"), + (TORCH, GPU, NUMPY, HOST, "host copy"), + (NUMPY, HOST, CUPY, GPU, "host copy"), + (NUMPY, HOST, PYCLESPERANTO, GPU, "host copy"), ], ids=str, ) @@ -323,7 +358,7 @@ def no_record(**fields): pytest.fail(f"a record was built: {fields}") monkeypatch.setattr(conversion, "ConversionRecord", no_record) - monkeypatch.setattr(libraries, "ConversionRecord", no_record) + assert not hasattr(libraries, "ConversionRecord") assert type(ConversionObserver.current) is conversion.IgnoreConversions assert ConversionObserver.current.wants_records is False @@ -367,7 +402,7 @@ def test_each_conversion_writes_one_record( source_device="host", target_library=target, target_device="host", - mechanism="host copy", + mechanism="shared host buffer", nbytes=self.values.nbytes, seconds=record.seconds, ) @@ -387,8 +422,31 @@ def test_returning_the_input_writes_no_record(self, torch_module, recorded): assert NUMPY.from_numpy(self.values, HOST) is self.values assert NUMPY.convert_to(self.values, NUMPY, HOST) is self.values assert TORCH.convert_to(tensor, TORCH, HOST) is tensor + assert TORCH.move_to_device(tensor, HOST) is tensor assert recorded == [] + assert planned(NUMPY, HOST, NUMPY, HOST) == planned(TORCH, HOST, TORCH, HOST) + assert planned(TORCH, HOST, TORCH, HOST) == "stays in place" + + def test_a_copy_where_sharing_was_planned_is_recorded_as_a_copy( + self, torch_module, recorded, monkeypatch + ): + monkeypatch.setattr( + torch_module, "from_numpy", lambda data: FakeTorchTensor(np.array(data)) + ) + assert planned(NUMPY, HOST, TORCH, HOST) == "shared host buffer" + + tensor = TORCH.from_numpy(self.values, HOST) + + assert not np.shares_memory(tensor.values, self.values) + assert [(each.mechanism, each.declined_mechanisms) for each in recorded] == [ + ("host copy", ("shared host buffer",)) + ] + + def test_an_empty_array_shares_nothing(self, torch_module, recorded): + TORCH.from_numpy(np.zeros((0, 3), dtype=np.uint8), HOST) + + assert [(each.mechanism, each.nbytes) for each in recorded] == [("host copy", 0)] def test_device_move_is_recorded(self, torch_module, recorded): moved = TORCH.convert_to(FakeTorchTensor(self.values), TORCH, GPU) @@ -408,6 +466,274 @@ def test_record_is_plain_data(self): assert hash(record) == hash(pickle.loads(pickle.dumps(record))) +def address(array): + return array.ctypes.data + + +def buffer_starting_at(alignment, offset, count=256): + """A writable float32 array whose first byte is ``offset`` past a multiple of ``alignment``.""" + + raw = np.zeros(count * 4 + 2 * alignment, dtype=np.uint8) + start = (-address(raw)) % alignment + offset + array = raw[start : start + count * 4].view(np.float32) + array[:] = np.arange(count, dtype=np.float32) + assert address(array) % alignment == offset + return array + + +class TestHostBufferIsSharedOrCopied: + """What the record says is what the memory shows, measured on the installed library.""" + + plain = staticmethod(lambda: np.arange(24, dtype=np.float32).reshape(4, 6)) + + @staticmethod + def one(recorded): + (record,) = recorded + recorded.clear() + return record.mechanism, record.declined_mechanisms + + @pytest.mark.torch + @pytest.mark.parametrize( + "make", + [ + lambda: np.arange(24, dtype=np.float32).reshape(4, 6), + lambda: np.arange(24, dtype=np.uint16).reshape(4, 6), + lambda: np.arange(24, dtype=np.int64).reshape(4, 6) > 3, + lambda: np.asfortranarray(np.arange(24, dtype=np.float64).reshape(4, 6)), + lambda: np.arange(24, dtype=np.float32).reshape(4, 6)[1:, ::2], + ], + ids=["float32", "uint16", "bool", "fortran order", "offset stepped view"], + ) + def test_torch_and_numpy_share_the_buffer_both_ways(self, make, recorded): + pytest.importorskip("torch") + values = make() + + tensor = TORCH.from_numpy(values, HOST) + assert tensor.data_ptr() == address(values) + assert self.one(recorded) == ("shared host buffer", ()) + + back = TORCH.to_numpy(tensor) + assert address(back) == tensor.data_ptr() + assert np.shares_memory(back, values) + assert self.one(recorded) == ("shared host buffer", ()) + + values[0, 0] = not values[0, 0] + assert back[0, 0] == values[0, 0] == tensor[0, 0].item() + + @pytest.mark.torch + def test_torch_copies_a_reversed_view_and_the_record_says_copy(self, recorded): + pytest.importorskip("torch") + values = self.plain()[::-1] + + tensor = TORCH.from_numpy(values, HOST) + + assert not np.shares_memory(tensor.numpy(), values) + assert self.one(recorded) == ("host copy", ("shared host buffer",)) + np.testing.assert_array_equal(tensor.numpy(), values) + + @pytest.mark.torch + def test_torch_shares_a_read_only_buffer(self, recorded): + pytest.importorskip("torch") + values = self.plain() + values.setflags(write=False) + + with pytest.warns(UserWarning, match="not writable"): + tensor = TORCH.from_numpy(values, HOST) + + assert tensor.data_ptr() == address(values) + assert self.one(recorded) == ("shared host buffer", ()) + + @pytest.mark.jax + def test_jax_hands_numpy_its_own_buffer(self, recorded): + pytest.importorskip("jax") + array = JAX.from_numpy(self.plain(), HOST) + recorded.clear() + + first, second = JAX.to_numpy(array), JAX.to_numpy(array) + + assert address(first) == address(second) == array.unsafe_buffer_pointer() + assert np.shares_memory(first, second) + assert not first.flags.writeable + assert [each.mechanism for each in recorded] == ["shared host buffer"] * 2 + + @pytest.mark.jax + @pytest.mark.parametrize( + ("offset", "expected"), + [(0, ("shared host buffer", ())), (4, ("host copy", ("shared host buffer",)))], + ids=["64-byte aligned", "not aligned"], + ) + def test_jax_wraps_a_numpy_buffer_only_when_it_is_aligned(self, offset, expected, recorded): + pytest.importorskip("jax") + values = buffer_starting_at(64, offset) + + array = JAX.from_numpy(values, HOST) + + shares = array.unsafe_buffer_pointer() == address(values) + assert shares == np.shares_memory(np.asarray(array), values) + assert shares == (expected[0] == "shared host buffer") + assert self.one(recorded) == expected + + @pytest.mark.tensorflow + @pytest.mark.parametrize("offset", [0, 4], ids=["64-byte aligned", "not aligned"]) + def test_tensorflow_copies_both_ways(self, offset, recorded): + pytest.importorskip("tensorflow") + values = buffer_starting_at(64, offset) + assert planned(NUMPY, HOST, TENSORFLOW, HOST) == "host copy" + assert planned(TENSORFLOW, HOST, NUMPY, HOST) == "host copy" + + tensor = TENSORFLOW.from_numpy(values, HOST) + assert self.one(recorded) == ("host copy", ()) + first, second = TENSORFLOW.to_numpy(tensor), TENSORFLOW.to_numpy(tensor) + + assert not np.shares_memory(first, values) + assert not np.shares_memory(first, second) + values[0] += 1 + assert first[0] == second[0] == values[0] - 1 + assert [each.mechanism for each in recorded] == ["host copy"] * 2 + + @pytest.mark.torch + @pytest.mark.tensorflow + def test_tensorflow_to_torch_on_the_host_is_a_copy(self, recorded): + pytest.importorskip("torch") + pytest.importorskip("tensorflow") + tensor = TENSORFLOW.from_numpy(self.plain(), HOST) + recorded.clear() + + result = TENSORFLOW.convert_to(tensor, TORCH, HOST) + + assert not np.shares_memory(result.numpy(), TENSORFLOW.unrecorded_to_numpy(tensor)) + assert self.one(recorded) == ("host copy", ()) + + +def counted(counters): + return [(row["mechanism"], row["nbytes"]) for row in counters.rows()] + + +class TestObserverBelongsToItsContext: + """Each thread and each asyncio task has its own ``ConversionObserver.current``.""" + + def test_threads_that_leave_in_entry_order_keep_their_own_observers(self, torch_module): + default = ConversionObserver.current + first_entered, second_entered, first_left = (threading.Event() for _ in range(3)) + observers = {name: ConversionCounters() for name in "AB"} + current_inside, current_after, failures = {}, {}, [] + + def work(name, nbytes): + try: + # Forced order: A enters, B enters, A leaves, B leaves. + assert name == "A" or first_entered.wait(10) + with observers[name].installed(): + (first_entered if name == "A" else second_entered).set() + assert (second_entered if name == "A" else first_left).wait(10) + current_inside[name] = ConversionObserver.current + TORCH.from_numpy(np.zeros(nbytes, dtype=np.uint8), HOST) + current_after[name] = ConversionObserver.current + if name == "A": + first_left.set() + except BaseException as error: # reported by the main thread + failures.append(error) + + threads = [threading.Thread(target=work, args=pair) for pair in (("A", 8), ("B", 16))] + for thread in threads: + thread.start() + for thread in threads: + thread.join(20) + + assert failures == [] + assert current_inside == observers + assert counted(observers["A"]) == [("shared host buffer", 8)] + assert counted(observers["B"]) == [("shared host buffer", 16)] + assert type(current_after["A"]) is type(current_after["B"]) is type(default) + assert current_after["A"] is current_after["B"] is default is ConversionObserver.current + + def test_a_new_thread_starts_with_the_default_observer(self, torch_module): + default = ConversionObserver.current + seen = [] + + def work(): + seen.append(ConversionObserver.current) + TORCH.from_numpy(np.zeros(8, dtype=np.uint8), HOST) + + with ConversionCounters().installed() as counters: + thread = threading.Thread(target=work) + thread.start() + thread.join(20) + + assert seen == [default] + assert counters.rows() == [] + + def test_a_thread_run_in_a_copied_context_reports_to_the_spawner(self, torch_module): + seen = [] + + def work(): + seen.append(ConversionObserver.current) + TORCH.from_numpy(np.zeros(8, dtype=np.uint8), HOST) + + with ConversionCounters().installed() as counters: + thread = threading.Thread(target=contextvars.copy_context().run, args=(work,)) + thread.start() + thread.join(20) + + assert seen == [counters] + assert counted(counters) == [("shared host buffer", 8)] + + def test_assigning_current_changes_this_context_only(self): + default = ConversionObserver.current + counters = ConversionCounters() + seen = [] + + def assign(): + ConversionObserver.current = counters + seen.append(ConversionObserver.current) + + contextvars.copy_context().run(assign) + thread = threading.Thread(target=assign) + thread.start() + thread.join(20) + + assert seen == [counters, counters] + assert ConversionObserver.current is default + assert isinstance(vars(type(ConversionObserver))["current"], property) + + def test_asyncio_tasks_keep_their_own_observers(self, torch_module): + default = ConversionObserver.current + + async def convert(nbytes, started, may_leave): + with ConversionCounters().installed() as counters: + started.set() + await may_leave.wait() + TORCH.from_numpy(np.zeros(nbytes, dtype=np.uint8), HOST) + assert ConversionObserver.current is counters + return counters, ConversionObserver.current + + async def inherited(): + TORCH.from_numpy(np.zeros(32, dtype=np.uint8), HOST) + return ConversionObserver.current + + async def main(): + a_started, b_started, a_leaves, b_leaves = (asyncio.Event() for _ in range(4)) + a = asyncio.create_task(convert(8, a_started, a_leaves)) + await a_started.wait() + b = asyncio.create_task(convert(16, b_started, b_leaves)) + await b_started.wait() + a_leaves.set() + a_result = await a + b_leaves.set() + b_result = await b + with ConversionCounters().installed() as spawner: + child_saw = await asyncio.create_task(inherited()) + return a_result, b_result, spawner, child_saw, ConversionObserver.current + + (a, a_after), (b, b_after), spawner, child_saw, after = asyncio.run(main()) + + assert counted(a) == [("shared host buffer", 8)] + assert counted(b) == [("shared host buffer", 16)] + assert a_after is b_after is after is default + assert child_saw is spawner + assert counted(spawner) == [("shared host buffer", 32)] + assert ConversionObserver.current is default + + class TestCounters: @staticmethod def record(source="numpy", target="torch", nbytes=10, seconds=0.25, mechanism="host copy"): @@ -455,7 +781,7 @@ def test_installed_counters_count_conversions(self, torch_module): TORCH.from_numpy(values, HOST) (row,) = counters.rows() - assert (row["count"], row["nbytes"], row["mechanism"]) == (3, 24, "host copy") + assert (row["count"], row["nbytes"], row["mechanism"]) == (3, 24, "shared host buffer") def test_merge_sums_counts_bytes_and_seconds(self): parent = self.counters(self.record(), self.record("cupy", nbytes=5)) diff --git a/tests/test_toy_library.py b/tests/test_toy_library.py index fd19e58..69411f2 100644 --- a/tests/test_toy_library.py +++ b/tests/test_toy_library.py @@ -146,7 +146,7 @@ def test_defining_the_class_is_the_whole_registration(toy_library_class): numpy_library = ArrayLibrary.named("numpy") assert ConversionRoute.planned(numpy_library, HOST, library, HOST).mechanism == "host copy" assert ConversionRoute.planned(library, HOST, numpy_library, HOST).mechanism == "host copy" - assert ConversionRoute.planned(library, HOST, library, HOST).mechanism == "device move" + assert ConversionRoute.planned(library, HOST, library, HOST).mechanism == "stays in place" @toy def halve(image):