nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/distributed/distribution.py

@@ -0,0 +1,770 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+import copy
+import math
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import TypeAlias, cast
+
+import nvmath.distributed as dist
+from nvmath.bindings import (
+    cublasMp,  # type: ignore
+    cufftMp,  # type: ignore
+)
+
+__all__ = ["ProcessGrid", "Distribution", "Slab", "Box", "BlockCyclic", "BlockNonCyclic"]
+
+
+class ProcessGrid:
+    """
+    N-dimensional grid of processes used by some distributions like the PBLAS block-cyclic
+    distribution.
+
+    Example 2D process grid for 4 processes, with processes arranged in column-major order::
+
+        ---------
+        | 0 | 2 |
+        ---------
+        | 1 | 3 |
+        ---------
+    """
+
+    Layout: TypeAlias = cublasMp.GridLayout
+
+    def __init__(
+        self,
+        *,
+        shape: Sequence[int] | None = None,
+        layout: ProcessGrid.Layout | None = None,
+        process_array=None,
+    ):
+        """
+        Create a new ProcessGrid object.
+
+        Args:
+            shape: Shape of the process grid.
+
+            layout: Layout of the process grid (column-major or row-major). This is optional
+                for 1D grid or when a custom grid is provided.
+
+            process_array: optional ndarray specifying custom arrangement of processes.
+        """
+        self._nranks = _get_process_group().nranks
+
+        if process_array is None:
+            if shape is None:
+                raise ValueError("shape must be provided when process_array=None")
+            self._shape = tuple(shape)
+            if layout is None:
+                if self._is_1d_distribution():
+                    layout = ProcessGrid.Layout.ROW_MAJOR  # layout doesn't matter in this case.
+                else:
+                    raise ValueError("layout must be provided when process_array=None and partitioning on multiple dimensions")
+            if not isinstance(layout, ProcessGrid.Layout):
+                raise TypeError(f"layout must be of type ProcessGrid.Layout, got {layout}")
+            self._layout = layout
+            self._process_array = None
+        else:
+            self._shape = tuple(process_array.shape)
+            if shape is not None and self._shape != tuple(shape):
+                raise ValueError(f"shape {shape} and process_array.shape ({process_array.shape}) don't match")
+            # TODO: Can set layout to COL_MAJOR or ROW_MAJOR automatically if the
+            # process_array matches.
+            if layout is not None:
+                raise NotImplementedError
+            self._layout = None
+            self._process_array = process_array
+
+        if math.prod(self._shape) != self._nranks:
+            raise ValueError(
+                f"Number of grid elements ({math.prod(self._shape)}) must equal the number of processes ({self._nranks})"
+            )
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Shape of process grid."""
+        return self._shape
+
+    @property
+    def layout(self) -> ProcessGrid.Layout | None:
+        """Layout of process grid if row-major or column-major, otherwise None."""
+        return self._layout
+
+    @property
+    def process_array(self):
+        return self._process_array
+
+    def __str__(self):
+        return f"ProcessGrid(shape={self._shape}, layout={self._layout.name}, process_array={self._process_array})"
+
+    def __hash__(self):
+        # NOTE: The layout isn't considered for the hash, to allow process grids
+        # partitioned on a single dimension with the same shape but different layout
+        # to be the same dictionary key.
+        return hash(self._shape)
+
+    def __eq__(self, other):
+        if self._process_array is None and other._process_array is None:
+            if self._shape == other._shape:
+                return self._is_1d_distribution() or self._layout == other._layout
+            return False
+        raise NotImplementedError
+
+    def _is_1d_distribution(self) -> bool:
+        """True if process grid partitions on a single dimension."""
+        return self._nranks in self._shape
+
+    def _is_row_wise(self) -> bool:
+        """True if 2D process grid partitioned only on rows."""
+        return self._shape == (self._nranks, 1)
+
+    def _is_col_wise(self) -> bool:
+        """True if 2D process grid partitioned only on columns."""
+        return self._shape == (1, self._nranks)
+
+
+class BindDistributionError(Exception):
+    pass
+
+
+class ConvertDistributionError(Exception):
+    """Errors converting a distribution instance to another distribution type"""
+
+    pass
+
+
+class Distribution(ABC):
+    """Specifies how a tensor is distributed across processes."""
+
+    def __init__(self):
+        self._bound = False
+
+    @property
+    @abstractmethod
+    def ndim(self) -> int | None:
+        """The number of dimensions of a distributed tensor for which this distribution
+        applies; None if it doesn't apply to any specific number of dimensions."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def shape(self, rank: int, global_shape: Sequence[int] | None = None) -> tuple[int, ...]:
+        """Get the local shape of data on the given rank according to this distribution.
+
+        Args:
+            rank: the process rank for which to calculate the local shape.
+
+            global_shape: Global shape of data. Required if the distribution is
+                not bound to a global shape, otherwise not required.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def to(
+        self,
+        cls: type[Distribution],
+        /,
+        *,
+        ndim: int | None = None,
+        copy: bool = False,
+    ) -> Distribution:
+        """Convert this distribution object to an equivalent distribution of the given type.
+
+        Args:
+            cls: the target distribution type.
+
+            ndim: dimensionality of the target distribution. Must be compatible with the
+                dimensionality of the source distribution. This may be required if the
+                source distribution doesn't have associated dimensionality.
+
+            copy: Returns a copy if the source and target type are the same.
+
+        Raises:
+            ConvertDistributionError: if the conversion is not possible.
+        """
+        raise NotImplementedError
+
+    def _to_checks(self, cls: type[Distribution], ndim: int | None) -> None:
+        if cls is Distribution or not issubclass(cls, Distribution):
+            raise ValueError(f"{cls} is not a valid distribution")
+        if ndim is not None and self.ndim is not None and ndim != self.ndim:
+            raise ValueError(f"ndim argument ({ndim}) doesn't match this distribution's dimensionality ({self.ndim})")
+
+    @abstractmethod
+    def _bind(
+        self,
+        global_shape: Sequence[int],
+        *,
+        shape: Sequence[int] | None = None,
+    ) -> Distribution:
+        """Binds this distribution object to a global shape, which determines how a
+        distributed tensor with that shape must be partitioned among processes (the local
+        shape on each process). You can also provide the local shape on this process to
+        check if it fits the distribution (the function will raise an exception if not).
+        **The exception may be raised on some ranks but not others (it's up to the caller
+        to handle this)**.
+
+        Args:
+            global_shape: global shape of the data.
+
+            shape: shape of the data on this process.
+
+        Returns:
+            self
+
+        Raises:
+            BindDistributionError: if distribution is already bound or local shape doesn't
+            fit the distribution.
+        """
+        raise NotImplementedError
+
+    def _binding_str(self):
+        return f"[bound: global_shape={self._data_global_shape}, shape={self._data_shape}]" if self._bound else ""
+
+    def copy(self) -> Distribution:
+        """This is a common implementation for those distributions that only require
+        a shallow copy."""
+        return copy.copy(self)
+
+    def _local_shape_checks(self, rank: int, nranks: int, global_shape: Sequence[int] | None = None):
+        if not isinstance(rank, int):
+            raise ValueError(f"rank must be an integer, got {rank}")
+
+        if rank < 0 or rank > nranks - 1:
+            raise ValueError(f"This is not a valid process rank: got rank={rank} with nranks={nranks}")
+
+        if global_shape is None and not self._bound:
+            raise RuntimeError("This distribution is unbound: please specify a global shape")
+
+        if global_shape is not None and self._bound and tuple(global_shape) != self._data_global_shape:  # type: ignore
+            raise ValueError(
+                "This distribution is already bound to a different global shape: provided "
+                f"{global_shape}, bound to {self._data_global_shape}"  # type: ignore
+            )
+
+
+class Slab(Distribution):
+    """
+    Slab distribution
+
+    Data is partitioned across processes on a single axis, such that:
+
+    - The shape of the slab on the first s_p % P processes is
+      (s_0, ..., s_p // P + 1, ..., s_{n-1})
+    - The shape of the slab on the remaining processes is (s_0, ..., s_p // P, ..., s_{n-1})
+    - Process 0 owns the first slab according to the global index order, process 1 owns
+      the second slab and so on.
+
+    where:
+
+    - s_i is the size of dimension i of the global array
+    - p is the partition dimension
+    - n is the number of dimensions of the array
+    - P is the number of processes
+    """
+
+    X: Slab
+    """Slab distribution on axis 0."""
+
+    Y: Slab
+    """Slab distribution on axis 1."""
+
+    Z: Slab
+    """Slab distribution on axis 2."""
+
+    def __init__(
+        self,
+        partition_dim: int,
+        ndim: int | None = None,
+    ):
+        super().__init__()
+        if not isinstance(partition_dim, int) or partition_dim < 0:
+            raise ValueError(f"partition_dim must be integer >= 0, got {partition_dim}")
+
+        if ndim is not None:
+            if not isinstance(ndim, int) or ndim < 1:
+                raise ValueError(f"ndim must be integer >= 1, got {ndim}")
+            if partition_dim >= ndim:
+                raise ValueError("partition_dim must be < ndim")
+
+        self._partition_dim = partition_dim
+        self._ndim = ndim
+
+    def __eq__(self, other):
+        if not isinstance(other, Slab):
+            return False
+        if self._partition_dim != other._partition_dim:
+            return False
+        if self._ndim is None or other._ndim is None:
+            return True
+        return self._ndim == other._ndim
+
+    def __hash__(self):
+        return self._partition_dim
+
+    def __str__(self):
+        return f"Slab(partition_dim={self._partition_dim}, ndim={self._ndim})" + self._binding_str()
+
+    @property
+    def name(self) -> str:
+        match self._partition_dim:
+            case 0:
+                return "Slab.X"
+            case 1:
+                return "Slab.Y"
+            case 2:
+                return "Slab.Z"
+            case _:
+                return "Slab"
+
+    @property
+    def partition_dim(self) -> int:
+        """Slab partition dimension"""
+        return self._partition_dim
+
+    @property
+    def ndim(self) -> int | None:
+        return self._ndim
+
+    def shape(self, rank: int, global_shape: Sequence[int] | None = None) -> tuple[int, ...]:
+        n = _get_process_group().nranks
+        self._local_shape_checks(rank, n, global_shape)
+        if global_shape is None:
+            global_shape = self._data_global_shape
+
+        S = global_shape[self._partition_dim]
+        partition_dim_local_size = S // n + bool(rank < S % n)
+        slab_shape = list(global_shape)
+        slab_shape[self._partition_dim] = partition_dim_local_size
+        return tuple(slab_shape)
+
+    @property
+    def _cufftmp_value(self):
+        if self._partition_dim not in (0, 1):
+            raise TypeError(f"Unsupported distribution {self} for cuFFTMp: partition dimension must be X or Y")
+        return cufftMp.XtSubFormat.FORMAT_INPLACE if self._partition_dim == 0 else cufftMp.XtSubFormat.FORMAT_INPLACE_SHUFFLED
+
+    def _bind(self, global_shape, *, shape=None) -> Slab:
+        if self._bound:
+            raise BindDistributionError(f"{self} is already bound")
+
+        if self._ndim is not None and len(global_shape) != self._ndim:
+            raise ValueError(f"The given shape doesn't have the same dimensionality as this {self} distribution")
+
+        if self._partition_dim >= len(global_shape):
+            raise ValueError("partition_dim must be < ndim")
+
+        rank = _get_process_group().rank
+        slab_shape = self.shape(rank, global_shape)
+
+        if shape is not None and tuple(shape) != slab_shape:
+            raise BindDistributionError(
+                f"The given shapes (global_shape={global_shape}, shape={shape}) don't fit distribution {str(self)}"
+            )
+
+        self._ndim = len(global_shape)
+        self._data_global_shape = tuple(global_shape)
+        self._data_shape = slab_shape
+        self._bound = True
+        return self
+
+    def to(self, cls, /, *, ndim=None, copy=False):
+        super()._to_checks(cls, ndim)
+        nranks = _get_process_group().nranks
+
+        if ndim is None:
+            ndim = self._ndim
+        elif self._partition_dim >= ndim:
+            raise ValueError(f"ndim ({ndim}) must be greater than the partition dimension ({self._partition_dim})")
+
+        if cls is Slab:
+            if copy or (ndim is not None and self._ndim is None):
+                d = cast(Slab, self.copy())
+                d._ndim = ndim
+                return d
+            return self
+        elif issubclass(cls, BlockCyclic):
+            if ndim is None:
+                raise ConvertDistributionError("Can't convert Slab distribution to BlockCyclic: unknown dimensionality")
+
+            process_grid_shape = tuple(1 if x != self._partition_dim else nranks for x in range(ndim))
+            # layout doesn't matter when partitioning on a single axis.
+            process_grid = ProcessGrid(shape=process_grid_shape, layout=ProcessGrid.Layout.ROW_MAJOR)
+            if not self._bound:
+                return BlockNonCyclic(process_grid)
+            else:
+                b = BlockNonCyclic(process_grid)
+                return b._bind(self._data_global_shape, shape=self._data_shape)
+        elif cls is Box:
+            raise NotImplementedError
+
+
+# Define alternate forms for the user to specify Slab on X, Y or Z.
+# NOTE: dimensionality is left unspecified when using these (but will be set when the
+# distribution is bound to data).
+Slab.X = Slab(0)
+Slab.Y = Slab(1)
+Slab.Z = Slab(2)
+
+
+class Box(Distribution):
+    """Box distribution"""
+
+    def __init__(
+        self,
+        lower: Sequence[int],
+        upper: Sequence[int],
+    ):
+        super().__init__()
+        if len(lower) != len(upper):
+            raise ValueError("lower and upper coordinates must have the same dimensionality")
+        for coords in (lower, upper):
+            if not all(isinstance(x, int) for x in coords):
+                raise ValueError("lower and upper coordinates must be integer")
+        if not all(upper[i] > lower[i] for i in range(len(upper))):
+            raise ValueError(
+                f"The upper coordinates must be larger than the lower coordinates, but got lower={lower} upper={upper}"
+            )
+        self._lower = tuple(lower)
+        self._upper = tuple(upper)
+
+    @property
+    def lower(self) -> tuple[int, ...]:
+        """Box lower coordinates"""
+        return self._lower
+
+    @property
+    def upper(self) -> tuple[int, ...]:
+        """Box upper coordinates"""
+        return self._upper
+
+    @property
+    def ndim(self) -> int:
+        return len(self._lower)
+
+    def shape(self, rank: int, global_shape: Sequence[int] | None = None) -> tuple[int, ...]:
+        process_group = _get_process_group()
+        if rank != process_group.rank:
+            raise RuntimeError("Can't calculate local shape of peer process with Box distribution")
+        nranks = process_group.nranks
+        self._local_shape_checks(rank, nranks, global_shape)
+        return tuple(self._upper[i] - self._lower[i] for i in range(self.ndim))
+
+    def __str__(self):
+        return f"Box(lower={self._lower}, upper={self._upper})" + self._binding_str()
+
+    def __eq__(self, other):
+        if not isinstance(other, Box):
+            return False
+        return self._lower == other._lower and self._upper == other._upper
+
+    def __hash__(self):
+        return hash((self._lower, self._upper))
+
+    def __iter__(self):
+        # To allow unpacking
+        yield self._lower
+        yield self._upper
+
+    def __getitem__(self, index):
+        if index not in (0, 1):
+            return IndexError(f"Index must be 0 or 1, got {index}")
+        return self._lower if index == 0 else self._upper
+
+    def _bind(self, global_shape, *, shape=None) -> Distribution:
+        if self._bound:
+            raise BindDistributionError(f"{self} is already bound")
+        my_shape = tuple(self._upper[i] - self._lower[i] for i in range(self.ndim))
+        if shape is not None and tuple(shape) != my_shape:
+            raise BindDistributionError(
+                f"The given shapes don't fit this Box distribution: {global_shape} "
+                f"and {shape}, lower={self._lower} upper={self._upper}"
+            )
+        self._data_global_shape = tuple(global_shape)
+        self._data_shape = my_shape
+        self._bound = True
+        return self
+
+    def to(self, cls, /, *, ndim=None, copy=False):
+        super()._to_checks(cls, ndim)
+        if cls is Box:
+            return self.copy() if copy else self
+        raise NotImplementedError
+
+
+class BlockCyclic(Distribution):
+    """Block-cyclic distribution"""
+
+    def __init__(
+        self,
+        process_grid: ProcessGrid,
+        block_sizes: Sequence[int],
+        *,
+        first_process: Sequence[int] | None = None,
+    ):
+        super().__init__()
+        if block_sizes is None or not all(isinstance(x, int) for x in block_sizes):
+            raise ValueError(f"Must provide a sequence of integer block sizes, got {block_sizes}")
+        if len(block_sizes) != len(process_grid.shape):
+            raise ValueError(
+                f"Number of block sizes ({len(block_sizes)}) doesn't match dimensionality ({len(process_grid.shape)})"
+            )
+        self._process_grid = process_grid
+        self._block_sizes = tuple(block_sizes)
+        if first_process is None:
+            self._first_process = (0,) * self.ndim
+        else:
+            if not all(isinstance(x, int) for x in first_process):
+                raise ValueError(f"first_process must be a sequence of integer coordinates, got {first_process}")
+            for i, x in enumerate(first_process):
+                if x < 0 or x >= process_grid.shape[i]:
+                    raise ValueError(
+                        f"first_process {first_process} is not a valid index into the process grid of "
+                        f"shape {process_grid.shape}"
+                    )
+            self._first_process = tuple(first_process)
+
+    @property
+    def process_grid(self) -> ProcessGrid:
+        """The process grid of this BlockCyclic distribution"""
+        return self._process_grid
+
+    @property
+    def ndim(self) -> int:
+        return len(self._process_grid.shape)
+
+    @property
+    def block_sizes(self) -> tuple[int, ...]:
+        """The block sizes of this BlockCyclic distribution"""
+        return self._block_sizes
+
+    @property
+    def first_process(self) -> tuple[int, ...]:
+        """Index in the process grid of the process who owns the first block of the
+        distributed tensor."""
+        return self._first_process
+
+    def __str__(self):
+        return (
+            f"{self.__class__.__name__}(process_grid={self._process_grid}, block_sizes={self._block_sizes})"
+            + self._binding_str()
+        )
+
+    def __eq__(self, other):
+        if not isinstance(other, BlockCyclic):
+            return False
+        return (
+            self._process_grid == other._process_grid
+            and self._block_sizes == other._block_sizes
+            and self._first_process == other._first_process
+        )
+
+    def _is_1d_distribution(self) -> bool:
+        """True if process grid partitions on a single dimension."""
+        return self._process_grid._is_1d_distribution()
+
+    def _is_row_wise(self) -> bool:
+        """True if 2D process grid partitioned only on rows."""
+        return self._process_grid._is_row_wise()
+
+    def _is_col_wise(self) -> bool:
+        """True if 2D process grid partitioned only on columns."""
+        return self._process_grid._is_col_wise()
+
+    def shape(self, rank: int, global_shape: Sequence[int] | None = None) -> tuple[int, ...]:
+        self._local_shape_checks(rank, self._process_grid._nranks, global_shape)
+        if global_shape is None:
+            global_shape = self._data_global_shape
+        return self._calc_local_shape(rank, self._block_sizes, global_shape)
+
+    def _calc_local_shape(self, rank, block_sizes, global_shape):
+        nprow, npcol = self._process_grid._shape
+        layout = self._process_grid._layout
+        if layout is not None:
+            myprow = rank % nprow if layout == ProcessGrid.Layout.COL_MAJOR else rank // npcol
+            mypcol = rank // nprow if layout == ProcessGrid.Layout.COL_MAJOR else rank % npcol
+            index = (myprow, mypcol)
+        else:
+            raise NotImplementedError
+        nrows = cublasMp.numroc(global_shape[0], block_sizes[0], index[0], self._first_process[0], nprow)
+        ncols = cublasMp.numroc(global_shape[1], block_sizes[1], index[1], self._first_process[1], npcol)
+        return (nrows, ncols)
+
+    def _bind(self, global_shape, *, shape=None) -> Distribution:
+        if self._bound:
+            raise BindDistributionError(f"{self} is already bound")
+
+        if self.ndim != 2:
+            raise NotImplementedError
+
+        if len(global_shape) != self.ndim:
+            raise ValueError(
+                f"Dimensionality of shapes ({len(global_shape)}) doesn't match dimensionality "
+                f"of this distribution ({self.ndim})"
+            )
+
+        rank = _get_process_group().rank
+        nrows, ncols = self._calc_local_shape(rank, self._block_sizes, global_shape)
+        if shape is not None and tuple(shape) != (nrows, ncols):
+            raise BindDistributionError(
+                f"The local shape {shape} on process {rank} is not the expected one based "
+                f"on the global shape {global_shape}, process grid {self._process_grid} and "
+                f"block sizes {self._block_sizes}: expected shape is {(nrows, ncols)}"
+            )
+
+        self._data_global_shape = global_shape
+        self._data_shape = (nrows, ncols)
+        self._bound = True
+        return self
+
+    def to(self, cls, /, *, ndim=None, copy=False):
+        super()._to_checks(cls, ndim)
+        nranks = _get_process_group().nranks
+        if issubclass(cls, BlockCyclic):
+            return self.copy() if copy else self
+        elif cls is Slab:
+            if not self._bound:
+                # Without binding, it's a stretch to assume that this is compatible with
+                # Slab (the "cyclic" nature means that it's much more likely that it isn't).
+                raise ConvertDistributionError(
+                    "Unbound BlockCyclic distribution can't be converted to Slab. "
+                    "Consider using BlockNonCyclic if there is no cyclic distribution of blocks."
+                )
+
+            if not self._is_1d_distribution():
+                raise ConvertDistributionError(
+                    "Can't convert this block distribution to Slab: partitioning must be on a single dimension"
+                )
+
+            # Data must be divisible on partition_dim and the block size must correspond
+            # to the slab size.
+            partition_dim = self._process_grid.shape.index(nranks)
+
+            if self._data_global_shape[partition_dim] % nranks != 0:
+                raise ConvertDistributionError(
+                    "Can't convert this distribution to Slab: data doesn't divide evenly on partition dimension."
+                    f" Global shape is {self._data_global_shape}, partition dimension is "
+                    f"{partition_dim} and number of processes is {nranks}."
+                )
+
+            if self._block_sizes[partition_dim] != self._data_global_shape[partition_dim] // nranks:
+                raise ConvertDistributionError(
+                    "This distribution can't be converted to Slab, because the block size in the "
+                    f"partition dimension {partition_dim} is not a factor of the global extent in "
+                    f"that dimension {self._data_global_shape[partition_dim]}"
+                )
+            d = Slab(partition_dim, ndim=self.ndim)
+            return d._bind(self._data_global_shape, self._data_shape)
+        elif cls is Box:
+            raise NotImplementedError
+
+
+class BlockNonCyclic(BlockCyclic):
+    """Block distribution without cycles"""
+
+    def __init__(
+        self,
+        process_grid: ProcessGrid,
+        *,
+        first_process: Sequence[int] | None = None,
+    ):
+        super().__init__(process_grid, (-1,) * len(process_grid.shape), first_process=first_process)
+
+    def __eq__(self, other):
+        if isinstance(other, BlockNonCyclic) and (not other._bound or not self._bound):
+            # Don't compare block sizes since one of them is not bound.
+            return self._process_grid == other._process_grid and self._first_process == other._first_process
+        return super().__eq__(other)
+
+    def shape(self, rank: int, global_shape: Sequence[int] | None = None) -> tuple[int, ...]:
+        self._local_shape_checks(rank, self._process_grid._nranks, global_shape)
+        if global_shape is None:
+            global_shape = self._data_global_shape
+        if self._bound:
+            block_sizes = self._block_sizes
+        else:
+            block_sizes = self._infer_block_sizes(global_shape)
+        return self._calc_local_shape(rank, block_sizes, global_shape)
+
+    def _infer_block_sizes(self, global_shape):
+        if all(x % self._process_grid._shape[i] == 0 for i, x in enumerate(global_shape)):
+            block_sizes = tuple(x // self._process_grid._shape[i] for i, x in enumerate(global_shape))
+        else:
+            # The logic to bind this global shape to this distribution isn't implemented yet
+            # (which doesn't necessarily mean that it isn't possible to fit the data to this
+            # distribution).
+            raise NotImplementedError(
+                "BlockNonCyclic is currently only implemented for uniform partition sizes. "
+                "Use BlockCyclic with explicit block sizes instead."
+            )
+        return block_sizes
+
+    def _bind(self, global_shape, *, shape=None) -> Distribution:
+        if self._bound:
+            raise BindDistributionError(f"{self} is already bound")
+
+        assert all(x == -1 for x in self._block_sizes)
+
+        if len(global_shape) != self.ndim:
+            raise ValueError(
+                f"Dimensionality of shapes ({len(global_shape)}) doesn't match dimensionality "
+                f"of this distribution ({self.ndim})"
+            )
+
+        # Now we assign block sizes to this distribution based on the shape of the data.
+        # NOTE: For dimensions that aren't partitioned, there are multiple blocks sizes that
+        # are valid and still fit the Block[Non]Cyclic model. This is because, in addition
+        # to block size being the full length L of the dimension, any L//N is also a valid
+        # block size (so a single block in that dimension is equivalent to N contiguous
+        # blocks in that dimension). The importance of this is that cuBLASMp actually
+        # requires configurations with these block sizes, since block sizes have to match
+        # across matrices A, B and C but they might be distributed differently.
+        # In other words, for distributed matmul with BlockNonCyclic, block sizes have to be
+        # inferred jointly with matrices A, B and C, but that is outside the scope of this
+        # method.
+
+        block_sizes = self._infer_block_sizes(global_shape)
+        if shape is not None and block_sizes != tuple(shape):
+            raise BindDistributionError("Data doesn't fit BlockNonCyclic distribution")
+
+        self._block_sizes = block_sizes
+        return super()._bind(global_shape, shape=shape)
+
+    def to(self, cls, /, *, ndim=None, copy=False):
+        super()._to_checks(cls, ndim)
+        nranks = _get_process_group().nranks
+        if issubclass(cls, BlockCyclic):
+            return self.copy() if copy else self
+        elif cls is Slab:
+            if not self._is_1d_distribution():
+                raise ConvertDistributionError(
+                    "Can't convert this block distribution to Slab: partitioning must be on a single dimension"
+                )
+
+            partition_dim = self._process_grid.shape.index(nranks)
+            d = Slab(partition_dim, ndim=self.ndim)
+            # For bound=False, we can allow the conversion and let Slab._bind() catch
+            # any potential errors later.
+            if self._bound:
+                if self._data_global_shape[partition_dim] % nranks != 0:
+                    raise ConvertDistributionError(
+                        "Can't convert this distribution to Slab: data doesn't divide evenly on partition dimension"
+                    )
+                d._bind(self._data_global_shape, shape=self._data_shape)
+            return d
+        elif cls is Box:
+            raise NotImplementedError
+
+
+def _get_process_group():
+    distributed_ctx = dist.get_context()
+    if distributed_ctx is None:
+        raise RuntimeError(
+            "nvmath.distributed has not been initialized. Refer to "
+            "https://docs.nvidia.com/cuda/nvmath-python/latest/distributed-apis/runtime.html"
+            " for more information."
+        )
+    return distributed_ctx.process_group