-
Notifications
You must be signed in to change notification settings - Fork 190
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Move external allocators into rmm.allocators module to defer imports (#…
…1221) RMM provides callbacks to configure third-party libraries to use RMM for memory allocation. Previously, these were defined in the top-level package, but that requires (potentially expensive) import of the package we're providing a hook for, since typically we must import that package to define the callback. This makes importing RMM expensive. To avoid this, move the callbacks into (not imported by default) sub-modules in `rmm.allocators`. So, if we want to configure the CuPy allocator, we now import `rmm_cupy_allocator` from `rmm.allocators.cupy`, and don't pay the price of importing pytorch. This change **deprecates** the use of the allocator callbacks in the top-level `rmm` module in favour of explicit imports from the relevant `rmm.allocators.XXX` sub-module. Before these changes, a sampling trace of `import rmm` with pyinstrument shows: $ pyinstrument -i 0.01 importrmm.py _ ._ __/__ _ _ _ _ _/_ Recorded: 10:19:56 Samples: 67 /_//_/// /_\ / //_// / //_'/ // Duration: 0.839 CPU time: 0.837 / _/ v4.4.0 Program: importrmm.py 0.839 <module> importrmm.py:1 └─ 0.839 <module> rmm/__init__.py:1 ├─ 0.315 <module> rmm/allocators/torch.py:1 │ └─ 0.315 <module> torch/__init__.py:1 │ [96 frames hidden] torch, <built-in>, enum, inspect, tok... ├─ 0.297 <module> rmm/mr.py:1 │ └─ 0.297 <module> rmm/_lib/__init__.py:1 │ ├─ 0.216 <module> numba/__init__.py:1 │ │ [140 frames hidden] numba, abc, <built-in>, importlib, em... │ ├─ 0.040 <module> numba/cuda/__init__.py:1 │ │ [34 frames hidden] numba, asyncio, ssl, <built-in>, re, ... │ ├─ 0.030 __new__ enum.py:180 │ │ [5 frames hidden] enum, <built-in> │ └─ 0.011 [self] None └─ 0.227 <module> rmm/allocators/cupy.py:1 └─ 0.227 <module> cupy/__init__.py:1 [123 frames hidden] cupy, pytest, _pytest, attr, <built-i... That is, almost a full second to import things, most of which is spent importing pytorch and cupy. These modules are not needed in normal usage of RMM, so we can defer the imports. Numba is a little bit trickier, but we can also defer up-front imports, with a final result that after these changes the same `import rmm` call takes just a tenth of a second: $ pyinstrument -i 0.01 importrmm.py _ ._ __/__ _ _ _ _ _/_ Recorded: 10:37:40 Samples: 9 /_//_/// /_\ / //_// / //_'/ // Duration: 0.099 CPU time: 0.099 / _/ v4.4.0 Program: importrmm.py 0.099 <module> importrmm.py:1 └─ 0.099 <module> rmm/__init__.py:1 └─ 0.099 <module> rmm/mr.py:1 └─ 0.099 <module> rmm/_lib/__init__.py:1 ├─ 0.059 <module> numpy/__init__.py:1 │ [31 frames hidden] numpy, re, sre_compile, <built-in>, s... ├─ 0.020 __new__ enum.py:180 │ [2 frames hidden] enum ├─ 0.010 <module> ctypes/__init__.py:1 │ [3 frames hidden] ctypes, <built-in> └─ 0.010 _EnumDict.__setitem__ enum.py:89 [3 frames hidden] enum Closes #1211. Authors: - Lawrence Mitchell (https://github.com/wence-) - Bradley Dice (https://github.com/bdice) Approvers: - Vyas Ramasubramani (https://github.com/vyasr) - Bradley Dice (https://github.com/bdice) URL: #1221
- Loading branch information
Showing
13 changed files
with
311 additions
and
201 deletions.
There are no files selected for viewing
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
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
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
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
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
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
Empty file.
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,44 @@ | ||
# Copyright (c) 2023, NVIDIA CORPORATION. | ||
# | ||
# Licensed under the Apache License, Version 2.0 (the "License"); | ||
# you may not use this file except in compliance with the License. | ||
# You may obtain a copy of the License at | ||
# | ||
# http://www.apache.org/licenses/LICENSE-2.0 | ||
# | ||
# 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. | ||
from rmm import _lib as librmm | ||
from rmm._cuda.stream import Stream | ||
|
||
try: | ||
import cupy | ||
except ImportError: | ||
cupy = None | ||
|
||
|
||
def rmm_cupy_allocator(nbytes): | ||
""" | ||
A CuPy allocator that makes use of RMM. | ||
Examples | ||
-------- | ||
>>> from rmm.allocators.cupy import rmm_cupy_allocator | ||
>>> import cupy | ||
>>> cupy.cuda.set_allocator(rmm_cupy_allocator) | ||
""" | ||
if cupy is None: | ||
raise ModuleNotFoundError("No module named 'cupy'") | ||
|
||
stream = Stream(obj=cupy.cuda.get_current_stream()) | ||
buf = librmm.device_buffer.DeviceBuffer(size=nbytes, stream=stream) | ||
dev_id = -1 if buf.ptr else cupy.cuda.device.get_device_id() | ||
mem = cupy.cuda.UnownedMemory( | ||
ptr=buf.ptr, size=buf.size, owner=buf, device_id=dev_id | ||
) | ||
ptr = cupy.cuda.memory.MemoryPointer(mem, 0) | ||
|
||
return ptr |
Oops, something went wrong.