Skip to content

brahmap.mpi.SharedMemoryManager

Bases: object

Manages MPI shared-memory communicators, window allocations, and tree group reductions.

This manager splits a base MPI communicator into a node-level shared-memory communicator and allocates MPI window-backed shared NumPy arrays. It also splits the node-level communicator into a tree group sub-communicators to orchestrate sequential accumulations and group-wise reductions within each node.

Parameters:

Name Type Description Default
base_comm Intracomm

The base MPI communicator (typically MPI.COMM_WORLD)

 required
nproc_reduce int

The size of each local tree group sub-communicator split from the node-level communicator. This group size determines how the node-local processes are partitioned into smaller sub-communicator chunks. It is typically used by calling containers to orchestrate sequential shared-memory accumulations within each sub-communicator and group-wise reductions across group roots. By default 1

 1
node_root int

The designated root rank within the node-level shared memory communicator. By default 0

 0

Attributes:

Name Type Description
base_comm Intracomm

The base MPI communicator
node_comm Intracomm

The node-level shared-memory MPI communicator
node_rank int

The process rank within the node-level communicator
node_size int

The total number of processes on the current node
nproc_reduce int

The size of each local tree group sub-communicator split from the node-level communicator
node_root int

The root rank on the current node communicator
node_root_comm Comm or None

Communicator containing only the root ranks of each node, used for inter-node communication
tree_grp_comm Intracomm

Sub-communicator group for local serialized accumulations
tree_grp_rank int

The rank within the local tree group communicator
tree_grp_size int

The size of the local tree group communicator
tree_grp_root int

The designated root rank of the tree group (usually 0)
tree_grp_root_comm Intracomm

Communicator containing only the tree group roots on this node, used for intra-node aggregation
list_windows dict[int, list[Win]]

Tracks allocated shared-memory MPI windows mapped by communicator handle
list_arrays dict[int, list[NDArray]]

Tracks allocated shared-memory NumPy array views mapped by communicator handle

Methods:

Name Description
alloc_shared_comm

Allocates a shared-memory MPI window-backed 1D NumPy array for a
alloc_shared_node

Allocates a shared-memory MPI window-backed 1D NumPy array for the
alloc_shared_zeros_comm

Allocates a shared-memory MPI window-backed 1D NumPy array for a
alloc_shared_zeros_node

Allocates a shared-memory MPI window-backed 1D NumPy array for the
alloc_shared_ones_comm

Allocates a shared-memory MPI window-backed 1D NumPy array for a
alloc_shared_ones_node

Allocates a shared-memory MPI window-backed 1D NumPy array for the
fence_comm_all

Call MPI.Win.Fence on all windows allocated on the given
free_shared_arrays_all

Frees all allocated shared-memory MPI windows and clears manager
free_shared_arrays_comm

Frees all shared-memory MPI windows allocated for a specific
free_shared_array

Frees a specific shared-memory MPI window and removes its associated
free_all_resources

