Repository navigation
Expand file tree
/
Copy pathbuffer.py
More file actions
478 lines (389 loc) · 16.9 KB
/
Copy pathbuffer.py
File metadata and controls
478 lines (389 loc) · 16.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
# ===----------------------------------------------------------------------=== #
# Copyright (c) 2026, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #
"""Provides extended utilities for :obj:`max.driver.Buffer`, including dtype casting, numpy interop, and DLPack integration."""
from __future__ import annotations
import struct
from collections.abc import Generator, Sequence
from itertools import product
from os import PathLike
from typing import Any
import numpy as np
import numpy.typing as npt
from max._core.driver import Buffer as Buffer
from max._core.driver import _batch_inplace_copy
from max._core.driver import (
_ChunkedStagingRegion as _ChunkedStagingRegion,
)
from max.dtype import DType
from .driver import CPU
_IdxElType = int | slice
IndexType = Sequence[_IdxElType] | _IdxElType
ShapeType = Sequence[int]
def _iterate_indices(self: Buffer) -> Generator[ShapeType]:
yield from product(*map(range, self.shape))
def _contiguous(self: Buffer) -> Buffer:
"""Creates a contiguous copy of the parent buffer."""
buffer_copy = Buffer(self.dtype, self.shape)
for idx in self._iterate_indices():
buffer_copy[idx] = self[idx].item()
return buffer_copy
def _repr(self: Buffer) -> str:
return f"max.driver.Buffer({self.dtype}, {self.shape}, {self.stream})"
# DTypes whose values NumPy can represent faithfully. These allow a fast bulk
# conversion for display. Other float types (bfloat16, float8_*, float4_*) are
# read element-wise so their values render as numbers rather than raw bytes.
_NUMPY_FAITHFUL_DTYPES = frozenset(
{
DType.bool,
DType.int8,
DType.int16,
DType.int32,
DType.int64,
DType.uint8,
DType.uint16,
DType.uint32,
DType.uint64,
DType.float16,
DType.float32,
DType.float64,
}
)
def _display_numpy(self: Buffer) -> npt.NDArray[Any]:
"""Returns a host numpy array of the buffer's values for display.
Faithfully-representable dtypes are converted in bulk via ``to_numpy``.
Other float dtypes (such as ``bfloat16`` and the ``float8`` family) are
read element-wise so their values are shown as decoded numbers instead of
the raw byte contents that ``to_numpy`` would expose.
"""
if self.dtype in _NUMPY_FAITHFUL_DTYPES:
return self.to_numpy()
# Copy to the host once so per-element reads don't trigger a device
# transfer for every element.
host = self if (self.pinned or self.device.is_host) else self.to(CPU())
if host.rank == 0:
return np.asarray(host.item())
values = [host[idx].item() for idx in host._iterate_indices()]
return np.asarray(values).reshape(self.shape)
def _str(self: Buffer) -> str:
"""Returns a human-readable string containing the buffer's data.
The data is copied to the host if the buffer lives on an accelerator, then
formatted like a numpy array (large buffers are summarized with ``...``).
The values are followed by the buffer's ``dtype``, ``shape``, and
``device``.
"""
prefix = "Buffer("
data = np.array2string(_display_numpy(self), prefix=prefix)
return (
f"{prefix}{data}, dtype={self.dtype}, shape={self.shape},"
f" device={self.device})"
)
def _view(self: Buffer, dtype: DType, shape: ShapeType | None = None) -> Buffer:
"""Return a new buffer with the given type and shape that shares the underlying memory.
If ``shape`` is not given, it is deduced by resizing the last axis to fit
the new dtype. A rank-0 buffer has no axis to resize, so it keeps shape
``()`` and only accepts a dtype of the same size. A ``ValueError`` is raised
when no shape can be deduced.
"""
if shape is None:
if self.rank == 0:
# Matches numpy, which rejects itemsize changes on 0-d arrays.
if dtype.size_in_bytes != self.element_size:
raise ValueError(
"Cannot deduce a shape when viewing a rank-0 buffer as a"
" dtype of a different size."
)
shape = ()
else:
last_axis_size = self.element_size * self.shape[-1]
if last_axis_size % dtype.size_in_bytes:
raise ValueError(
"When changing to a different dtype, its size must be a"
" divisor of the total size in bytes of the last axis of"
" the array."
)
shape = (*self.shape[:-1], last_axis_size // dtype.size_in_bytes)
return self._view(dtype, shape)
def inplace_copy_from(self: Buffer, src: Buffer) -> None:
"""Copy the contents of another buffer into this one.
These buffers may be on different devices.
Requires that both buffers are contiguous and have same size.
"""
if self is src:
return
# check that both buffers are contiguous
if not self.is_contiguous:
raise ValueError("Cannot copy from non-contiguous buffer")
if not src.is_contiguous:
raise ValueError("Cannot copy to non-contiguous buffer")
# check that both buffers have same size
if self.num_elements != src.num_elements:
raise ValueError("Cannot copy buffers of different sizes")
# check that both buffers have the same dtype
if self.dtype != src.dtype:
raise ValueError("Cannot copy buffers of different dtypes")
self._inplace_copy_from(src)
def _from_numpy(arr: npt.NDArray[Any]) -> Buffer:
"""Creates a buffer from a provided numpy array on the host device.
The underlying data is not copied unless the array is noncontiguous. If
it is, a contiguous copy will be returned.
Args:
arr: The numpy array to wrap.
Returns:
A new buffer that views or copies the numpy array data.
"""
# NOTE: np.ascontiguousarray only copies if needed.
# Skip np.contiguousarray for scalars since it converts them to rank-1.
return Buffer.from_dlpack(np.ascontiguousarray(arr) if arr.shape else arr)
def _to_numpy(self: Buffer) -> npt.NDArray[Any]:
"""Converts the buffer to a numpy array.
If the buffer is not on the host, a copy will be issued.
"""
if self.pinned or self.device.is_host:
cpu_buf = self
else:
cpu_buf = self.to(CPU())
try:
arr = np.from_dlpack(cpu_buf)
except RuntimeError as e:
if str(e).startswith("Unsupported device in DLTensor"):
raise RuntimeError(
f"Cannot convert buffer on {self.device} to numpy; move to"
" the host using `Buffer.to`"
) from e
raise
# DLPack v0 capsules may produce read-only arrays in numpy >= 1.26.
# Callers expect a writable view, so return a writable copy when needed.
if not arr.flags.writeable:
arr = arr.copy()
return arr
def _from_dlpack(array: Any, *, copy: bool | None = None) -> Buffer:
"""Create a buffer from an object implementing the dlpack protocol.
This usually does not result in a copy, and the producer of the object
retains ownership of the underlying memory.
Args:
array: An object that implements the dlpack protocol.
copy: Whether to create a copy of the data. Defaults to ``None``,
which copies only when the producer requires it.
Returns:
A new buffer that views or copies the data.
"""
if isinstance(array, np.ndarray):
if not array.flags.c_contiguous:
raise ValueError(
"driver buffer's from_dlpack only accepts contiguous arrays. "
"First call np.ascontiguousarray(array)"
)
# TODO(MSDK-976): Older version of numpy don't support exporting
# read-only arrays, so we copy if we can, and leave a hint if not.
if copy is None and not array.flags.writeable:
copy = True
if copy:
array = array.copy()
# Numpy's dlpack implementation cannot handle its own bool types, so
# we trick it into thinking it is uint8.
is_bool = array.dtype == bool
if is_bool:
array = array.view(np.uint8)
try:
buffer = Buffer._from_dlpack(array)
except BufferError as e:
msg = str(e)
if msg.startswith("Cannot export readonly array"):
raise type(e)( # noqa: B904
msg
+ " Consider passing `copy = True` to `Buffer.from_dlpack`."
)
raise
return buffer.view(DType.bool) if is_bool else buffer
# Short circuit if it's our type.
if isinstance(array, Buffer):
return array.copy() if copy else array
# Check for torch tensors by looking for the is_contiguous method rather
# than importing torch.
if hasattr(array, "is_contiguous") and callable(
getattr(array, "is_contiguous") # noqa: B009
):
if not array.is_contiguous():
raise ValueError(
"driver buffer's from_dlpack only accepts contiguous tensors. "
"First call .contiguous() on the tensor"
)
if copy is not None:
raise ValueError(
"`Buffer.from_dlpack` supports the `copy` flag only for numpy"
" array and `Buffer` inputs"
)
return Buffer._from_dlpack(array)
# TODO(MAXPLAT-206): re-enable @wraps
# @wraps(Buffer.mmap)
def _mmap(
filename: PathLike[str] | str,
dtype: DType,
shape: ShapeType | int,
mode: np._MemMapModeKind = "copyonwrite",
offset: int = 0,
) -> Buffer:
"""Creates a memory-mapped buffer from a binary file on disk.
The argument semantics follow ``np.memmap``.
Args:
filename: Path of the binary file to map.
dtype: Element type of the mapped data.
shape: Shape of the mapped buffer, either a tuple or an ``int``.
mode: Memory-map mode, following ``np.memmap`` semantics.
offset: Byte offset into the file where the mapping starts.
Returns:
A new buffer backed by the memory-mapped file.
"""
arr: np.memmap[Any, Any] = np.memmap(
filename,
dtype.to_numpy(),
mode,
offset,
# NOTE: prior to NumPy 2.0, `shape` must be `tuple` or `int`.
shape if isinstance(shape, int) else tuple(shape),
order="C",
)
assert arr.flags["C_CONTIGUOUS"]
return Buffer.from_dlpack(arr)
Buffer._iterate_indices = _iterate_indices # type: ignore[method-assign]
Buffer.contiguous = _contiguous # type: ignore[method-assign]
Buffer.__repr__ = _repr # type: ignore[method-assign, assignment]
Buffer.__str__ = _str # type: ignore[method-assign, assignment]
Buffer.view = _view # type: ignore[method-assign]
Buffer.inplace_copy_from = inplace_copy_from # type: ignore[method-assign]
Buffer.from_numpy = staticmethod(_from_numpy) # type: ignore[method-assign]
Buffer.to_numpy = _to_numpy # type: ignore[method-assign]
Buffer.from_dlpack = staticmethod(_from_dlpack) # type: ignore[method-assign]
Buffer.mmap = staticmethod(_mmap) # type: ignore[method-assign]
def batch_inplace_copy(dsts: Sequence[Buffer], srcs: Sequence[Buffer]) -> None:
"""Copies ``srcs[i]`` into ``dsts[i]`` in as few driver submissions as possible.
Sources may be host, pinned, same-device or peer memory in any mix.
Identity pairs (``dst is src``), zero-byte pairs and repeats of an earlier
pair's destination, source and size (on the same devices) are skipped.
One submission is one stream, and a stream only orders the writes it
performs, so destinations are grouped by device and each device's group is
submitted on that device's own stream. Callers do not need to group.
Args:
dsts: Destination buffers, in any mix of devices.
srcs: Source buffers, paired with ``dsts`` by index.
Raises:
ValueError: If ``dsts``/``srcs`` lengths differ.
"""
if len(dsts) != len(srcs):
raise ValueError(
"batch_inplace_copy: dsts and srcs must have the same length"
f" (got {len(dsts)} dsts, {len(srcs)} srcs)"
)
# Dedupe identical copy operations.
seen: set[tuple[str, int, int, str, int, int, int]] = set()
unique_dsts: list[Buffer] = []
unique_srcs: list[Buffer] = []
for dst, src in zip(dsts, srcs, strict=True):
nbytes = dst.num_elements * dst.dtype.size_in_bytes
key = (
dst.device.label,
dst.device.id,
dst._data_ptr(),
src.device.label,
src.device.id,
src._data_ptr(),
nbytes,
)
if key not in seen:
seen.add(key)
unique_dsts.append(dst)
unique_srcs.append(src)
_batch_inplace_copy(unique_dsts, unique_srcs)
def copy_pinned_to_destinations(
source: Buffer, destinations: Sequence[Buffer]
) -> None:
"""Copies a host buffer into destination buffers safely.
Copies ``source`` into each buffer in ``destinations``. When ``source`` is
pinned and the destinations span multiple devices, its deferred free is
enqueued only on its owning device's stream, so the memory manager could
recycle the pinned staging while copies on the other devices' streams are
still reading it. Making the owning device's stream wait for the other
destination devices' streams closes that window. The cross-stream wait is a
no-op for a non-pinned ``source`` or a single device.
Args:
source: The host buffer to copy from (pinned for async H2D copies).
destinations: The per-device buffers to copy ``source`` into.
"""
for destination in destinations:
destination.inplace_copy_from(source)
if not source.pinned:
return
source_device = source.device
source_stream = source_device.default_queue
for destination in destinations:
if destination.device.id != source_device.id:
source_stream.wait_for(destination.device.default_queue)
def load_max_buffer(path: PathLike[str]) -> Buffer:
"""Experimental method for loading serialized MAX buffers.
Max buffers can be exported by creating a graph and calling
``Value.print()`` with the ``BINARY_MAX_CHECKPOINT`` option.
Args:
path: Path to buffer (should end with .max)
Returns:
A :obj:`~max.driver.Buffer` created from the path. The shape and
dtype are read from the file.
Raises:
``ValueError``: If the file format is not the MAX checkpoint format.
"""
with open(path, "rb") as f:
header = f.read(8)
if header != b"\x93\xf0\x9f\x94\xa5\x2b\x2b\x93":
raise ValueError(
f"{path} is not a max checkpoint. If this file was saved "
'from the "BINARY" debug print option (and not '
'"BINARY_MAX_CHECKPOINT"), please initialize `Buffer.mmap` '
"directly."
)
# '2I' = 2 4-byte integers (major_version, minor_version).
major_version, _minor_version = struct.unpack("2I", f.read(8))
# Hardcoded but we should move to a robust versioning system if this
# method is ever used outside of debugging.
if major_version > 0:
raise ValueError("Unable to read from version > 0.")
# 'Q' = 8-byte unsigned long long (metadata_size).
metadata_size = struct.unpack("Q", f.read(8))[0]
# 'I' = 4-byte unsigned integer (key_size).
key_size = struct.unpack("I", f.read(4))[0]
f.read(key_size) # skip key
# '2B' = 2 unsigned bytes (dtype, rank).
dtype, rank = struct.unpack("2B", f.read(2))
dtype = DType(dtype)
# 'I' = 4-byte unsigned integer (each dimension in shape tuple).
shape = tuple(struct.unpack("I", f.read(4))[0] for _ in range(rank))
# 'Q' = 8-byte unsigned long long (offset).
offset = struct.unpack("Q", f.read(8))[0]
bytes_read = 4 + key_size + 2 + 4 * rank + 8
if bytes_read != metadata_size:
raise ValueError(
"Multiple buffers found in .max file. This is currently not supported."
)
if dtype == DType.bfloat16:
# Only modify last dimension for byte expansion.
new_shape = list(shape)
if len(new_shape) == 0:
# Handle scalar case.
new_shape = [2]
else:
# Expand last dimension for uint8 bytes.
new_shape[-1] *= 2
buffer = Buffer.mmap(
path, DType.uint8, new_shape, mode="r", offset=offset
)
return buffer.view(DType.bfloat16)
else:
return Buffer.mmap(path, dtype, shape, mode="r", offset=offset)