nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/distributed/reshape/reshape.py

@@ -0,0 +1,1256 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["Reshape", "reshape"]
+
+import logging
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Final, Literal, cast
+
+import numpy as np
+
+import nvmath.distributed
+from nvmath.bindings import (
+    cufftMp,  # type: ignore
+    nvshmem,  # type: ignore
+)
+from nvmath.distributed._internal import tensor_wrapper
+from nvmath.distributed._internal.nvshmem import NvshmemMemoryManager
+from nvmath.distributed._internal.tensor_ifc import DistributedTensor
+from nvmath.distributed.distribution import Box
+from nvmath.internal import formatters, utils
+from nvmath.internal.package_wrapper import AnyStream, StreamHolder
+
+from ._configuration import ReshapeOptions
+
+
+@dataclass
+class TensorLayout:
+    """An internal data class for capturing the tensor layout."""
+
+    shape: Sequence[int]
+    strides: Sequence[int]
+
+
+@dataclass
+class _ProblemSpec:
+    """This is used in a custom reduction to check that the Reshape problem specification
+    is consistent across processes, and to infer global information (e.g shape and memory
+    layout)."""
+
+    @dataclass
+    class Options:
+        """
+        This is used for _ProblemSpec instead of ._configuration.ReshapeOptions
+        because it's going to be serialized as part of the custom reduction of the
+        _ProblemSpec, and we want to control which fields are included (for example
+        we don't need the logger).
+        """
+
+        def __init__(self, options: ReshapeOptions):
+            self.blocking = options.blocking
+
+        blocking: Literal[True, "auto"]
+
+    shape: tuple[int, ...]  # operand shape
+    is_F: bool  # Is Fortran memory layout
+    is_C: bool  # Is C memory layout
+    operand_dtype: str  # str because TensorHolder.dtype returns str
+    package: Literal["numpy", "cupy", "torch"]  # operand package
+    memory_space: Literal["cuda", "cpu"]  # operand memory space
+    boxes: list[Box]  # Reshape input and output box
+    options: Options  # Reshape options
+
+    # Global number of elements in the operand (calculated as part of the reduction).
+    global_size: int = 0
+    # is_leaf=True means that this is the _ProblemSpec of a process before reducing
+    # with that of another process.
+    is_leaf: bool = True
+
+
+SHARED_RESHAPE_DOCUMENTATION = utils.COMMON_SHARED_DOC_MAP.copy()
+SHARED_RESHAPE_DOCUMENTATION.update(
+    {
+        "operand": SHARED_RESHAPE_DOCUMENTATION["operand"],
+        #
+        "operand_admonitions": """
+            .. important::
+                GPU operands must be on the symmetric heap (for example, allocated with
+                ``nvmath.distributed.allocate_symmetric_memory()``).
+""",
+        #
+        "options": """\
+Specify options for the Reshape as a :class:`ReshapeOptions` object. Alternatively, a `dict` containing the parameters for the
+``ReshapeOptions`` constructor can also be provided. If not specified, the value will be set to the default-constructed
+``ReshapeOptions`` object.""".replace("\n", " "),
+        #
+        "input_box": """\
+The box specifying the distribution of the input operand across processes, where each process
+specifies which portion of the global array it holds. A box is a pair of coordinates specifying
+the lower and upper extent for each dimension.""".replace("\n", " "),
+        #
+        "output_box": """\
+The box specifying the distribution of the result across processes, where each process
+specifies which portion of the global array it will hold after reshaping. A box is a pair
+of coordinates specifying the lower and upper extent for each dimension.""".replace("\n", " "),
+        #
+        "sync_symmetric_memory": """\
+Indicates whether to issue a symmetric memory synchronization operation on the execute stream
+before the reshape operation. Note that before the Reshape starts executing, it is
+required that the source operand be ready on all processes. A symmetric memory synchronization
+ensures completion and visibility by all processes of previously issued local stores to
+symmetric memory. Advanced users who choose to manage the synchronization on their own using
+the appropriate NVSHMEM API, or who know that GPUs are already synchronized on the source
+operand, can set this to False.""".replace("\n", " "),
+        #
+        "function_signature": """\
+operand,
+input_box: Box,
+output_box: Box,
+sync_symmetric_memory: bool = True,
+options: ReshapeOptions | None = None,
+stream: AnyStream | None = None
+""".replace("\n", " "),
+    }
+)
+
+
+def _calculate_strides(shape, axis_order):
+    """
+    Calculate the strides for the provided shape and axis order.
+    """
+    strides = [None] * len(shape)
+
+    stride = 1
+    for axis in axis_order:
+        strides[axis] = stride
+        stride *= shape[axis]
+
+    return strides
+
+
+def _copy_operand_perhaps(
+    operand: DistributedTensor,
+    stream_holder,
+    execution_space,
+    memory_space,
+    device_id: int | Literal["cpu"],
+):
+    if execution_space == memory_space:
+        return operand, None
+    else:
+        # Copy the `operand` to memory that matches the exec space.
+        # Currently, reshape only runs on GPU.
+        assert execution_space == "cuda"
+        exec_space_copy = operand.to(device_id, stream_holder, symmetric_memory=True)
+        return exec_space_copy, operand
+
+
+def _problem_spec_reducer(p1: _ProblemSpec, p2: _ProblemSpec):
+    try:
+        if isinstance(p1, Exception):
+            return p1  # propagate exception
+
+        if isinstance(p2, Exception):
+            return p2  # propagate exception
+
+        if len(p1.shape) != len(p2.shape):
+            return ValueError("The number of dimensions of the input operand is inconsistent across processes")
+
+        if len(p1.shape) <= 0 or len(p1.shape) > 3:
+            return ValueError(
+                "Reshape currently only supports 1-D, 2-D and 3-D tensors."
+                f" The number of dimensions of the operand is {len(p1.shape)}."
+            )
+
+        if p1.operand_dtype != p2.operand_dtype:
+            return ValueError("The operand dtype is inconsistent across processes")
+
+        if p1.package != p2.package:
+            return ValueError("operand doesn't belong to the same package on all processes")
+
+        if p1.memory_space != p2.memory_space:
+            return ValueError('operand is not on the same memory space ("cpu", "cuda") on all processes')
+
+        if p1.options != p2.options:
+            return ValueError(f"options are inconsistent across processes: {p1.options} != {p2.options}")
+
+        # Determine the memory layout shared by all processes. Note that it's possible for
+        # a process to have both C and F layout (e.g. (2, 1) input shape).
+        p1.is_C &= p2.is_C
+        p1.is_F &= p2.is_F
+        if not p1.is_F and not p1.is_C:
+            return ValueError("The input memory layout is not C or Fortran, or is inconsistent across processes")
+
+        for p_spec in (p1, p2):
+            if p_spec.is_leaf:
+                for box in p_spec.boxes:
+                    if not isinstance(box, Box):
+                        return ValueError(f"{box} is not a Box distribution")
+                    if box.ndim != len(p_spec.shape):
+                        return ValueError(
+                            f"The dimensionality of {box} doesn't match the dimensionality of "
+                            f"the reshape operand ({len(p_spec.shape)})"
+                        )
+
+        if p1 is not p2:  # with nranks=1 p1 is p2
+            p1.global_size += p2.global_size
+
+        def reduce_boxes(box1, box2):
+            """This function returns the smallest box that encompasses `box1` and `box2`"""
+            lower = np.minimum(np.array(box1.lower), np.array(box2.lower)).tolist()
+            upper = np.maximum(np.array(box1.upper), np.array(box2.upper)).tolist()
+            return Box(lower, upper)
+
+        # Merge the boxes to get the global operand shape. Note that this is applied
+        # progressively throughout the reduction, starting with the local boxes.
+        p1.boxes = [reduce_boxes(p1.boxes[0], p2.boxes[0]), reduce_boxes(p1.boxes[1], p2.boxes[1])]
+
+    except Exception as e:
+        return e
+    p1.is_leaf = False
+    return p1
+
+
+class InvalidReshapeState(Exception):
+    pass
+
+
+@utils.docstring_decorator(SHARED_RESHAPE_DOCUMENTATION, skip_missing=False)
+class Reshape:
+    """
+    Create a stateful object that encapsulates the specified distributed Reshape
+    and required resources. This object ensures the validity of resources during use and
+    releases them when they are no longer needed to prevent misuse.
+
+    This object encompasses all functionalities of the function-form API :func:`reshape`,
+    which is a convenience wrapper around it. The stateful object also allows for the
+    amortization of preparatory costs when the same Reshape operation is to be performed
+    on multiple operands with the same problem specification (see :meth:`reset_operand`
+    for more details).
+
+    Using the stateful object typically involves the following steps:
+
+    1. **Problem Specification**: Initialize the object with a defined operation and
+       options.
+    2. **Preparation**: Use :meth:`plan` to determine the best distributed algorithmic
+       implementation for this specific Reshape operation.
+    3. **Execution**: Perform the Reshape with :meth:`execute`.
+    4. **Resource Management**: Ensure all resources are released either by explicitly
+       calling :meth:`free` or by managing the stateful object within a context manager.
+
+    Detailed information on each step described above can be obtained by passing in a
+    :class:`logging.Logger` object to :class:`ReshapeOptions` or by setting the appropriate
+    options in the root logger object, which is used by default:
+
+        >>> import logging
+        >>> logging.basicConfig(
+        ...     level=logging.INFO,
+        ...     format="%(asctime)s %(levelname)-8s %(message)s",
+        ...     datefmt="%m-%d %H:%M:%S",
+        ... )
+
+    Args:
+        operand: {operand}
+            {operand_admonitions}
+
+        input_box: {input_box}
+
+        output_box: {output_box}
+
+        options: {options}
+
+        stream: {stream}
+
+    .. seealso::
+        :meth:`plan`, :meth:`reset_operand`, :meth:`execute`,
+        :meth:`release_operand`
+
+    Examples:
+
+        >>> import cupy as cp
+        >>> import nvmath.distributed
+
+        Get process group used to initialize nvmath.distributed (for information on
+        initializing nvmath.distributed, you can refer to the documentation or to the
+        Reshape examples in `nvmath/examples/distributed/reshape
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/distributed/reshape>`_):
+
+        >>> process_group = nvmath.distributed.get_context().process_group
+
+        Let's create a 3D floating-point ndarray on GPU, distributed across a certain number
+        of processes, with each holding a portion of the ndarray. As an example, process 0
+        holds a 3D box of the global 3D array of shape (4, 4, 4).
+
+        >>> shape = 4, 4, 4
+
+        Reshape uses the NVSHMEM PGAS model, which requires GPU operands to be on the
+        symmetric heap:
+
+        >>> if process_group.rank == 0:
+        ...     a[:] = cp.random.rand(*shape)
+        ... else:
+        ...     a = ...  # each process holds a different section of the global array.
+        ... a = nvmath.distributed.allocate_symmetric_memory(shape, cp)
+
+        With Reshape, we will change how the ndarray is distributed, by having each
+        process specify the input and output section of the global array. For process 0,
+        let's assume that it holds the 3D box that goes from the lower corner given by
+        coordinates (0, 0, 0) to the upper corner (4, 4, 4).
+
+        NOTE: each process has its own input and output boxes which are different to those
+        of other processes, as each holds a different section of the global array.
+
+        >>> from nvmath.distributed.distribution import Box
+        >>> if process_group.rank == 0:
+        ...     input_lower = (0, 0, 0)
+        ...     input_upper = (4, 4, 4)
+        ...     input_box = Box(input_lower, input_upper)
+        ...     output_box = ...
+        ... else:
+        ...     input_box = ...  # the input box depends on the process.
+        ...     output_box = ...  # the output box depends on the process.
+
+        Create a Reshape object encapsulating the problem specification above:
+
+        >>> r = nvmath.distributed.reshape.Reshape(a, input_box, output_box)
+
+        Options can be provided above to control the behavior of the operation using the
+        `options` argument (see :class:`ReshapeOptions`).
+
+        Next, plan the Reshape:
+
+        >>> r.plan()
+
+        Now execute the Reshape, and obtain the result `b` as a CuPy ndarray. Reshape always
+        performs the distributed operation on the GPU.
+
+        >>> b = r.execute()
+
+        Finally, free the Reshape object's resources. To avoid this explicit call, it's
+        recommended to use the Reshape object as a context manager as shown below, if
+        possible.
+
+        >>> r.free()
+
+        Any symmetric memory that is owned by the user must be deleted explicitly (this is
+        a collective call and must be called by all processes):
+
+        >>> nvmath.distributed.free_symmetric_memory(a, b)
+
+        Note that all :class:`Reshape` methods execute on the current stream by default.
+        Alternatively, the `stream` argument can be used to run a method on a specified
+        stream.
+
+        Let's now look at the same problem with NumPy ndarrays on the CPU.
+
+        >>> import numpy as np
+        >>> a = np.random.rand(*shape)  # each process holds a different section
+
+        Create a Reshape object encapsulating the problem specification described earlier
+        and use it as a context manager.
+
+        >>> with nvmath.distributed.reshape.Reshape(a, input_box, output_box) as r:
+        ...     r.plan()
+        ...
+        ...     # Execute the Reshape to redistribute the ndarray.
+        ...     b = r.execute()
+
+        All the resources used by the object are released at the end of the block.
+
+        Reshape always executes on the GPU. In this case, because ``a`` resides in host
+        memory, the NumPy array is temporarily copied to device memory (on the symmetric
+        memory heap), re-distributed on the GPU, and the result is copied to host memory
+        as a NumPy array.
+
+        Further examples can be found in the `nvmath/examples/distributed/reshape
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/distributed/reshape>`_
+        directory.
+    """
+
+    def _free_internal_sheap(self, exception: Exception | None = None) -> bool:
+        # This is a fail-safe to free NVSHMEM internal memory in case of invalid
+        # state (Reshape constructor fails). Since we might call nvshmem_free here, we're
+        # assuming that all processes equally failed in the ctor, which might not be true,
+        # but if it weren't true they would end up in deadlock most likely anyway.
+        if (
+            hasattr(self, "memory_space")
+            and self.memory_space == "cpu"
+            and self.operand is not None
+            and self.operand.device == "cuda"
+        ):
+            with utils.device_ctx(self.device_id):
+                self.operand.free_symmetric()
+        return True
+
+    @utils.atomic(_free_internal_sheap, method=True)
+    def __init__(
+        self,
+        operand,
+        /,
+        input_box: Box,
+        output_box: Box,
+        *,
+        options: ReshapeOptions | None = None,
+        stream: AnyStream | None = None,
+    ):
+        distributed_ctx = nvmath.distributed.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."
+            )
+        if not distributed_ctx.nvshmem_available:
+            raise RuntimeError("nvmath.distributed wasn't initialized with NVSHMEM backend")
+        self.process_group = distributed_ctx.process_group
+        nranks = self.process_group.nranks
+
+        self.operand = operand = tensor_wrapper.wrap_operand(operand)
+        self.options = options = cast(
+            ReshapeOptions, utils.check_or_create_options(ReshapeOptions, options, "Distributed Reshape options")
+        )
+        self.package = operand.name
+
+        is_C = sorted(operand.strides, reverse=True) == list(operand.strides)
+        is_F = sorted(operand.strides) == list(operand.strides)
+
+        input_box = cast(Box, input_box.copy())
+        output_box = cast(Box, output_box.copy())
+
+        # Merge the problem specification across processes to make sure that there are no
+        # inconsistencies and to calculate the global shape. Importantly, this also does
+        # collective error checking of the Reshape input parameters, to ensure that all
+        # processes fail on error of any one process, thus preventing deadlock.
+        problem_spec = _ProblemSpec(
+            boxes=[input_box, output_box],
+            shape=tuple(operand.shape),
+            operand_dtype=operand.dtype,
+            options=_ProblemSpec.Options(options),
+            package=self.package,
+            memory_space=operand.device,
+            global_size=math.prod(operand.shape),
+            is_C=is_C,
+            is_F=is_F,
+        )
+        if nranks > 1:
+            problem_spec = self.process_group.allreduce_object(problem_spec, op=_problem_spec_reducer)
+        else:
+            # Ensure we error-check with one rank.
+            problem_spec = _problem_spec_reducer(problem_spec, problem_spec)
+        if isinstance(problem_spec, Exception):
+            # There is an error or inconsistency in the problem spec across processes.
+            # Note that since this comes from an allreduce, all processes will have
+            # received the same exception.
+            raise problem_spec
+
+        if problem_spec.is_C:
+            self.layout: Final = "C"
+        else:
+            assert problem_spec.is_F, "Internal Error."  # The reducer is supposed to have detected this
+            self.layout: Final = "F"  # type: ignore
+
+        self.operand_dim = len(operand.shape)
+
+        self.logger = options.logger if options.logger is not None else logging.getLogger()
+
+        # The merged (global) input and output box must be the same.
+        if problem_spec.boxes[0] != problem_spec.boxes[1]:
+            raise ValueError(
+                "The global box derived from the input and output boxes doesn't match: "
+                f"{problem_spec.boxes[0]} != {problem_spec.boxes[1]}"
+            )
+
+        # Infer the global shape from the global input box. Note that cuFFTMp does not
+        # disallow lower coordinates for the merged (global) boxes that are not 0.
+        lower, upper = problem_spec.boxes[0]
+        global_shape = tuple(int(upper[i] - lower[i]) for i in range(self.operand_dim))
+        self.logger.info(f"The global shape of the operand is {global_shape}.")
+
+        # This can't throw error since the local operand shape was already checked
+        # against the box shape in the ProblemSpec reducer.
+        input_box._bind(global_shape, shape=operand.shape)
+
+        # The global number of elements must be compatible with the global shape.
+        if problem_spec.global_size != math.prod(global_shape):
+            raise ValueError(f"The global number of elements is incompatible with the inferred global shape {global_shape}")
+
+        # Store the local input and output box.
+        self.input_box: Box = input_box
+        self.output_box: Box = output_box
+
+        self.operand_data_type = operand.dtype
+        # TODO: change to `operand.dtype.itemsize` once operand is StridedMemoryView.
+        itemsize = operand.itemsize
+        if itemsize not in (4, 8, 16):
+            raise ValueError(
+                f"Reshape only supports element sizes in (4, 8, 16) bytes. The operand's element size is {itemsize}"
+            )
+
+        self.logger.info(f"The Reshape data type is {self.operand_data_type}.")
+
+        # Infer execution and memory space.
+        execution_device_id: int = distributed_ctx.device_id
+        if operand.device_id != "cpu":  # exec space matches the mem space
+            self.memory_space = "cuda"
+            self.device_id = operand.device_id
+            assert operand.device_id == execution_device_id
+        else:  # we need to move inputs cpu -> gpu and outputs gpu -> cpu
+            self.memory_space = "cpu"
+            self.device_id = execution_device_id
+        self.execution_space = "cuda"
+        self.operand_device_id = operand.device_id
+        self.internal_op_package = self._internal_operand_package(self.package)
+        stream_holder: StreamHolder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+
+        if self.memory_space == "cuda" and not operand.is_symmetric_memory:
+            raise TypeError("Distributed reshape requires GPU operand to be on symmetric memory")
+
+        self.logger.info(
+            f"The input tensor's memory space is {self.memory_space}, and the execution space "
+            f"is {self.execution_space}, with device {self.device_id}."
+        )
+
+        self.logger.info(f"The specified stream for the Reshape ctor is {stream_holder and stream_holder.obj}")
+
+        # Copy the operand to execution_space's device if needed.
+        self.operand, self.operand_backup = _copy_operand_perhaps(
+            operand,
+            stream_holder,
+            self.execution_space,
+            self.memory_space,
+            self.device_id,
+        )
+        operand = self.operand
+        # Capture operand layout for consistency checks when resetting operands.
+        self.operand_layout = TensorLayout(shape=operand.shape, strides=operand.strides)
+
+        self.result_layout: TensorLayout | None = None
+        # We'll infer the result layout at plan time.
+        self.result_class: DistributedTensor = operand.__class__
+        self.result_data_type = operand.dtype
+
+        # Set blocking or non-blocking behavior.
+        self.blocking = self.options.blocking is True or self.memory_space == "cpu"
+        if self.blocking:
+            self.call_prologue = "This call is blocking and will return only after the operation is complete."
+        else:
+            self.call_prologue = (
+                "This call is non-blocking and will return immediately after the operation is launched on the device."
+            )
+
+        # Set memory allocator.
+        self.allocator = NvshmemMemoryManager(self.device_id, self.logger)
+
+        # Create handle.
+        with utils.device_ctx(self.device_id):
+            self.handle = cufftMp.create_reshape()
+
+        self.reshape_planned = False
+
+        # Workspace attributes.
+        self.workspace_ptr, self.workspace_size = None, None
+        self._workspace_allocated_here = False
+
+        # Attributes to establish stream ordering.
+        self.workspace_stream = None
+        self.last_compute_event = None
+
+        # Track whether the user has called release_operand(). This flag is
+        # checked in _check_valid_operand to prevent execution after the user
+        # has released their operand. It is cleared by reset_operand().
+        self._operand_released = False
+
+        self.valid_state = True
+        self.logger.info("The distributed Reshape operation has been created.")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.free()
+
+    def _check_valid_reshape(self, *args, **kwargs):
+        """
+        Check if Reshape object is alive and well.
+        """
+        if not self.valid_state:
+            raise InvalidReshapeState("The Reshape object cannot be used after resources are free'd")
+
+    def _free_plan_resources(self, exception: Exception | None = None) -> bool:
+        """
+        Free resources allocated in planning.
+        """
+        self.reshape_planned = False
+        return True
+
+    def _internal_operand_package(self, package_name):
+        return package_name if package_name != "numpy" else "cuda"
+
+    def _allocate_result_operand(self, exec_stream_holder, log_debug):
+        if log_debug:
+            self.logger.debug("Beginning output (empty) tensor creation...")
+            self.logger.debug(
+                f"The output tensor shape = {self.result_layout.shape} with strides = "
+                f"{self.result_layout.strides} and data type '{self.result_data_type}'."
+            )
+        result = utils.create_empty_tensor(
+            self.result_class,
+            self.result_layout.shape,
+            self.result_data_type,
+            self.device_id,
+            exec_stream_holder,
+            verify_strides=False,
+            strides=self.result_layout.strides,
+            symmetric_memory=True,
+            make_symmetric=True,
+            logger=self.logger,
+        )
+        if log_debug:
+            self.logger.debug("The output (empty) tensor has been created.")
+        return result
+
+    @utils.precondition(_check_valid_reshape)
+    @utils.atomic(_free_plan_resources, method=True)
+    def plan(self, stream: AnyStream | None = None):
+        """Plan the Reshape.
+
+        Args:
+            stream: {stream}
+        """
+        log_info = self.logger.isEnabledFor(logging.INFO)
+
+        if self.reshape_planned:
+            self.logger.debug("The Reshape has already been planned, and redoing the plan is not supported.")
+            return
+
+        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+        self.workspace_stream = stream_holder.obj
+
+        if log_info:
+            self.logger.info("Starting distributed Reshape planning...")
+
+        lower_input, upper_input = self.input_box
+        lower_output, upper_output = self.output_box
+
+        def calculate_reshape_params(lower, upper, operand, order: Literal["C", "F"]):
+            """Calculate shape and strides for input or output of Reshape."""
+            local_shape = tuple(upper[i] - lower[i] for i in range(self.operand_dim))
+            if operand is not None:
+                # Take the strides from the operand.
+                strides = operand.strides
+            elif order == "F":
+                # This is out=None with F layout.
+                strides = _calculate_strides(local_shape, range(self.operand_dim))
+            else:
+                # This is out=None with C layout.
+                strides = _calculate_strides(local_shape, reversed(range(self.operand_dim)))
+            return local_shape, strides
+
+        input_local_shape, strides_input = calculate_reshape_params(lower_input, upper_input, self.operand, self.layout)
+        output_local_shape, strides_output = calculate_reshape_params(lower_output, upper_output, None, self.layout)
+        self.result_layout = TensorLayout(shape=output_local_shape, strides=strides_output)
+
+        if self.operand_dim < 3:
+            # cufftMp.make_reshape only supports 3D, so we broadcast the local operands.
+            N = 3 - self.operand_dim
+            lower_input = tuple(lower_input) + (0,) * N
+            upper_input = tuple(upper_input) + (1,) * N
+            lower_output = tuple(lower_output) + (0,) * N
+            upper_output = tuple(upper_output) + (1,) * N
+
+            strides_input = self.operand.reshape(input_local_shape + (1,) * N, copy=False).strides
+            if self.layout == "F":
+                strides_output = _calculate_strides(output_local_shape + (1,) * N, (0, 1, 2))
+            else:
+                strides_output = _calculate_strides(output_local_shape + (1,) * N, (2, 1, 0))
+
+        # cuFFTMp only supports decreasing strides. To support increasing strides, we'll
+        # pass the required metadata to cuFFTMp with axes transposed (but won't touch
+        # the data).
+        if self.layout == "F":
+            lower_input = tuple(reversed(lower_input))
+            upper_input = tuple(reversed(upper_input))
+            strides_input = tuple(reversed(strides_input))
+            lower_output = tuple(reversed(lower_output))
+            upper_output = tuple(reversed(upper_output))
+            strides_output = tuple(reversed(strides_output))
+
+        with utils.host_call_ctx(timing=log_info) as elapsed, utils.device_ctx(self.device_id):
+            nullptr = 0
+            cufftMp.make_reshape(
+                self.handle,
+                # TODO: change to `operand.dtype.itemsize` once operand is StridedMemoryView
+                self.operand.itemsize,
+                3,
+                lower_input,
+                upper_input,
+                strides_input,
+                lower_output,
+                upper_output,
+                strides_output,
+                nullptr,
+                cufftMp.MpCommType.COMM_NONE,
+            )
+            self.workspace_size = cufftMp.get_reshape_size(self.handle)
+
+        self.reshape_planned = True
+
+        if log_info and elapsed.data is not None:
+            self.logger.info(f"The Reshape planning phase took {elapsed.data:.3f} ms to complete.")
+
+    @utils.precondition(_check_valid_reshape)
+    def reset_operand(self, operand, *, stream: AnyStream | None = None):
+        """
+        Reset the operand held by this :class:`Reshape` instance. This method is used
+        to provide a new operand for execution.
+
+        Args:
+            operand: A tensor (ndarray-like object) compatible with the previous one.
+                The new operand is considered compatible if all the following properties
+                match with the previous one:
+
+                - The operand distribution, which must be (input_box, output_box)
+                  where input_box and output_box are the boxes specified at plan time.
+                - The package that the new operand belongs to.
+                - The dtype of the new operand.
+                - The shape and strides of the new operand.
+                - The memory space of the new operand (CPU or GPU).
+                - The device that new operand belongs to if it is on GPU.
+
+            stream: {stream}
+
+        .. seealso::
+            :meth:`release_operand`
+
+        Examples:
+
+            >>> import cupy as cp
+            >>> import nvmath.distributed
+
+            Get process group used to initialize nvmath.distributed (for information on
+            initializing nvmath.distributed, you can refer to the documentation or to the
+            Reshape examples in `nvmath/examples/distributed/reshape
+            <https://github.com/NVIDIA/nvmath-python/tree/main/examples/distributed/reshape>`_):
+
+            >>> process_group = nvmath.distributed.get_context().process_group
+            >>> nranks = process_group.nranks
+
+            Create a 3-D complex128 ndarray on GPU symmetric memory, initially partitioned
+            on the X axis (the global shape is (128, 128, 128)):
+
+            >>> shape = 128 // nranks, 128, 128
+            >>> dtype = cp.complex128
+            >>> a = nvmath.distributed.allocate_symmetric_memory(shape, cp, dtype=dtype)
+            >>> a[:] = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)
+
+            Compute the input and output box for the desired re-distribution:
+
+            >>> input_box = ...
+            >>> output_box = ...
+
+            Create a Reshape object as a context manager
+
+            >>> with nvmath.distributed.reshape.Reshape(a, input_box, output_box) as f:
+            ...     # Plan the Reshape
+            ...     r.plan()
+            ...
+            ...     # Execute the Reshape to get the first result.
+            ...     r1 = r.execute()
+            ...
+            ...     # Reset the operand to a new CuPy ndarray.
+            ...     b = nvmath.distributed.allocate_symmetric_memory(shape, cp, dtype=dtype)
+            ...     b[:] = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)
+            ...     f.reset_operand(b)
+            ...
+            ...     # Execute to get the new result corresponding to the updated operand.
+            ...     r2 = f.execute()
+
+            With :meth:`reset_operand`, minimal overhead is achieved as problem
+            specification and planning are only performed once.
+
+            For the particular example above, explicitly calling :meth:`reset_operand` is
+            equivalent to updating the operand in-place, i.e, replacing
+            ``f.reset_operand(b)`` with ``a[:]=b``. Note that updating the operand in-place
+            should be adopted with caution as it can only yield the expected result and
+            incur no additional copies under the additional constraints below:
+
+            - The operand's distribution is the same.
+
+            For more details, please refer to `inplace update example
+            <https://github.com/NVIDIA/nvmath-python/tree/main/examples/distributed/reshape/example05_stateful_reset_inplace.py>`_.
+        """
+
+        if operand is None:
+            raise ValueError("Resetting operand requires a valid operand. Use release_operand() to release the operand.")
+
+        self.logger.info("Resetting operand...")
+        # First wrap operand.
+        operand = tensor_wrapper.wrap_operand(operand)
+
+        # Check package match.
+        package = utils.infer_object_package(operand.tensor)
+        if self.package != package:
+            message = f"Library package mismatch: '{self.package}' => '{package}'"
+            raise TypeError(message)
+
+        # Check the dtype.
+        utils.check_attribute_match(self.operand_data_type, operand.dtype, "data type")
+
+        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+        self.logger.info(f"The specified stream for reset_operand() is {stream_holder and stream_holder.obj}.")
+
+        # In principle, we could support memory_space change,
+        # but to handle it properly we need to update self.memory_space and
+        # some dependent properties, like self.blocking, which may be error-prone
+        # from the user perspective. It would prevent inplace optimizations as well.
+        operand_device_id = operand.device_id
+        if operand_device_id != self.operand_device_id:
+
+            def device_str(device_id: int | Literal["cpu"]) -> str:
+                return f"cuda:{device_id}" if isinstance(device_id, int) else f"{device_id}"
+
+            raise ValueError(
+                f"The new operand must be on the same device as the original one. "
+                f"The new operand's device is {device_str(operand_device_id)}, "
+                f"the original device is {device_str(self.operand_device_id)}"
+            )
+
+        if self.memory_space == "cuda" and not operand.is_symmetric_memory:
+            raise TypeError("Distributed reshape requires GPU operand to be on symmetric memory")
+
+        # The plan was made for a specific input box and strides, so the new operand must
+        # match.
+        if operand.shape != self.operand_layout.shape:
+            raise ValueError(
+                f"The shape of the new operand ({operand.shape}) does not match "
+                f"the original shape ({self.operand_layout.shape})."
+            )
+        if operand.strides != self.operand_layout.strides:
+            raise ValueError(
+                f"The strides of the new operand ({operand.strides}) do not match "
+                f"the original strides ({self.operand_layout.strides})."
+            )
+        self.logger.info(f"The reset operand shape = {self.operand_layout.shape}, and strides = {self.operand_layout.strides}.")
+
+        if self.memory_space == "cpu":
+            if self.operand is not None:
+                self.operand.copy_(operand, stream_holder=stream_holder)
+                self.operand_backup = operand
+            else:
+                self.operand, self.operand_backup = _copy_operand_perhaps(
+                    operand,
+                    stream_holder,
+                    self.execution_space,
+                    self.memory_space,
+                    self.device_id,
+                )
+        else:
+            self.operand = operand
+
+        self._operand_released = False
+
+        self.logger.info("The operand has been reset to the specified operand.")
+
+    @utils.precondition(_check_valid_reshape)
+    def release_operand(self):
+        """
+        {release_operand}
+        """
+        if self._operand_released:
+            self.logger.info("Operand has already been released; nothing to do.")
+            return
+
+        if self.memory_space == self.execution_space:
+            # Same-space (GPU input): self.operand is the user's tensor.
+            self.operand = None
+        else:
+            # Cross-space (CPU input): self.operand_backup is the user's tensor,
+            # self.operand is an internal nvshmem device mirror.
+            # Release user reference and free the nvshmem mirror.
+            # Cross-space execution is always blocking, so no synchronization
+            # is needed before freeing.
+            self.operand_backup = None
+            with utils.device_ctx(self.device_id):
+                self.operand.free_symmetric()
+            self.operand = None
+
+        self._operand_released = True
+        self.logger.info("User-provided operand has been released.")
+
+    def _check_planned(self, *args, **kwargs):
+        """ """
+        what = kwargs["what"]
+        if not self.reshape_planned:
+            raise RuntimeError(f"{what} cannot be performed before plan() has been called.")
+
+    def _check_valid_operand(self, *args, **kwargs):
+        """ """
+        what = kwargs["what"]
+        if self._operand_released:
+            raise RuntimeError(
+                f"{what} cannot be performed after the operand has been released. Use reset_operand() to provide a new "
+                f"operand before performing the {what.lower()}."
+            )
+
+    def _free_workspace_memory(self, exception: Exception | None = None) -> bool:
+        """
+        Free workspace by releasing the MemoryPointer object.
+        """
+        if self.workspace_ptr is None:
+            return True
+
+        with utils.device_ctx(self.device_id):
+            # Calling nvshmem_free on memory that's still in use is not safe
+            # (nvshmem_free is not stream-ordered), so we need to wait for the
+            # computation to finish.
+            if self.workspace_stream is not None:
+                self.workspace_stream.sync()
+            self.workspace_ptr.free()
+        self.workspace_ptr = None
+        self.logger.debug("[_free_workspace_memory] The workspace has been released.")
+
+        return True
+
+    @utils.precondition(_check_valid_reshape)
+    @utils.precondition(_check_planned, "Workspace memory allocation")
+    @utils.atomic(_free_workspace_memory, method=True)
+    def _allocate_workspace_memory(self, stream_holder):
+        """
+        Allocate workspace memory using the specified allocator.
+        """
+
+        assert self._workspace_allocated_here is False, "Internal Error."
+
+        self.logger.debug("Allocating workspace for performing the Reshape...")
+        with utils.device_ctx(self.device_id), stream_holder.ctx:
+            try:
+                self.workspace_ptr = self.allocator.memalloc(self.workspace_size)
+                self._workspace_allocated_here = True
+            except TypeError as e:
+                message = (
+                    "The method 'memalloc' in the allocator object must conform to the interface in the "
+                    "'BaseCUDAMemoryManager' protocol."
+                )
+                raise TypeError(message) from e
+
+        self.workspace_stream = stream_holder.obj
+        self.logger.debug(
+            f"Finished allocating device workspace of size {formatters.MemoryStr(self.workspace_size)} in the context "
+            f"of stream {self.workspace_stream}."
+        )
+
+    def _allocate_workspace_memory_perhaps(self, stream_holder):
+        """
+        Allocate workspace memory using the specified allocator, if it hasn't already been
+        done.
+        """
+        if self.execution_space != "cuda" or self.workspace_ptr is not None:
+            return
+
+        return self._allocate_workspace_memory(stream_holder)
+
+    @utils.precondition(_check_valid_reshape)
+    def _free_workspace_memory_perhaps(self, release_workspace):
+        """
+        Free workspace memory if if 'release_workspace' is True.
+        """
+        if not release_workspace:
+            return
+
+        # Establish ordering wrt the computation and free workspace if it's more than the
+        # specified cache limit.
+        if self.last_compute_event is not None:
+            with utils.device_ctx(self.device_id):
+                self.workspace_stream.wait(self.last_compute_event)
+            self.logger.debug("Established ordering with respect to the computation before releasing the workspace.")
+            self.last_compute_event = None
+
+        self.logger.debug("[_free_workspace_memory_perhaps] The workspace memory will be released.")
+        self._free_workspace_memory()
+
+        return True
+
+    def _release_workspace_memory_perhaps(self, exception: Exception | None = None) -> bool:
+        """
+        Free workspace memory if it was allocated in this call
+        (self._workspace_allocated_here == True) when an exception occurs.
+        """
+        release_workspace = self._workspace_allocated_here
+        self.logger.debug(
+            f"[_release_workspace_memory_perhaps] The release_workspace flag is set to {release_workspace} based upon "
+            "the value of 'workspace_allocated_here'."
+        )
+        self._free_workspace_memory_perhaps(release_workspace)
+        self._workspace_allocated_here = False
+        return True
+
+    @utils.precondition(_check_valid_reshape)
+    @utils.precondition(_check_planned, "Execution")
+    @utils.precondition(_check_valid_operand, "Execution")
+    @utils.atomic(_release_workspace_memory_perhaps, method=True)
+    def execute(self, stream: AnyStream | None = None, release_workspace: bool = False, sync_symmetric_memory: bool = True):
+        """
+        Execute the Reshape operation.
+
+        Args:
+            stream: {stream}
+
+            release_workspace: A value of `True` specifies that the Reshape object
+                should release workspace memory back to the symmetric memory pool on
+                function return, while a value of `False` specifies that the object
+                should retain the memory. This option may be set to `True` if the
+                application performs other operations that consume a lot of memory between
+                successive calls to the (same or different) :meth:`execute` API, but incurs
+                an overhead due to obtaining and releasing workspace memory from and
+                to the symmetric memory pool on every call. The default is `False`.
+                **NOTE: All processes must use the same value or the application can
+                deadlock.**
+
+            sync_symmetric_memory: {sync_symmetric_memory}
+
+        Returns:
+            The reshaped operand, which remains on the same device and utilizes the same
+            package as the input operand. For GPU operands, the result will be in symmetric
+            memory and the user is responsible for explicitly deallocating it (for example,
+            using ``nvmath.distributed.free_symmetric_memory(tensor)``).
+        """
+
+        log_info = self.logger.isEnabledFor(logging.INFO)
+        log_debug = self.logger.isEnabledFor(logging.DEBUG)
+
+        stream_holder: StreamHolder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+
+        # Allocate workspace if needed.
+        self._allocate_workspace_memory_perhaps(stream_holder)
+
+        # Allocate output operand.
+        with utils.device_ctx(self.device_id):
+            result = self._allocate_result_operand(stream_holder, log_debug)
+
+        if log_info:
+            self.logger.info("Starting distributed Reshape...")
+            self.logger.info(f"{self.call_prologue}")
+
+        with utils.cuda_call_ctx(stream_holder, self.blocking, timing=log_info) as (
+            self.last_compute_event,
+            elapsed,
+        ):
+            raw_workspace_ptr = utils.get_ptr_from_memory_pointer(self.workspace_ptr)
+            if sync_symmetric_memory:
+                nvshmem.sync_all_on_stream(stream_holder.ptr)
+                if log_info:
+                    self.logger.info(
+                        "sync_symmetric_memory is enabled (this may incur redundant multi-GPU "
+                        "synchronization, please refer to the documentation for more information)"
+                    )
+            elif log_info:
+                self.logger.info("sync_symmetric_memory is disabled")
+            cufftMp.exec_reshape_async(
+                self.handle, result.data_ptr, self.operand.data_ptr, raw_workspace_ptr, stream_holder.ptr
+            )
+
+        if log_info and elapsed.data is not None:
+            self.logger.info(f"The distributed Reshape took {elapsed.data:.3f} ms to complete.")
+
+        # Establish ordering wrt the computation and free workspace if it's more than the
+        # specified cache limit.
+        self._free_workspace_memory_perhaps(release_workspace)
+
+        # reset workspace allocation tracking to False at the end of the methods where
+        # workspace memory is potentially allocated. This is necessary to prevent any
+        # exceptions raised before method entry from using stale tracking values.
+        self._workspace_allocated_here = False
+
+        if self.memory_space == "cpu":
+            out = result.to("cpu", stream_holder=stream_holder).tensor
+            with utils.device_ctx(self.device_id):
+                # Since the execution when user passes CPU operands is blocking, it's
+                # safe to call nvshmem_free here without additional synchronization.
+                result.free_symmetric()
+        else:
+            out = result.tensor
+
+        return out
+
+    def free(self):
+        """Free Reshape resources.
+
+        It is recommended that the :class:`Reshape` object be used within a context, but if
+        it is not possible then this method must be called explicitly to ensure that the
+        Reshape resources (especially internal library objects) are properly cleaned up.
+        """
+
+        if not self.valid_state:
+            return
+
+        try:
+            # Future operations on the workspace stream should be ordered after the
+            # computation.
+            if self.last_compute_event is not None and self.workspace_stream is not None:
+                with utils.device_ctx(self.device_id):
+                    self.workspace_stream.wait(self.last_compute_event)
+                self.last_compute_event = None
+
+            self._free_workspace_memory()
+
+            with utils.device_ctx(self.device_id):
+                if self.handle is not None:
+                    cufftMp.destroy_reshape(self.handle)
+                    self.handle = None
+
+                if self.memory_space == "cpu" and self.operand is not None:
+                    # In this case, self.operand is an internal GPU operand owned by Reshape
+                    # Since the execution when user passes CPU operands is blocking, it's
+                    # safe to call nvshmem_free here without additional synchronization.
+                    self.operand.free_symmetric()
+
+            # Set all attributes to None except for logger and valid_state
+            _keep = {"logger", "valid_state"}
+            for attr in list(vars(self)):
+                if attr not in _keep:
+                    setattr(self, attr, None)
+
+        except Exception as e:
+            self.logger.critical("Internal error: only part of the Reshape object's resources have been released.")
+            self.logger.critical(str(e))
+            raise e
+        finally:
+            self.valid_state = False
+
+        self.logger.info("The Reshape object's resources have been released.")
+
+
+@utils.docstring_decorator(SHARED_RESHAPE_DOCUMENTATION, skip_missing=False)
+def reshape(
+    operand,
+    /,
+    input_box: Box,
+    output_box: Box,
+    *,
+    sync_symmetric_memory: bool = True,
+    options: ReshapeOptions | None = None,
+    stream: AnyStream | None = None,
+):
+    r"""
+    reshape({function_signature})
+
+    Perform a distributed reshape on the provided operand to change its distribution across
+    processes.
+
+    Args:
+        operand: {operand}
+            {operand_admonitions}
+
+        input_box: {input_box}
+
+        output_box: {output_box}
+
+        sync_symmetric_memory: {sync_symmetric_memory}
+
+        options: {options}
+
+        stream: {stream}
+
+    Returns:
+        A tensor that remains on the same device and belongs to the same package as
+        the input operand, with shape according to output_box.
+
+    .. seealso::
+        :class:`Reshape`.
+
+    Examples:
+
+        >>> import cupy as cp
+        >>> import nvmath.distributed
+
+        Get process group used to initialize nvmath.distributed (for information on
+        initializing nvmath.distributed, you can refer to the documentation or to the
+        Reshape examples in `nvmath/examples/distributed/reshape
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/distributed/reshape>`_):
+
+        >>> process_group = nvmath.distributed.get_context().process_group
+
+        Let's create a 3D floating-point ndarray on GPU, distributed across a certain number
+        of processes, with each holding a portion of the ndarray. As an example, process 0
+        holds a 3D box of the global 3D array of shape (4, 4, 4).
+
+        >>> shape = 4, 4, 4
+
+        Reshape uses the NVSHMEM PGAS model, which requires GPU operands to be on the
+        symmetric heap:
+
+        >>> if process_group.rank == 0:
+        ...     a[:] = cp.random.rand(*shape)
+        ... else:
+        ...     a = ...  # each process holds a different section of the global array.
+        ... a = nvmath.distributed.allocate_symmetric_memory(shape, cp)
+
+        With Reshape, we will change how the ndarray is distributed, by having each
+        process specify the input and output section of the global array. For process 0,
+        let's assume that it holds the 3D box that goes from the lower corner given by
+        coordinates (0, 0, 0) to the upper corner (4, 4, 4).
+
+        NOTE: each process has its own input and output boxes which are different to those
+        of other processes, as each holds a different section of the global array.
+
+        >>> from nvmath.distributed.distribution import Box
+        >>> if process_group.rank == 0:
+        ...     input_lower = (0, 0, 0)
+        ...     input_upper = (4, 4, 4)
+        ...     input_box = Box(input_lower, input_upper)
+        ...     output_box = ...
+        ... else:
+        ...     input_box = ...  # the input box depends on the process.
+        ...     output_box = ...  # the output box depends on the process.
+
+        Perform the distributed reshape using :func:`reshape`:
+
+        >>> r = nvmath.distributed.reshape.reshape(a, input_box, output_box)
+
+        See :class:`ReshapeOptions` for the complete list of available options.
+
+        The package current stream is used by default, but a stream can be explicitly
+        provided to the Reshape operation. This can be done if the Reshape operand is
+        computed on a different stream, for example:
+
+        >>> s = cp.cuda.Stream()
+        >>> with s:
+        ...     a = nvmath.distributed.allocate_symmetric_memory(shape, cp)
+        ...     a[:] = cp.random.rand(*shape)
+        >>> r = nvmath.distributed.reshape.reshape(a, stream=s)
+
+        The operation above runs on stream `s` and is ordered with respect to the input
+        computation.
+
+        Create a NumPy ndarray on the CPU.
+
+        >>> import numpy as np
+        >>> b = np.random.rand(*shape)
+
+        Provide the NumPy ndarray to :func:`reshape`, with the result also being a NumPy
+        ndarray:
+
+        >>> r = nvmath.distributed.reshape.reshape(b, input_box, output_box)
+
+    Notes:
+        - This function is a convenience wrapper around :class:`Reshape` and is
+          specifically meant for *single* use. The same computation can be performed
+          with the stateful API.
+
+    Further examples can be found in the `nvmath/examples/distributed/reshape
+    <https://github.com/NVIDIA/nvmath-python/tree/main/examples/distributed/reshape>`_
+    directory.
+    """
+
+    with Reshape(operand, input_box, output_box, options=options, stream=stream) as reshape_obj:
+        # Plan the Reshape.
+        reshape_obj.plan(stream=stream)
+
+        # Execute the Reshape.
+        result = reshape_obj.execute(sync_symmetric_memory=sync_symmetric_memory, stream=stream)
+
+    return result