Frees all allocated shared-memory MPI windows and all split MPI
Source code in brahmap/mpi.py 
 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
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
class SharedMemoryManager(object):
    """Manages MPI shared-memory communicators, window allocations, and tree
    group reductions.

    This manager splits a base MPI communicator into a node-level
    shared-memory communicator and allocates MPI window-backed shared NumPy
    arrays. It also splits the node-level communicator into a tree group
    sub-communicators to orchestrate sequential accumulations and group-wise
    reductions within each node.

    Parameters
    ----------
    base_comm : Intracomm
        The base MPI communicator (typically `MPI.COMM_WORLD`)
    nproc_reduce : int, optional
        The size of each local tree group sub-communicator split from the
        node-level communicator. This group size determines how the
        node-local processes are partitioned into smaller sub-communicator
        chunks. It is typically used by calling containers to orchestrate
        sequential shared-memory accumulations within each sub-communicator
        and group-wise reductions across group roots. By default `1`
    node_root : int, optional
        The designated root rank within the node-level shared memory
        communicator. By default `0`

    Attributes
    ----------
    base_comm : Intracomm
        The base MPI communicator
    node_comm : Intracomm
        The node-level shared-memory MPI communicator
    node_rank : int
        The process rank within the node-level communicator
    node_size : int
        The total number of processes on the current node
    nproc_reduce : int
        The size of each local tree group sub-communicator split from the
        node-level communicator
    node_root : int
        The root rank on the current node communicator
    node_root_comm : Comm or None
        Communicator containing only the root ranks of each node, used for
        inter-node communication
    tree_grp_comm : Intracomm
        Sub-communicator group for local serialized accumulations
    tree_grp_rank : int
        The rank within the local tree group communicator
    tree_grp_size : int
        The size of the local tree group communicator
    tree_grp_root : int
        The designated root rank of the tree group (usually 0)
    tree_grp_root_comm : Intracomm
        Communicator containing only the tree group roots on this node, used
        for intra-node aggregation
    list_windows : dict[int, list[MPI.Win]]
        Tracks allocated shared-memory MPI windows mapped by communicator handle
    list_arrays : dict[int, list[npt.NDArray]]
        Tracks allocated shared-memory NumPy array views mapped by
        communicator handle
    """

    def __init__(
        self,
        base_comm: Intracomm,
        nproc_reduce: int = 1,
        node_root: int = 0,
    ) -> None:
        self._base_comm = base_comm
        self._node_comm: Intracomm = cast(
            Intracomm, self._base_comm.Split_type(MPI.COMM_TYPE_SHARED)
        )
        self._node_rank = self._node_comm.rank
        self._node_size = self._node_comm.size
        self._nproc_reduce = min(max(nproc_reduce, 1), self._node_size)
        self._node_root = node_root

        # node root communicator: a communicator that contains the
        # root/master/leaders of all node communicators
        root_color = 0 if self._node_rank == self._node_root else MPI.UNDEFINED
        self._node_root_comm: Comm | None = self._base_comm.Split(
            root_color, self._base_comm.rank
        )

        # This is block grouping, so `--map-by core` option would be most optimal
        tree_grp_color = self._node_rank // self._nproc_reduce
        self._tree_grp_comm = self._node_comm.Split(
            color=tree_grp_color, key=self._node_rank
        )
        self._tree_grp_rank = self._tree_grp_comm.rank
        self._tree_grp_size = self._tree_grp_comm.size
        self._tree_grp_root = 0
        tree_root_color = (
            0 if self._tree_grp_rank == self._tree_grp_root else MPI.UNDEFINED
        )
        self._tree_grp_root_comm = self._node_comm.Split(
            tree_root_color, self._tree_grp_comm.rank
        )

        # List of MPI shared memory windows
        self._list_windows: dict[int, list[MPI.Win]] = {}
        self._list_arrays: dict[int, list[npt.NDArray]] = {}

    @property
    def base_comm(self) -> Intracomm:
        """The base global MPI communicator

        Returns
        -------
        MPI.Intracomm
            The base MPI communicator
        """
        return self._base_comm

    @property
    def node_comm(self) -> Intracomm:
        """The node-level shared-memory MPI communicator

        Returns
        -------
        MPI.Intracomm
            The node-level MPI communicator
        """
        return self._node_comm

    @property
    def node_rank(self) -> int:
        """The process rank within the node-level communicator

        Returns
        -------
        int
            Node-local rank corresponding to the node-level communicator
        """
        return self._node_rank

    @property
    def node_size(self) -> int:
        """The total number of processes on the current node

        Returns
        -------
        int
            Size of this node-level communicator
        """
        return self._node_size

    @property
    def nproc_reduce(self) -> int:
        """The size of each local tree group sub-communicator split from the
        node-level communicator

        Returns
        -------
        int
            Size of the local tree group
        """
        return self._nproc_reduce

    @property
    def node_root(self) -> int:
        """The root rank on the current node

        Returns
        -------
        int
            Root rank of the node-level communicator
        """
        return self._node_root

    @property
    def node_root_comm(self) -> Comm | None:
        """The communicator containing only the root ranks of each
        node-level communicator, or None

        Returns
        -------
        MPI.Comm
            The node root communicator
        """
        return self._node_root_comm

    @property
    def tree_grp_comm(self) -> Intracomm:
        """The sub-communicator group for local tree-like serialized
        accumulations

        Returns
        -------
        MPI.Comm
            The local tree group communicator
        """
        return cast(Intracomm, self._tree_grp_comm)

    @property
    def tree_grp_rank(self) -> int:
        """The rank within the local tree group communicator

        Returns
        -------
        int
            Local tree group rank
        """
        return self._tree_grp_rank

    @property
    def tree_grp_size(self) -> int:
        """The size of the local tree group communicator

        Returns
        -------
        int
            Local tree group size
        """
        return self._tree_grp_size

    @property
    def tree_grp_root(self) -> int:
        """The designated root rank of the tree group

        Returns
        -------
        int
            Local tree group root
        """
        return self._tree_grp_root

    @property
    def tree_grp_root_comm(self) -> Comm | None:
        """The communicator containing only the tree group roots on this
        node, or None

        Returns
        -------
        MPI.Comm
            The tree group root communicator
        """
        return self._tree_grp_root_comm

    @property
    def list_windows(self) -> dict:
        """The dictionary mapping communicators to list of allocated
        shared-memory windows for that communicator

        Returns
        -------
        dict
            The dictionary mapping communicators to list of allocated
            shared-memory windows for that communicator
        """
        return self._list_windows

    @property
    def list_arrays(self) -> dict:
        """The dictionary mapping communicators to list of allocated
        shared-memory arrays for that communicator

        Returns
        -------
        dict
            The dictionary mapping communicators to list of allocated
            shared-memory arrays for that communicator
        """
        return self._list_arrays

    def alloc_shared_comm(
        self,
        size: int,
        dtype: npt.DTypeLike,
        comm: Intracomm,
        comm_root: int = 0,
    ) -> tuple[npt.NDArray, MPI.Win]:
        """Allocates a shared-memory MPI window-backed 1D NumPy array for a
        communicator.

        Parameters
        ----------
        size : int
            The size of the array
        dtype : npt.DTypeLike
            The data type of the array
        comm : Intracomm
            The MPI communicator over which the shared memory window is
            allocated
        comm_root : int, optional
            The root rank in `comm` that allocates the actual memory buffer.
            By default `0`

        Returns
        -------
        tuple[npt.NDArray, MPI.Win]
            A tuple containing the shared NumPy array view and the backing
            MPI window object
        """
        dtype = np.dtype(dtype)
        dtype_bytes = dtype.itemsize
        arr_bytes = size * dtype_bytes if comm.rank == comm_root else 0

        win = MPI.Win.Allocate_shared(
            arr_bytes,
            dtype_bytes,
            comm=comm,
        )
        buf, _ = win.Shared_query(rank=comm_root)
        # np.ndarray provides the view, it doesn't owns the memory
        array = np.ndarray(shape=size, dtype=dtype, buffer=buf)

        handle = comm.handle
        if handle not in self._list_windows:
            self._list_windows[handle] = []
        if handle not in self._list_arrays:
            self._list_arrays[handle] = []

        self._list_windows[handle].append(win)
        self._list_arrays[handle].append(array)
        return array, win

    def alloc_shared_node(
        self,
        size: int,
        dtype: npt.DTypeLike,
    ) -> tuple[npt.NDArray, MPI.Win]:
        """Allocates a shared-memory MPI window-backed 1D NumPy array for the
        node-level communicator

        Parameters
        ----------
        size : int
            The size of the array
        dtype : npt.DTypeLike
            The data type of the array

        Returns
        -------
        tuple[npt.NDArray, MPI.Win]
            A tuple containing the shared NumPy array view and the backing
            MPI window object
        """
        return self.alloc_shared_comm(
            size=size,
            dtype=dtype,
            comm=self.node_comm,
            comm_root=self.node_root,
        )

    def alloc_shared_zeros_comm(
        self,
        size: int,
        dtype: npt.DTypeLike,
        comm: Intracomm,
        comm_root: int = 0,
    ):
        """Allocates a shared-memory MPI window-backed 1D NumPy array for a
        communicator, initialized to zeros.

        Parameters
        ----------
        size : int
            The size of the array
        dtype : npt.DTypeLike
            The data type of the array
        comm : Intracomm
            The MPI communicator over which the shared memory window is
            allocated
        comm_root : int, optional
            The root rank in `comm` that allocates the actual memory buffer.
            By default `0`

        Returns
        -------
        tuple[npt.NDArray, MPI.Win]
            A tuple containing the shared NumPy array view and the backing
            MPI window object
        """
        array, win = self.alloc_shared_comm(
            size=size,
            dtype=dtype,
            comm=comm,
            comm_root=comm_root,
        )

        if comm.rank == 0:
            array[:] = 0

        return array, win

    def alloc_shared_zeros_node(
        self,
        size: int,
        dtype: npt.DTypeLike,
    ):
        """Allocates a shared-memory MPI window-backed 1D NumPy array for the
        node-level communicator, initialized to zeros.

        Parameters
        ----------
        size : int
            The size of the array
        dtype : npt.DTypeLike
            The data type of the array

        Returns
        -------
        tuple[npt.NDArray, MPI.Win]
            A tuple containing the shared NumPy array view and the backing
            MPI window object
        """
        return self.alloc_shared_zeros_comm(
            size=size,
            dtype=dtype,
            comm=self.node_comm,
            comm_root=self.node_root,
        )

    def alloc_shared_ones_comm(
        self,
        size: int,
        dtype: npt.DTypeLike,
        comm: Intracomm,
        comm_root: int = 0,
    ):
        """Allocates a shared-memory MPI window-backed 1D NumPy array for a
        communicator, initialized to ones.

        Parameters
        ----------
        size : int
            The size of the array
        dtype : npt.DTypeLike
            The data type of the array
        comm : Intracomm
            The MPI communicator over which the shared memory window is
            allocated
        comm_root : int, optional
            The root rank in `comm` that allocates the actual memory buffer.
            By default `0`

        Returns
        -------
        tuple[npt.NDArray, MPI.Win]
            A tuple containing the shared NumPy array view and the backing
            MPI window object
        """
        array, win = self.alloc_shared_comm(
            size=size,
            dtype=dtype,
            comm=comm,
            comm_root=comm_root,
        )

        if comm.rank == 0:
            array[:] = 1

        return array, win

    def alloc_shared_ones_node(
        self,
        size: int,
        dtype: npt.DTypeLike,
    ):
        """Allocates a shared-memory MPI window-backed 1D NumPy array for the
        node-level communicator, initialized to ones.

        Parameters
        ----------
        size : int
            The size of the array
        dtype : npt.DTypeLike
            The data type of the array

        Returns
        -------
        tuple[npt.NDArray, MPI.Win]
            A tuple containing the shared NumPy array view and the backing
            MPI window object
        """
        return self.alloc_shared_ones_comm(
            size=size,
            dtype=dtype,
            comm=self.node_comm,
            comm_root=self.node_root,
        )

    def fence_comm_all(self, comm: Intracomm, assertion: int = 0) -> None:
        """Call MPI.Win.Fence on all windows allocated on the given
        communicator.

        Parameters
        ----------
        comm : Intracomm
            The communicator for which to fence the windows
        assertion : int, optional
            The assertion flag for the fence call, by default 0

        Returns
        -------
        None
        """
        handle = comm.handle
        if handle in self._list_windows:
            for win in self._list_windows[handle]:
                win.Fence(assertion)

    def free_shared_arrays_all(self) -> None:
        """Frees all allocated shared-memory MPI windows and clears manager
        state.

        Since the window owns the actual memory buffer, freeing the windows
        also deallocates the underlying buffers of all tracked shared-memory
        arrays.

        Returns
        -------
        None
        """
        # np.ndarray() simply provides the view, it doesn't transfer the
        # memory ownership. The buffer is owned by the window, so freeing
        # the window frees the buffer as well.

        for comm, wins in self._list_windows.items():
            for win in wins:
                win.Free()
        self._list_windows = {}
        self._list_arrays = {}

    def free_shared_arrays_comm(self, comm: Intracomm) -> None:
        """Frees all shared-memory MPI windows allocated for a specific
        communicator.

        Parameters
        ----------
        comm : Intracomm
            The MPI communicator whose shared memory windows should be freed.

        Returns
        -------
        None
        """
        handle = comm.handle
        if handle in self._list_windows:
            for win in self._list_windows[handle]:
                win.Free()
            del self._list_windows[handle]
        if handle in self._list_arrays:
            del self._list_arrays[handle]

    def free_shared_array(self, comm: Intracomm, win: MPI.Win) -> None:
        """Frees a specific shared-memory MPI window and removes its associated
        array view and window from the manager's tracking lists.

        Parameters
        ----------
        comm : Intracomm
            The MPI communicator over which the shared memory window was
            allocated
        win : MPI.Win
            The MPI window object to be freed

        Returns
        -------
        None
        """
        handle = comm.handle
        if handle in self._list_windows and win in self._list_windows[handle]:
            idx = self._list_windows[handle].index(win)
            win.Free()
            self._list_windows[handle].pop(idx)
            self._list_arrays[handle].pop(idx)
            if not self._list_windows[handle]:
                del self._list_windows[handle]
            if not self._list_arrays[handle]:
                del self._list_arrays[handle]

    def free_all_resources(self) -> None:
        """Frees all allocated shared-memory MPI windows and all split MPI
        communicators created by the manager.

        Returns
        -------
        None
        """
        self.free_shared_arrays_all()
        if hasattr(self, "_tree_grp_root_comm") and self._tree_grp_root_comm:
            self._tree_grp_root_comm.Free()
        if hasattr(self, "_tree_grp_comm") and self._tree_grp_comm:
            self._tree_grp_comm.Free()
        if hasattr(self, "_node_root_comm") and self._node_root_comm:
            self._node_root_comm.Free()
        if hasattr(self, "_node_comm") and self._node_comm:
            self._node_comm.Free()

