refactor: refactor storage module and add support for NumPy >= 2.0 #1559
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
egparedes
changed the title
havogt
reviewed
Jun 25, 2024
| TensorShape: TypeAlias = Sequence[ | ||
| int | ||
| ] # TODO(egparedes) figure out if PositiveIntegral can be made to work | ||
| BufferStrides: TypeAlias = Sequence[IntScalar] |
| @dataclasses.dataclass(frozen=True, init=False) | ||
| class _BaseNDArrayBufferAllocator(abc.ABC, Generic[core_defs.DeviceTypeT]): | ||
| """Base class for buffer allocators using NumPy-like modules.""" | ||
| class MemoryResourceHandler(Protocol[core_defs.DeviceTypeLiteralT]): |
There was a problem hiding this comment.
Maybe just
Suggested change
| class MemoryResourceHandler(Protocol[core_defs.DeviceTypeLiteralT]): | |
| class Allocator(Protocol[core_defs.DeviceTypeLiteralT]): |
or even _Allocator?
There was a problem hiding this comment.
or BufferAllocator if you think Allocator is too generic.
There was a problem hiding this comment.
After reading more in this file, my suggestion doesn't make sense for what this is intended to do, but I am not sure I like this grouping into a class with classmethods.
The trigger for this change was byte_bounds, the following solutions looks cleaner to me:
mallocreturns buffer and address, because that's the only place we use it (if I didn't miss anything)address_ofas single dispatch and we register cupy and numpy buffer
tensorize is duplicated, it seems cleaner to use a free function and pass the namespace
There was a problem hiding this comment.
While looking at this I am trying to understand at which abstraction level would we add a low level pool allocator.
havogt
added a commit
that referenced
this pull request
Jun 27, 2024
requires gridtools_cpp >= 2.3.4 for nanobind 2.x Additional changes: freezes numpy to 1.x until #1559
|
Close in favor of #1563 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Refactor the storage module to encapsulate the required functionality from array libraries to allocate memory buffers in a new concept, which simplifies the design at the same time by removing one layer.
Add:
MemoryResourceHandlerconcept for low-level memory buffer handling, and two different implementations using NumPy and CuPy, which support newest changes in NumPy 2.0 interfacestoragemodules and a TODO note to migrate the old unit tests fromcartesianChange:
DeviceTypeliteral values in_core.definitionsstorage.allocatorsinstead of being duplicated incartesianandnextallocatorsDelete:
ValidNumPyLikeAllocationNSfromstorage.allocators, which is not longer needed