zgc-runtime 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- zgc_runtime-0.1.0/PKG-INFO +5 -0
- zgc_runtime-0.1.0/pyproject.toml +15 -0
- zgc_runtime-0.1.0/setup.cfg +4 -0
- zgc_runtime-0.1.0/src/zgc/__init__.py +14 -0
- zgc_runtime-0.1.0/src/zgc/_abi.py +172 -0
- zgc_runtime-0.1.0/src/zgc/py.typed +1 -0
- zgc_runtime-0.1.0/src/zgc/runtime.py +472 -0
- zgc_runtime-0.1.0/src/zgc_runtime.egg-info/PKG-INFO +5 -0
- zgc_runtime-0.1.0/src/zgc_runtime.egg-info/SOURCES.txt +10 -0
- zgc_runtime-0.1.0/src/zgc_runtime.egg-info/dependency_links.txt +1 -0
- zgc_runtime-0.1.0/src/zgc_runtime.egg-info/top_level.txt +1 -0
- zgc_runtime-0.1.0/tests/test_buffers.py +113 -0
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "zgc-runtime"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Python runtime for ZGC executable artifacts"
|
|
9
|
+
requires-python = ">=3.10"
|
|
10
|
+
|
|
11
|
+
[tool.setuptools.packages.find]
|
|
12
|
+
where = ["src"]
|
|
13
|
+
|
|
14
|
+
[tool.setuptools.package-data]
|
|
15
|
+
zgc = ["py.typed"]
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
from ._abi import ABI_VERSION, SourceBinding, SourceKind, dtype
|
|
2
|
+
from .runtime import Runtime, Source, Tensor, TensorInfo, load
|
|
3
|
+
|
|
4
|
+
__all__ = [
|
|
5
|
+
"ABI_VERSION",
|
|
6
|
+
"dtype",
|
|
7
|
+
"Runtime",
|
|
8
|
+
"Source",
|
|
9
|
+
"SourceBinding",
|
|
10
|
+
"SourceKind",
|
|
11
|
+
"Tensor",
|
|
12
|
+
"TensorInfo",
|
|
13
|
+
"load",
|
|
14
|
+
]
|
|
@@ -0,0 +1,172 @@
|
|
|
1
|
+
"""ZGC ABI bindings and helpers."""
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
|
|
4
|
+
import ctypes
|
|
5
|
+
from enum import IntEnum
|
|
6
|
+
from typing import Protocol, cast, final
|
|
7
|
+
|
|
8
|
+
ABI_VERSION = 1
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class Status(IntEnum):
|
|
12
|
+
OK = 0
|
|
13
|
+
INVALID_ARGUMENT = 1
|
|
14
|
+
INVALID_RUNTIME_STORAGE = 2
|
|
15
|
+
INVALID_SOURCE = 3
|
|
16
|
+
INVALID_OUTPUT = 4
|
|
17
|
+
SIZE_MISMATCH = 5
|
|
18
|
+
ALIGNMENT_MISMATCH = 6
|
|
19
|
+
INCOMPATIBLE_BINDING = 7
|
|
20
|
+
MISSING_BINDING = 8
|
|
21
|
+
BUFFER_TOO_SMALL = 9
|
|
22
|
+
OUT_OF_MEMORY = 10
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class dtype(IntEnum):
|
|
26
|
+
f32 = 1
|
|
27
|
+
f16 = 2
|
|
28
|
+
i8 = 3
|
|
29
|
+
bool = 4
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class SourceKind(IntEnum):
|
|
33
|
+
INPUT = 1
|
|
34
|
+
PARAMETER = 2
|
|
35
|
+
CONSTANT = 3
|
|
36
|
+
STATE = 4
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class SourceBinding(IntEnum):
|
|
40
|
+
OWNED = 1
|
|
41
|
+
BOUND = 2
|
|
42
|
+
EMBEDDED = 3
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
@final
|
|
46
|
+
class TensorDescriptor(ctypes.Structure):
|
|
47
|
+
_fields_ = [
|
|
48
|
+
("dtype", ctypes.c_uint32),
|
|
49
|
+
("rank", ctypes.c_uint32),
|
|
50
|
+
("element_count", ctypes.c_size_t),
|
|
51
|
+
("logical_byte_count", ctypes.c_size_t),
|
|
52
|
+
("offset_elements", ctypes.c_size_t),
|
|
53
|
+
("shape", ctypes.POINTER(ctypes.c_size_t)),
|
|
54
|
+
("strides", ctypes.POINTER(ctypes.c_ssize_t)),
|
|
55
|
+
]
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
@final
|
|
59
|
+
class SourceDescriptor(ctypes.Structure):
|
|
60
|
+
_fields_ = [
|
|
61
|
+
("key", ctypes.c_uint32),
|
|
62
|
+
("name", ctypes.c_char_p),
|
|
63
|
+
("kind", ctypes.c_uint32),
|
|
64
|
+
("binding", ctypes.c_uint32),
|
|
65
|
+
("tensor", TensorDescriptor),
|
|
66
|
+
]
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def configure(library: ctypes.CDLL) -> None:
|
|
70
|
+
library.zgc_abi_version.restype = ctypes.c_uint32
|
|
71
|
+
library.zgc_executable_structural_fingerprint.restype = ctypes.c_uint64
|
|
72
|
+
library.zgc_runtime_size.restype = ctypes.c_size_t
|
|
73
|
+
library.zgc_runtime_alignment.restype = ctypes.c_size_t
|
|
74
|
+
library.zgc_runtime_mutable_bytes.restype = ctypes.c_size_t
|
|
75
|
+
|
|
76
|
+
library.zgc_runtime_init.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
|
|
77
|
+
library.zgc_runtime_init.restype = ctypes.c_int
|
|
78
|
+
library.zgc_runtime_deinit.argtypes = [ctypes.c_void_p]
|
|
79
|
+
library.zgc_runtime_deinit.restype = None
|
|
80
|
+
library.zgc_runtime_run.argtypes = [ctypes.c_void_p]
|
|
81
|
+
library.zgc_runtime_run.restype = ctypes.c_int
|
|
82
|
+
|
|
83
|
+
library.zgc_executable_source_count.restype = ctypes.c_size_t
|
|
84
|
+
library.zgc_executable_get_source_descriptor.argtypes = [
|
|
85
|
+
ctypes.c_size_t,
|
|
86
|
+
ctypes.POINTER(SourceDescriptor),
|
|
87
|
+
]
|
|
88
|
+
library.zgc_executable_get_source_descriptor.restype = ctypes.c_int
|
|
89
|
+
library.zgc_runtime_copy_source.argtypes = [
|
|
90
|
+
ctypes.c_void_p,
|
|
91
|
+
ctypes.c_uint32,
|
|
92
|
+
ctypes.c_void_p,
|
|
93
|
+
ctypes.c_size_t,
|
|
94
|
+
]
|
|
95
|
+
library.zgc_runtime_copy_source.restype = ctypes.c_int
|
|
96
|
+
library.zgc_runtime_bind_source.argtypes = [
|
|
97
|
+
ctypes.c_void_p,
|
|
98
|
+
ctypes.c_uint32,
|
|
99
|
+
ctypes.c_void_p,
|
|
100
|
+
ctypes.c_size_t,
|
|
101
|
+
]
|
|
102
|
+
library.zgc_runtime_bind_source.restype = ctypes.c_int
|
|
103
|
+
|
|
104
|
+
library.zgc_executable_output_count.restype = ctypes.c_size_t
|
|
105
|
+
library.zgc_executable_get_output_descriptor.argtypes = [
|
|
106
|
+
ctypes.c_size_t,
|
|
107
|
+
ctypes.POINTER(TensorDescriptor),
|
|
108
|
+
]
|
|
109
|
+
library.zgc_executable_get_output_descriptor.restype = ctypes.c_int
|
|
110
|
+
library.zgc_runtime_copy_output.argtypes = [
|
|
111
|
+
ctypes.c_void_p,
|
|
112
|
+
ctypes.c_size_t,
|
|
113
|
+
ctypes.c_void_p,
|
|
114
|
+
ctypes.c_size_t,
|
|
115
|
+
]
|
|
116
|
+
library.zgc_runtime_copy_output.restype = ctypes.c_int
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
class Library(Protocol):
|
|
120
|
+
def zgc_abi_version(self) -> int: ...
|
|
121
|
+
def zgc_executable_structural_fingerprint(self) -> int: ...
|
|
122
|
+
def zgc_runtime_size(self) -> int: ...
|
|
123
|
+
def zgc_runtime_alignment(self) -> int: ...
|
|
124
|
+
def zgc_runtime_mutable_bytes(self) -> int: ...
|
|
125
|
+
def zgc_runtime_init(self, runtime: ctypes.c_void_p, size: int) -> int: ...
|
|
126
|
+
def zgc_runtime_deinit(self, runtime: ctypes.c_void_p) -> None: ...
|
|
127
|
+
def zgc_runtime_run(self, runtime: ctypes.c_void_p) -> int: ...
|
|
128
|
+
def zgc_executable_source_count(self) -> int: ...
|
|
129
|
+
def zgc_executable_get_source_descriptor(self, index: int, output: object) -> int: ...
|
|
130
|
+
def zgc_runtime_copy_source(
|
|
131
|
+
self,
|
|
132
|
+
runtime: ctypes.c_void_p,
|
|
133
|
+
key: int,
|
|
134
|
+
data: ctypes.c_void_p,
|
|
135
|
+
size: int,
|
|
136
|
+
) -> int: ...
|
|
137
|
+
def zgc_runtime_bind_source(
|
|
138
|
+
self,
|
|
139
|
+
runtime: ctypes.c_void_p,
|
|
140
|
+
key: int,
|
|
141
|
+
data: ctypes.c_void_p,
|
|
142
|
+
size: int,
|
|
143
|
+
) -> int: ...
|
|
144
|
+
def zgc_executable_output_count(self) -> int: ...
|
|
145
|
+
def zgc_executable_get_output_descriptor(self, index: int, output: object) -> int: ...
|
|
146
|
+
def zgc_runtime_copy_output(
|
|
147
|
+
self,
|
|
148
|
+
runtime: ctypes.c_void_p,
|
|
149
|
+
index: int,
|
|
150
|
+
data: ctypes.c_void_p,
|
|
151
|
+
size: int,
|
|
152
|
+
) -> int: ...
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
def load_library(path: str) -> Library:
|
|
156
|
+
library = ctypes.CDLL(path)
|
|
157
|
+
library.zgc_abi_version.restype = ctypes.c_uint32
|
|
158
|
+
version = library.zgc_abi_version()
|
|
159
|
+
if version != ABI_VERSION:
|
|
160
|
+
raise RuntimeError(f"unsupported ZGC ABI version {version}")
|
|
161
|
+
configure(library)
|
|
162
|
+
return cast(Library, cast(object, library))
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
def require_ok(status: int, operation: str) -> None:
|
|
166
|
+
if status == Status.OK:
|
|
167
|
+
return
|
|
168
|
+
try:
|
|
169
|
+
name = Status(status).name.lower()
|
|
170
|
+
except ValueError:
|
|
171
|
+
name = f"unknown status {status}"
|
|
172
|
+
raise RuntimeError(f"{operation} failed: {name}")
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1,472 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import ctypes
|
|
4
|
+
import struct
|
|
5
|
+
import sys
|
|
6
|
+
from collections.abc import Iterable, Sequence
|
|
7
|
+
from dataclasses import dataclass
|
|
8
|
+
from os import PathLike
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
from typing import Literal, Protocol, Self, cast, final
|
|
11
|
+
|
|
12
|
+
from ._abi import (
|
|
13
|
+
Library,
|
|
14
|
+
SourceBinding,
|
|
15
|
+
SourceDescriptor,
|
|
16
|
+
SourceKind,
|
|
17
|
+
TensorDescriptor,
|
|
18
|
+
load_library,
|
|
19
|
+
require_ok,
|
|
20
|
+
)
|
|
21
|
+
from ._abi import (
|
|
22
|
+
dtype as dtype_enum,
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class _TensorDescriptorFields(Protocol):
|
|
27
|
+
dtype: int
|
|
28
|
+
rank: int
|
|
29
|
+
element_count: int
|
|
30
|
+
logical_byte_count: int
|
|
31
|
+
offset_elements: int
|
|
32
|
+
shape: Sequence[int]
|
|
33
|
+
strides: Sequence[int]
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class _SourceDescriptorFields(Protocol):
|
|
37
|
+
key: int
|
|
38
|
+
name: bytes | None
|
|
39
|
+
kind: int
|
|
40
|
+
binding: int
|
|
41
|
+
tensor: TensorDescriptor
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
@dataclass(frozen=True)
|
|
45
|
+
class TensorInfo:
|
|
46
|
+
dtype: dtype_enum
|
|
47
|
+
shape: tuple[int, ...]
|
|
48
|
+
strides: tuple[int, ...]
|
|
49
|
+
element_count: int
|
|
50
|
+
byte_count: int
|
|
51
|
+
offset: int
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@dataclass(frozen=True)
|
|
55
|
+
class Source:
|
|
56
|
+
key: int
|
|
57
|
+
name: str
|
|
58
|
+
kind: SourceKind
|
|
59
|
+
binding: SourceBinding
|
|
60
|
+
tensor: TensorInfo
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
@dataclass(frozen=True)
|
|
64
|
+
class Tensor:
|
|
65
|
+
dtype: dtype_enum
|
|
66
|
+
shape: tuple[int, ...]
|
|
67
|
+
_storage: bytearray
|
|
68
|
+
|
|
69
|
+
@property
|
|
70
|
+
def buffer(self) -> memoryview:
|
|
71
|
+
raw = memoryview(self._storage)
|
|
72
|
+
return cast(
|
|
73
|
+
memoryview,
|
|
74
|
+
raw.cast( # pyright: ignore[reportCallIssue]
|
|
75
|
+
_format(self.dtype), # pyright: ignore[reportArgumentType]
|
|
76
|
+
shape=list(self.shape),
|
|
77
|
+
),
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
@property
|
|
81
|
+
def values(self) -> tuple[float | int | bool, ...]:
|
|
82
|
+
return _decode(self.dtype, len(self._storage), self._storage)
|
|
83
|
+
|
|
84
|
+
@property
|
|
85
|
+
def strides(self) -> tuple[int, ...]:
|
|
86
|
+
return tuple(self.buffer.strides or ())
|
|
87
|
+
|
|
88
|
+
def tolist(self) -> list[float | int | bool]:
|
|
89
|
+
return list(self.values)
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def _tensor_info(descriptor: TensorDescriptor) -> TensorInfo:
|
|
93
|
+
fields = cast(_TensorDescriptorFields, cast(object, descriptor))
|
|
94
|
+
return TensorInfo(
|
|
95
|
+
dtype=dtype_enum(fields.dtype),
|
|
96
|
+
shape=tuple(fields.shape[index] for index in range(fields.rank)),
|
|
97
|
+
strides=tuple(fields.strides[index] for index in range(fields.rank)),
|
|
98
|
+
element_count=fields.element_count,
|
|
99
|
+
byte_count=fields.logical_byte_count,
|
|
100
|
+
offset=fields.offset_elements,
|
|
101
|
+
)
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def _format(dtype: dtype_enum) -> Literal["f", "e", "b", "?"]:
|
|
105
|
+
if dtype is dtype_enum.f32:
|
|
106
|
+
return "f"
|
|
107
|
+
if dtype is dtype_enum.f16:
|
|
108
|
+
return "e"
|
|
109
|
+
if dtype is dtype_enum.i8:
|
|
110
|
+
return "b"
|
|
111
|
+
return "?"
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def _item_size(dtype: dtype_enum) -> int:
|
|
115
|
+
return struct.calcsize(f"={_format(dtype)}")
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def _decode(
|
|
119
|
+
dtype: dtype_enum,
|
|
120
|
+
byte_count: int,
|
|
121
|
+
payload: bytes | bytearray | memoryview,
|
|
122
|
+
) -> tuple[float | int | bool, ...]:
|
|
123
|
+
element_count = byte_count // _item_size(dtype)
|
|
124
|
+
return struct.unpack(f"={element_count}{_format(dtype)}", payload)
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
def _native_format(view: memoryview) -> str | None:
|
|
128
|
+
value = view.format
|
|
129
|
+
if not value:
|
|
130
|
+
return None
|
|
131
|
+
if value[0] in "@=":
|
|
132
|
+
return value[1:]
|
|
133
|
+
if value[0] in "<>!":
|
|
134
|
+
if view.itemsize == 1:
|
|
135
|
+
return value[1:]
|
|
136
|
+
native_prefix = "<" if sys.byteorder == "little" else ">"
|
|
137
|
+
return value[1:] if value[0] == native_prefix else None
|
|
138
|
+
return value
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
def _buffer_view(info: TensorInfo, values: object) -> memoryview | None:
|
|
142
|
+
try:
|
|
143
|
+
view = memoryview(values) # pyright: ignore[reportArgumentType]
|
|
144
|
+
except TypeError:
|
|
145
|
+
return None
|
|
146
|
+
|
|
147
|
+
if view.nbytes != info.byte_count:
|
|
148
|
+
raise ValueError(
|
|
149
|
+
f"expected {info.byte_count} source bytes, received {view.nbytes}"
|
|
150
|
+
)
|
|
151
|
+
|
|
152
|
+
# Byte-oriented buffers preserve the raw-byte input supported by the ABI.
|
|
153
|
+
if view.itemsize == 1 and _native_format(view) in {"b", "B", "c"}:
|
|
154
|
+
return view
|
|
155
|
+
|
|
156
|
+
expected_format = _format(info.dtype)
|
|
157
|
+
if _native_format(view) != expected_format or view.itemsize != _item_size(
|
|
158
|
+
info.dtype
|
|
159
|
+
):
|
|
160
|
+
raise TypeError(
|
|
161
|
+
f"expected a buffer with format {expected_format!r} for {info.dtype.name.lower()}"
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
shape = tuple(view.shape or ())
|
|
165
|
+
if shape != info.shape and shape != (info.element_count,):
|
|
166
|
+
raise ValueError(f"expected source shape {info.shape}, received {shape}")
|
|
167
|
+
return view
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def _logical_bytes(info: TensorInfo, values: object) -> bytes:
|
|
171
|
+
view = _buffer_view(info, values)
|
|
172
|
+
if view is not None:
|
|
173
|
+
return view.tobytes(order="C")
|
|
174
|
+
|
|
175
|
+
try:
|
|
176
|
+
materialized = tuple(cast(Iterable[object], values))
|
|
177
|
+
except TypeError as error:
|
|
178
|
+
raise TypeError(
|
|
179
|
+
"source must be a buffer or an iterable of scalar values"
|
|
180
|
+
) from error
|
|
181
|
+
if len(materialized) != info.element_count:
|
|
182
|
+
message = (
|
|
183
|
+
f"expected {info.element_count} elements for shape {info.shape}; "
|
|
184
|
+
f"received {len(materialized)}"
|
|
185
|
+
)
|
|
186
|
+
raise ValueError(message)
|
|
187
|
+
return struct.pack(f"={info.element_count}{_format(info.dtype)}", *materialized)
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
def _writable_pointer(
|
|
191
|
+
view: memoryview,
|
|
192
|
+
) -> tuple[ctypes.c_ubyte, ctypes.c_void_p] | None:
|
|
193
|
+
if view.readonly or not view.c_contiguous or view.nbytes == 0:
|
|
194
|
+
return None
|
|
195
|
+
anchor = ctypes.c_ubyte.from_buffer(view)
|
|
196
|
+
return anchor, ctypes.c_void_p(ctypes.addressof(anchor))
|
|
197
|
+
|
|
198
|
+
|
|
199
|
+
def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]:
|
|
200
|
+
strides = [0] * len(shape)
|
|
201
|
+
stride = 1
|
|
202
|
+
for axis in range(len(shape) - 1, -1, -1):
|
|
203
|
+
strides[axis] = stride
|
|
204
|
+
stride *= shape[axis]
|
|
205
|
+
return tuple(strides)
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
def _physical_buffer_matches(info: TensorInfo, view: memoryview) -> bool:
|
|
209
|
+
if info.offset != 0 or not view.c_contiguous:
|
|
210
|
+
return False
|
|
211
|
+
if view.itemsize == 1 and _native_format(view) in {"b", "B", "c"}:
|
|
212
|
+
return (
|
|
213
|
+
view.ndim == 1
|
|
214
|
+
and view.nbytes == info.byte_count
|
|
215
|
+
and info.strides == _contiguous_strides(info.shape)
|
|
216
|
+
)
|
|
217
|
+
expected_strides = tuple(stride * view.itemsize for stride in info.strides)
|
|
218
|
+
return (
|
|
219
|
+
tuple(view.shape or ()) == info.shape
|
|
220
|
+
and tuple(view.strides or ()) == expected_strides
|
|
221
|
+
)
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def _pack_bound(info: TensorInfo, logical: bytes) -> bytes:
|
|
225
|
+
logical_values = _decode(info.dtype, info.byte_count, logical)
|
|
226
|
+
physical_values: list[float | int | bool] = [0] * info.element_count
|
|
227
|
+
for linear_index, value in enumerate(logical_values):
|
|
228
|
+
remaining = linear_index
|
|
229
|
+
physical_index = info.offset
|
|
230
|
+
for axis in range(len(info.shape) - 1, -1, -1):
|
|
231
|
+
coordinate = remaining % info.shape[axis]
|
|
232
|
+
remaining //= info.shape[axis]
|
|
233
|
+
physical_index += coordinate * info.strides[axis]
|
|
234
|
+
if physical_index < 0 or physical_index >= info.element_count:
|
|
235
|
+
raise ValueError("bound source layout exceeds its ABI storage extent")
|
|
236
|
+
physical_values[physical_index] = value
|
|
237
|
+
return struct.pack(
|
|
238
|
+
f"={info.element_count}{_format(info.dtype)}",
|
|
239
|
+
*physical_values,
|
|
240
|
+
)
|
|
241
|
+
|
|
242
|
+
|
|
243
|
+
def _aligned_storage(
|
|
244
|
+
size: int,
|
|
245
|
+
alignment: int,
|
|
246
|
+
) -> tuple[ctypes.Array[ctypes.c_ubyte], ctypes.c_void_p]:
|
|
247
|
+
if alignment <= 0 or alignment & (alignment - 1):
|
|
248
|
+
raise RuntimeError(f"invalid runtime alignment: {alignment}")
|
|
249
|
+
allocation = (ctypes.c_ubyte * (size + alignment - 1))()
|
|
250
|
+
address = (ctypes.addressof(allocation) + alignment - 1) & -alignment
|
|
251
|
+
return allocation, ctypes.c_void_p(address)
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
@final
|
|
255
|
+
class Runtime:
|
|
256
|
+
def __init__(self, path: str | PathLike[str]) -> None:
|
|
257
|
+
self._library: Library = load_library(str(Path(path).expanduser().resolve()))
|
|
258
|
+
self.structural_fingerprint: int = (
|
|
259
|
+
self._library.zgc_executable_structural_fingerprint()
|
|
260
|
+
)
|
|
261
|
+
self.runtime_size: int = self._library.zgc_runtime_size()
|
|
262
|
+
self.runtime_alignment: int = self._library.zgc_runtime_alignment()
|
|
263
|
+
self.mutable_bytes: int = self._library.zgc_runtime_mutable_bytes()
|
|
264
|
+
self._allocation, self._runtime = _aligned_storage(
|
|
265
|
+
self.runtime_size, self.runtime_alignment
|
|
266
|
+
)
|
|
267
|
+
self._closed = False
|
|
268
|
+
self._bound_storage: dict[int, object] = {}
|
|
269
|
+
require_ok(
|
|
270
|
+
self._library.zgc_runtime_init(self._runtime, self.runtime_size),
|
|
271
|
+
"zgc_runtime_init",
|
|
272
|
+
)
|
|
273
|
+
self.sources = self._read_sources()
|
|
274
|
+
self.outputs = self._read_outputs()
|
|
275
|
+
self._sources_by_name = {source.name: source for source in self.sources}
|
|
276
|
+
|
|
277
|
+
def _read_sources(self) -> tuple[Source, ...]:
|
|
278
|
+
sources: list[Source] = []
|
|
279
|
+
for index in range(self._library.zgc_executable_source_count()):
|
|
280
|
+
descriptor = SourceDescriptor()
|
|
281
|
+
require_ok(
|
|
282
|
+
self._library.zgc_executable_get_source_descriptor(
|
|
283
|
+
index, ctypes.byref(descriptor)
|
|
284
|
+
),
|
|
285
|
+
"zgc_executable_get_source_descriptor",
|
|
286
|
+
)
|
|
287
|
+
fields = cast(_SourceDescriptorFields, cast(object, descriptor))
|
|
288
|
+
sources.append(
|
|
289
|
+
Source(
|
|
290
|
+
key=fields.key,
|
|
291
|
+
name=fields.name.decode() if fields.name is not None else "",
|
|
292
|
+
kind=SourceKind(fields.kind),
|
|
293
|
+
binding=SourceBinding(fields.binding),
|
|
294
|
+
tensor=_tensor_info(fields.tensor),
|
|
295
|
+
)
|
|
296
|
+
)
|
|
297
|
+
return tuple(sources)
|
|
298
|
+
|
|
299
|
+
def _read_outputs(self) -> tuple[TensorInfo, ...]:
|
|
300
|
+
outputs: list[TensorInfo] = []
|
|
301
|
+
for index in range(self._library.zgc_executable_output_count()):
|
|
302
|
+
descriptor = TensorDescriptor()
|
|
303
|
+
require_ok(
|
|
304
|
+
self._library.zgc_executable_get_output_descriptor(
|
|
305
|
+
index, ctypes.byref(descriptor)
|
|
306
|
+
),
|
|
307
|
+
"zgc_executable_get_output_descriptor",
|
|
308
|
+
)
|
|
309
|
+
outputs.append(_tensor_info(descriptor))
|
|
310
|
+
return tuple(outputs)
|
|
311
|
+
|
|
312
|
+
def set_source(
|
|
313
|
+
self,
|
|
314
|
+
name: str,
|
|
315
|
+
values: object,
|
|
316
|
+
) -> None:
|
|
317
|
+
"""Copy logical source values into storage owned by the runtime.
|
|
318
|
+
|
|
319
|
+
``values`` may implement the Python buffer protocol or be an iterable
|
|
320
|
+
of scalar values. For a compiled bound source, the values are packed
|
|
321
|
+
into retained internal storage. Later changes to ``values`` therefore
|
|
322
|
+
cannot affect the runtime. Use :meth:`bind` to explicitly share caller
|
|
323
|
+
storage without copying.
|
|
324
|
+
"""
|
|
325
|
+
self._require_open()
|
|
326
|
+
try:
|
|
327
|
+
source = self._sources_by_name[name]
|
|
328
|
+
except KeyError as error:
|
|
329
|
+
raise KeyError(f"unknown runtime source {name!r}") from error
|
|
330
|
+
if source.binding is SourceBinding.EMBEDDED:
|
|
331
|
+
raise ValueError(f"source {name!r} is embedded and cannot be replaced")
|
|
332
|
+
|
|
333
|
+
if source.binding is SourceBinding.BOUND:
|
|
334
|
+
payload = _pack_bound(
|
|
335
|
+
source.tensor,
|
|
336
|
+
_logical_bytes(source.tensor, values),
|
|
337
|
+
)
|
|
338
|
+
storage = bytearray(payload)
|
|
339
|
+
storage_view = memoryview(storage)
|
|
340
|
+
address = _writable_pointer(storage_view)
|
|
341
|
+
if address is None:
|
|
342
|
+
raise ValueError("cannot bind empty source storage")
|
|
343
|
+
anchor, pointer = address
|
|
344
|
+
require_ok(
|
|
345
|
+
self._library.zgc_runtime_bind_source(
|
|
346
|
+
self._runtime,
|
|
347
|
+
source.key,
|
|
348
|
+
pointer,
|
|
349
|
+
len(storage),
|
|
350
|
+
),
|
|
351
|
+
"zgc_runtime_bind_source",
|
|
352
|
+
)
|
|
353
|
+
self._bound_storage[source.key] = (storage, storage_view, anchor)
|
|
354
|
+
return
|
|
355
|
+
|
|
356
|
+
view = _buffer_view(source.tensor, values)
|
|
357
|
+
direct = _writable_pointer(view) if view is not None else None
|
|
358
|
+
if direct is not None and view is not None and view.c_contiguous:
|
|
359
|
+
anchor, pointer = direct
|
|
360
|
+
byte_count = view.nbytes
|
|
361
|
+
keepalive: object = (values, view, anchor)
|
|
362
|
+
else:
|
|
363
|
+
payload = _logical_bytes(source.tensor, values)
|
|
364
|
+
pointer = ctypes.cast(ctypes.c_char_p(payload), ctypes.c_void_p)
|
|
365
|
+
byte_count = len(payload)
|
|
366
|
+
keepalive = payload
|
|
367
|
+
require_ok(
|
|
368
|
+
self._library.zgc_runtime_copy_source(
|
|
369
|
+
self._runtime,
|
|
370
|
+
source.key,
|
|
371
|
+
pointer,
|
|
372
|
+
byte_count,
|
|
373
|
+
),
|
|
374
|
+
"zgc_runtime_copy_source",
|
|
375
|
+
)
|
|
376
|
+
_ = keepalive
|
|
377
|
+
|
|
378
|
+
def bind(self, name: str, values: object) -> None:
|
|
379
|
+
"""Bind caller-owned physical storage without copying.
|
|
380
|
+
|
|
381
|
+
``values`` must export a writable, contiguous buffer whose metadata exactly match the
|
|
382
|
+
compiled bound source. The runtime retains the buffer until the source
|
|
383
|
+
is rebound or the runtime is closed.
|
|
384
|
+
"""
|
|
385
|
+
self._require_open()
|
|
386
|
+
try:
|
|
387
|
+
source = self._sources_by_name[name]
|
|
388
|
+
except KeyError as error:
|
|
389
|
+
raise KeyError(f"unknown runtime source {name!r}") from error
|
|
390
|
+
if source.binding is not SourceBinding.BOUND:
|
|
391
|
+
raise ValueError(f"source {name!r} is not a bound source")
|
|
392
|
+
|
|
393
|
+
view = _buffer_view(source.tensor, values)
|
|
394
|
+
if view is None:
|
|
395
|
+
raise TypeError("bound sources require an object exporting a buffer")
|
|
396
|
+
if not _physical_buffer_matches(source.tensor, view):
|
|
397
|
+
raise ValueError("buffer does not match the compiled source layout")
|
|
398
|
+
address = _writable_pointer(view)
|
|
399
|
+
if address is None:
|
|
400
|
+
raise ValueError("zero-copy binding requires a writable contiguous buffer")
|
|
401
|
+
anchor, pointer = address
|
|
402
|
+
require_ok(
|
|
403
|
+
self._library.zgc_runtime_bind_source(
|
|
404
|
+
self._runtime,
|
|
405
|
+
source.key,
|
|
406
|
+
pointer,
|
|
407
|
+
view.nbytes,
|
|
408
|
+
),
|
|
409
|
+
"zgc_runtime_bind_source",
|
|
410
|
+
)
|
|
411
|
+
self._bound_storage[source.key] = (values, view, anchor)
|
|
412
|
+
|
|
413
|
+
def run(self, **sources: object) -> tuple[Tensor, ...]:
|
|
414
|
+
self._require_open()
|
|
415
|
+
for name, values in sources.items():
|
|
416
|
+
self.set_source(name, values)
|
|
417
|
+
require_ok(self._library.zgc_runtime_run(self._runtime), "zgc_runtime_run")
|
|
418
|
+
return tuple(self.output(index) for index in range(len(self.outputs)))
|
|
419
|
+
|
|
420
|
+
def output(self, index: int) -> Tensor:
|
|
421
|
+
"""Copied output from the runtime."""
|
|
422
|
+
self._require_open()
|
|
423
|
+
if index < 0:
|
|
424
|
+
raise IndexError(f"runtime has no output {index}")
|
|
425
|
+
try:
|
|
426
|
+
info = self.outputs[index]
|
|
427
|
+
except IndexError as error:
|
|
428
|
+
raise IndexError(f"runtime has no output {index}") from error
|
|
429
|
+
storage = bytearray(info.byte_count)
|
|
430
|
+
view = memoryview(storage)
|
|
431
|
+
address = _writable_pointer(view)
|
|
432
|
+
if address is None:
|
|
433
|
+
raise ValueError("cannot create empty output storage")
|
|
434
|
+
anchor, pointer = address
|
|
435
|
+
require_ok(
|
|
436
|
+
self._library.zgc_runtime_copy_output(
|
|
437
|
+
self._runtime,
|
|
438
|
+
index,
|
|
439
|
+
pointer,
|
|
440
|
+
info.byte_count,
|
|
441
|
+
),
|
|
442
|
+
"zgc_runtime_copy_output",
|
|
443
|
+
)
|
|
444
|
+
_ = anchor
|
|
445
|
+
return Tensor(info.dtype, info.shape, storage)
|
|
446
|
+
|
|
447
|
+
def close(self) -> None:
|
|
448
|
+
if self._closed:
|
|
449
|
+
return
|
|
450
|
+
self._library.zgc_runtime_deinit(self._runtime)
|
|
451
|
+
self._bound_storage.clear()
|
|
452
|
+
self._closed = True
|
|
453
|
+
|
|
454
|
+
def _require_open(self) -> None:
|
|
455
|
+
if self._closed:
|
|
456
|
+
raise RuntimeError("runtime is closed")
|
|
457
|
+
|
|
458
|
+
def __enter__(self) -> Self:
|
|
459
|
+
self._require_open()
|
|
460
|
+
return self
|
|
461
|
+
|
|
462
|
+
def __exit__(self, *_: object) -> None:
|
|
463
|
+
self.close()
|
|
464
|
+
|
|
465
|
+
def __del__(self) -> None:
|
|
466
|
+
if hasattr(self, "_closed"):
|
|
467
|
+
self.close()
|
|
468
|
+
|
|
469
|
+
|
|
470
|
+
def load(path: str | PathLike[str]) -> Runtime:
|
|
471
|
+
"""Load a pre-compiled runtime from a local path."""
|
|
472
|
+
return Runtime(path)
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
pyproject.toml
|
|
2
|
+
src/zgc/__init__.py
|
|
3
|
+
src/zgc/_abi.py
|
|
4
|
+
src/zgc/py.typed
|
|
5
|
+
src/zgc/runtime.py
|
|
6
|
+
src/zgc_runtime.egg-info/PKG-INFO
|
|
7
|
+
src/zgc_runtime.egg-info/SOURCES.txt
|
|
8
|
+
src/zgc_runtime.egg-info/dependency_links.txt
|
|
9
|
+
src/zgc_runtime.egg-info/top_level.txt
|
|
10
|
+
tests/test_buffers.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
zgc
|
|
@@ -0,0 +1,113 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
# pyright: reportPrivateUsage=false
|
|
4
|
+
import ctypes
|
|
5
|
+
import struct
|
|
6
|
+
import unittest
|
|
7
|
+
from array import array
|
|
8
|
+
from typing import final
|
|
9
|
+
|
|
10
|
+
from zgc import Runtime, Source, SourceBinding, SourceKind, Tensor, TensorInfo, dtype
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@final
|
|
14
|
+
class _Library:
|
|
15
|
+
copied_source: bytes = b""
|
|
16
|
+
bound_pointer: int = 0
|
|
17
|
+
bound_size: int = 0
|
|
18
|
+
output: bytes = struct.pack("=4f", 5.0, 6.0, 7.0, 8.0)
|
|
19
|
+
|
|
20
|
+
def zgc_runtime_copy_source(
|
|
21
|
+
self,
|
|
22
|
+
_runtime: ctypes.c_void_p,
|
|
23
|
+
_source_key: int,
|
|
24
|
+
data: ctypes.c_void_p,
|
|
25
|
+
size: int,
|
|
26
|
+
) -> int:
|
|
27
|
+
self.copied_source = ctypes.string_at(data, size)
|
|
28
|
+
return 0
|
|
29
|
+
|
|
30
|
+
def zgc_runtime_bind_source(
|
|
31
|
+
self,
|
|
32
|
+
_runtime: ctypes.c_void_p,
|
|
33
|
+
_source_key: int,
|
|
34
|
+
data: ctypes.c_void_p,
|
|
35
|
+
size: int,
|
|
36
|
+
) -> int:
|
|
37
|
+
self.bound_pointer = data.value or 0
|
|
38
|
+
self.bound_size = size
|
|
39
|
+
return 0
|
|
40
|
+
|
|
41
|
+
def zgc_runtime_copy_output(
|
|
42
|
+
self,
|
|
43
|
+
_runtime: ctypes.c_void_p,
|
|
44
|
+
_output_index: int,
|
|
45
|
+
destination: ctypes.c_void_p,
|
|
46
|
+
size: int,
|
|
47
|
+
) -> int:
|
|
48
|
+
_result = ctypes.memmove(
|
|
49
|
+
destination,
|
|
50
|
+
self.output,
|
|
51
|
+
min(size, len(self.output)),
|
|
52
|
+
)
|
|
53
|
+
return 0
|
|
54
|
+
|
|
55
|
+
def zgc_runtime_deinit(self, _: ctypes.c_void_p) -> None:
|
|
56
|
+
pass
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def _runtime(binding: SourceBinding) -> tuple[Runtime, _Library]:
|
|
60
|
+
info = TensorInfo(dtype.f32, (2, 2), (2, 1), 4, 16, 0)
|
|
61
|
+
source = Source(0, "input", SourceKind.INPUT, binding, info)
|
|
62
|
+
library = _Library()
|
|
63
|
+
runtime = Runtime.__new__(Runtime)
|
|
64
|
+
runtime._library = library # pyright: ignore[reportAttributeAccessIssue]
|
|
65
|
+
runtime._runtime = ctypes.c_void_p(1)
|
|
66
|
+
runtime._closed = False
|
|
67
|
+
runtime._bound_storage = {}
|
|
68
|
+
runtime._sources_by_name = {"input": source}
|
|
69
|
+
runtime.outputs = (info,)
|
|
70
|
+
return runtime, library
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class BufferTests(unittest.TestCase):
|
|
74
|
+
def test_list_source_uses_scalar_fallback(self) -> None:
|
|
75
|
+
runtime, library = _runtime(SourceBinding.OWNED)
|
|
76
|
+
runtime.set_source("input", [1.0, 2.0, 3.0, 4.0])
|
|
77
|
+
self.assertEqual(library.copied_source, struct.pack("=4f", 1, 2, 3, 4))
|
|
78
|
+
|
|
79
|
+
def test_array_source_uses_buffer_contents(self) -> None:
|
|
80
|
+
runtime, library = _runtime(SourceBinding.OWNED)
|
|
81
|
+
values = array("f", [1.0, 2.0, 3.0, 4.0])
|
|
82
|
+
runtime.set_source("input", values)
|
|
83
|
+
self.assertEqual(library.copied_source, values.tobytes())
|
|
84
|
+
|
|
85
|
+
def test_bound_array_remains_borrowed(self) -> None:
|
|
86
|
+
runtime, library = _runtime(SourceBinding.BOUND)
|
|
87
|
+
values = array("f", [1.0, 2.0, 3.0, 4.0])
|
|
88
|
+
shaped = memoryview(values).cast("B").cast("f", shape=[2, 2])
|
|
89
|
+
runtime.bind("input", shaped)
|
|
90
|
+
values[0] = 9.0
|
|
91
|
+
payload = ctypes.string_at(library.bound_pointer, library.bound_size)
|
|
92
|
+
self.assertEqual(payload, values.tobytes())
|
|
93
|
+
|
|
94
|
+
def test_set_source_copies_bound_values(self) -> None:
|
|
95
|
+
runtime, library = _runtime(SourceBinding.BOUND)
|
|
96
|
+
values = array("f", [1.0, 2.0, 3.0, 4.0])
|
|
97
|
+
runtime.set_source("input", values)
|
|
98
|
+
values[0] = 9.0
|
|
99
|
+
payload = ctypes.string_at(library.bound_pointer, library.bound_size)
|
|
100
|
+
self.assertEqual(payload, struct.pack("=4f", 1.0, 2.0, 3.0, 4.0))
|
|
101
|
+
|
|
102
|
+
def test_output_owns_a_shaped_buffer(self) -> None:
|
|
103
|
+
runtime, _ = _runtime(SourceBinding.OWNED)
|
|
104
|
+
output = runtime.output(0)
|
|
105
|
+
self.assertIsInstance(output, Tensor)
|
|
106
|
+
self.assertEqual(output.buffer.format, "f")
|
|
107
|
+
self.assertEqual(output.buffer.shape, (2, 2))
|
|
108
|
+
self.assertEqual(output.strides, (8, 4))
|
|
109
|
+
self.assertEqual(output.tolist(), [5.0, 6.0, 7.0, 8.0])
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
if __name__ == "__main__":
|
|
113
|
+
_ = unittest.main()
|