Attributes

base_comm: Intracomm property

The base global MPI communicator

Returns:

Type Description
Intracomm

The base MPI communicator

node_comm: Intracomm property

The node-level shared-memory MPI communicator

Returns:

Type Description
Intracomm

The node-level MPI communicator

node_rank: int property

The process rank within the node-level communicator

Returns:

Type Description
int

Node-local rank corresponding to the node-level communicator

node_size: int property

The total number of processes on the current node

Returns:

Type Description
int

Size of this node-level communicator

nproc_reduce: int property

The size of each local tree group sub-communicator split from the node-level communicator

Returns:

Type Description
int

Size of the local tree group

node_root: int property

The root rank on the current node

Returns:

Type Description
int

Root rank of the node-level communicator

node_root_comm: Comm | None property

The communicator containing only the root ranks of each node-level communicator, or None

Returns:

Type Description
Comm

The node root communicator

tree_grp_comm: Intracomm property

The sub-communicator group for local tree-like serialized accumulations

Returns:

Type Description
Comm

The local tree group communicator

tree_grp_rank: int property

The rank within the local tree group communicator

Returns:

Type Description
int

Local tree group rank

tree_grp_size: int property

The size of the local tree group communicator

Returns:

Type Description
int

Local tree group size

tree_grp_root: int property

The designated root rank of the tree group

Returns:

Type Description
int

Local tree group root

tree_grp_root_comm: Comm | None property

The communicator containing only the tree group roots on this node, or None

Returns:

Type Description
Comm

The tree group root communicator

list_windows: dict property

The dictionary mapping communicators to list of allocated shared-memory windows for that communicator

Returns:

Type Description
dict

The dictionary mapping communicators to list of allocated shared-memory windows for that communicator

list_arrays: dict property

The dictionary mapping communicators to list of allocated shared-memory arrays for that communicator

Returns:

Type Description
dict

The dictionary mapping communicators to list of allocated shared-memory arrays for that communicator

Methods:

alloc_shared_comm(size: int, dtype: npt.DTypeLike, comm: Intracomm, comm_root: int = 0) -> tuple[npt.NDArray, MPI.Win]

Allocates a shared-memory MPI window-backed 1D NumPy array for a communicator.

Parameters:

Name Type Description Default
size int

The size of the array

 required
dtype DTypeLike

The data type of the array

 required
comm Intracomm

The MPI communicator over which the shared memory window is allocated

 required
comm_root int

The root rank in comm that allocates the actual memory buffer. By default 0

 0

Returns:

Type Description
tuple[NDArray, Win]

A tuple containing the shared NumPy array view and the backing MPI window object
Source code in brahmap/mpi.py 
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
def alloc_shared_comm(
    self,
    size: int,
    dtype: npt.DTypeLike,
    comm: Intracomm,
    comm_root: int = 0,
) -> tuple[npt.NDArray, MPI.Win]:
    """Allocates a shared-memory MPI window-backed 1D NumPy array for a
    communicator.

    Parameters
    ----------
    size : int
        The size of the array
    dtype : npt.DTypeLike
        The data type of the array
    comm : Intracomm
        The MPI communicator over which the shared memory window is
        allocated
    comm_root : int, optional
        The root rank in `comm` that allocates the actual memory buffer.
        By default `0`

    Returns
    -------
    tuple[npt.NDArray, MPI.Win]
        A tuple containing the shared NumPy array view and the backing
        MPI window object
    """
    dtype = np.dtype(dtype)
    dtype_bytes = dtype.itemsize
    arr_bytes = size * dtype_bytes if comm.rank == comm_root else 0

    win = MPI.Win.Allocate_shared(
        arr_bytes,
        dtype_bytes,
        comm=comm,
    )
    buf, _ = win.Shared_query(rank=comm_root)
    # np.ndarray provides the view, it doesn't owns the memory
    array = np.ndarray(shape=size, dtype=dtype, buffer=buf)

    handle = comm.handle
    if handle not in self._list_windows:
        self._list_windows[handle] = []
    if handle not in self._list_arrays:
        self._list_arrays[handle] = []

    self._list_windows[handle].append(win)
    self._list_arrays[handle].append(array)
    return array, win

alloc_shared_node(size: int, dtype: npt.DTypeLike) -> tuple[npt.NDArray, MPI.Win]

Allocates a shared-memory MPI window-backed 1D NumPy array for the node-level communicator

Parameters:

Name Type Description Default
size int

The size of the array

 required
dtype DTypeLike

The data type of the array

 required

Returns:

Type Description
tuple[NDArray, Win]

A tuple containing the shared NumPy array view and the backing MPI window object
Source code in brahmap/mpi.py 
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
def alloc_shared_node(
    self,
    size: int,
    dtype: npt.DTypeLike,
) -> tuple[npt.NDArray, MPI.Win]:
    """Allocates a shared-memory MPI window-backed 1D NumPy array for the
    node-level communicator

    Parameters
    ----------
    size : int
        The size of the array
    dtype : npt.DTypeLike
        The data type of the array

    Returns
    -------
    tuple[npt.NDArray, MPI.Win]
        A tuple containing the shared NumPy array view and the backing
        MPI window object
    """
    return self.alloc_shared_comm(
        size=size,
        dtype=dtype,
        comm=self.node_comm,
        comm_root=self.node_root,
    )

alloc_shared_zeros_comm(size: int, dtype: npt.DTypeLike, comm: Intracomm, comm_root: int = 0)

Allocates a shared-memory MPI window-backed 1D NumPy array for a communicator, initialized to zeros.

Parameters:

Name Type Description Default
size int

The size of the array

 required
dtype DTypeLike

The data type of the array

 required
comm Intracomm

The MPI communicator over which the shared memory window is allocated

 required
comm_root int

The root rank in comm that allocates the actual memory buffer. By default 0

 0

Returns:

Type Description
tuple[NDArray, Win]

A tuple containing the shared NumPy array view and the backing MPI window object
Source code in brahmap/mpi.py 
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
def alloc_shared_zeros_comm(
    self,
    size: int,
    dtype: npt.DTypeLike,
    comm: Intracomm,
    comm_root: int = 0,
):
    """Allocates a shared-memory MPI window-backed 1D NumPy array for a
    communicator, initialized to zeros.

    Parameters
    ----------
    size : int
        The size of the array
    dtype : npt.DTypeLike
        The data type of the array
    comm : Intracomm
        The MPI communicator over which the shared memory window is
        allocated
    comm_root : int, optional
        The root rank in `comm` that allocates the actual memory buffer.
        By default `0`

    Returns
    -------
    tuple[npt.NDArray, MPI.Win]
        A tuple containing the shared NumPy array view and the backing
        MPI window object
    """
    array, win = self.alloc_shared_comm(
        size=size,
        dtype=dtype,
        comm=comm,
        comm_root=comm_root,
    )

    if comm.rank == 0:
        array[:] = 0

    return array, win

alloc_shared_zeros_node(size: int, dtype: npt.DTypeLike)

Allocates a shared-memory MPI window-backed 1D NumPy array for the node-level communicator, initialized to zeros.

Parameters:

Name Type Description Default
size int

The size of the array

 required
dtype DTypeLike

The data type of the array

 required

Returns:

Type Description
tuple[NDArray, Win]

A tuple containing the shared NumPy array view and the backing MPI window object
Source code in brahmap/mpi.py 
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
def alloc_shared_zeros_node(
    self,
    size: int,
    dtype: npt.DTypeLike,
):
    """Allocates a shared-memory MPI window-backed 1D NumPy array for the
    node-level communicator, initialized to zeros.

    Parameters
    ----------
    size : int
        The size of the array
    dtype : npt.DTypeLike
        The data type of the array

    Returns
    -------
    tuple[npt.NDArray, MPI.Win]
        A tuple containing the shared NumPy array view and the backing
        MPI window object
    """
    return self.alloc_shared_zeros_comm(
        size=size,
        dtype=dtype,
        comm=self.node_comm,
        comm_root=self.node_root,
    )

alloc_shared_ones_comm(size: int, dtype: npt.DTypeLike, comm: Intracomm, comm_root: int = 0)

Allocates a shared-memory MPI window-backed 1D NumPy array for a communicator, initialized to ones.

Parameters:

Name Type Description Default
size int

The size of the array

 required
dtype DTypeLike

The data type of the array

 required
comm Intracomm

The MPI communicator over which the shared memory window is allocated

 required
comm_root int

The root rank in comm that allocates the actual memory buffer. By default 0

 0

Returns:

Type Description
tuple[NDArray, Win]

A tuple containing the shared NumPy array view and the backing MPI window object
Source code in brahmap/mpi.py 
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
def alloc_shared_ones_comm(
    self,
    size: int,
    dtype: npt.DTypeLike,
    comm: Intracomm,
    comm_root: int = 0,
):
    """Allocates a shared-memory MPI window-backed 1D NumPy array for a
    communicator, initialized to ones.

    Parameters
    ----------
    size : int
        The size of the array
    dtype : npt.DTypeLike
        The data type of the array
    comm : Intracomm
        The MPI communicator over which the shared memory window is
        allocated
    comm_root : int, optional
        The root rank in `comm` that allocates the actual memory buffer.
        By default `0`

    Returns
    -------
    tuple[npt.NDArray, MPI.Win]
        A tuple containing the shared NumPy array view and the backing
        MPI window object
    """
    array, win = self.alloc_shared_comm(
        size=size,
        dtype=dtype,
        comm=comm,
        comm_root=comm_root,
    )

    if comm.rank == 0:
        array[:] = 1

    return array, win

alloc_shared_ones_node(size: int, dtype: npt.DTypeLike)

Allocates a shared-memory MPI window-backed 1D NumPy array for the node-level communicator, initialized to ones.

Parameters:

Name Type Description Default
size int

The size of the array

 required
dtype DTypeLike

The data type of the array

 required

Returns:

Type Description
tuple[NDArray, Win]

A tuple containing the shared NumPy array view and the backing MPI window object
Source code in brahmap/mpi.py 
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
def alloc_shared_ones_node(
    self,
    size: int,
    dtype: npt.DTypeLike,
):
    """Allocates a shared-memory MPI window-backed 1D NumPy array for the
    node-level communicator, initialized to ones.

    Parameters
    ----------
    size : int
        The size of the array
    dtype : npt.DTypeLike
        The data type of the array

    Returns
    -------
    tuple[npt.NDArray, MPI.Win]
        A tuple containing the shared NumPy array view and the backing
        MPI window object
    """
    return self.alloc_shared_ones_comm(
        size=size,
        dtype=dtype,
        comm=self.node_comm,
        comm_root=self.node_root,
    )

fence_comm_all(comm: Intracomm, assertion: int = 0) -> None

Call MPI.Win.Fence on all windows allocated on the given communicator.

Parameters:

Name Type Description Default
comm Intracomm

The communicator for which to fence the windows

 required
assertion int

The assertion flag for the fence call, by default 0

 0

Returns:

Type Description
None
Source code in brahmap/mpi.py 
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
def fence_comm_all(self, comm: Intracomm, assertion: int = 0) -> None:
    """Call MPI.Win.Fence on all windows allocated on the given
    communicator.

    Parameters
    ----------
    comm : Intracomm
        The communicator for which to fence the windows
    assertion : int, optional
        The assertion flag for the fence call, by default 0

    Returns
    -------
    None
    """
    handle = comm.handle
    if handle in self._list_windows:
        for win in self._list_windows[handle]:
            win.Fence(assertion)

free_shared_arrays_all() -> None

Frees all allocated shared-memory MPI windows and clears manager state.

Since the window owns the actual memory buffer, freeing the windows also deallocates the underlying buffers of all tracked shared-memory arrays.

Returns:

Type Description
None
Source code in brahmap/mpi.py 
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
def free_shared_arrays_all(self) -> None:
    """Frees all allocated shared-memory MPI windows and clears manager
    state.

    Since the window owns the actual memory buffer, freeing the windows
    also deallocates the underlying buffers of all tracked shared-memory
    arrays.

    Returns
    -------
    None
    """
    # np.ndarray() simply provides the view, it doesn't transfer the
    # memory ownership. The buffer is owned by the window, so freeing
    # the window frees the buffer as well.

    for comm, wins in self._list_windows.items():
        for win in wins:
            win.Free()
    self._list_windows = {}
    self._list_arrays = {}

free_shared_arrays_comm(comm: Intracomm) -> None

Frees all shared-memory MPI windows allocated for a specific communicator.

Parameters:

Name Type Description Default
comm Intracomm

The MPI communicator whose shared memory windows should be freed.

 required

Returns:

Type Description
None
Source code in brahmap/mpi.py 
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
def free_shared_arrays_comm(self, comm: Intracomm) -> None:
    """Frees all shared-memory MPI windows allocated for a specific
    communicator.

    Parameters
    ----------
    comm : Intracomm
        The MPI communicator whose shared memory windows should be freed.

    Returns
    -------
    None
    """
    handle = comm.handle
    if handle in self._list_windows:
        for win in self._list_windows[handle]:
            win.Free()
        del self._list_windows[handle]
    if handle in self._list_arrays:
        del self._list_arrays[handle]

free_shared_array(comm: Intracomm, win: MPI.Win) -> None

Frees a specific shared-memory MPI window and removes its associated array view and window from the manager's tracking lists.

Parameters:

Name Type Description Default
comm Intracomm

The MPI communicator over which the shared memory window was allocated

 required
win Win

The MPI window object to be freed

 required

Returns:

Type Description
None
Source code in brahmap/mpi.py 
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
def free_shared_array(self, comm: Intracomm, win: MPI.Win) -> None:
    """Frees a specific shared-memory MPI window and removes its associated
    array view and window from the manager's tracking lists.

    Parameters
    ----------
    comm : Intracomm
        The MPI communicator over which the shared memory window was
        allocated
    win : MPI.Win
        The MPI window object to be freed

    Returns
    -------
    None
    """
    handle = comm.handle
    if handle in self._list_windows and win in self._list_windows[handle]:
        idx = self._list_windows[handle].index(win)
        win.Free()
        self._list_windows[handle].pop(idx)
        self._list_arrays[handle].pop(idx)
        if not self._list_windows[handle]:
            del self._list_windows[handle]
        if not self._list_arrays[handle]:
            del self._list_arrays[handle]

free_all_resources() -> None

Frees all allocated shared-memory MPI windows and all split MPI communicators created by the manager.

Returns:

Type Description
None
Source code in brahmap/mpi.py 
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
def free_all_resources(self) -> None:
    """Frees all allocated shared-memory MPI windows and all split MPI
    communicators created by the manager.

    Returns
    -------
    None
    """
    self.free_shared_arrays_all()
    if hasattr(self, "_tree_grp_root_comm") and self._tree_grp_root_comm:
        self._tree_grp_root_comm.Free()
    if hasattr(self, "_tree_grp_comm") and self._tree_grp_comm:
        self._tree_grp_comm.Free()
    if hasattr(self, "_node_root_comm") and self._node_root_comm:
        self._node_root_comm.Free()
    if hasattr(self, "_node_comm") and self._node_comm:
        self._node_comm.Free()