Skip to content

Add automatic workspace management feature. #159

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
romerojosh merged 2 commits into from
Aug 25, 2026
Merged

Add automatic workspace management feature. #159

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a9d274c
Add automatic workspace management
romerojosh Aug 25, 2026
e441d81
Fix cuTENSOR dependencies for CTest targets
romerojosh Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions docs/api/f_api.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -471,7 +471,7 @@ _____________________
:p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor.
:p T input(*) [in]: Device array containing input X-axis aligned pencil data.
:p T output(*) [out]: Device array to write output Y-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place
:p T work(*) [in]: Device array to use for transpose workspace.
:p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management.
:p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation.
:p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos.
:p integer output_halo_extents(3) [in,optional]: Similar to :code:`input_halo_extents` but for the output data. If not provided, output data is assumed to have no halos.
Expand All @@ -497,7 +497,7 @@ _____________________
:p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor.
:p T input(*) [in]: Device array containing input Y-axis aligned pencil data.
:p T output(*) [out]: Device array to write output Z-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place
:p T work(*) [in]: Device array to use for transpose workspace.
:p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management.
:p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation.
:p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos.
:p integer output_halo_extents(3) [in,optional]: Similar to :code:`intput_halo_extents` but for the output data. If not provided, output data is assumed to have no halos.
Expand All @@ -523,7 +523,7 @@ _____________________
:p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor.
:p T input(*) [in]: Device array containing input Z-axis aligned pencil data.
:p T output(*) [out]: Device array to write output Y-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place
:p T work(*) [in]: Device array to use for transpose workspace.
:p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management.
:p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation.
:p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos.
:p integer output_halo_extents(3) [in,optional]: Similar to :code:`intput_halo_extents` but for the output data. If not provided, output data is assumed to have no halos.
Expand All @@ -550,7 +550,7 @@ _____________________
:p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor.
:p T input(*) [in]: Device array containing input Y-axis aligned pencil data.
:p T output(*) [out]: Device array to write output X-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place
:p T work(*) [in]: Device array to use for transpose workspace.
:p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management.
:p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation.
:p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos.
:p integer output_halo_extents(3) [in,optional]: Similar to :code:`intput_halo_extents` but for the output data. If not provided, output data is assumed to have no halos.
Expand Down Expand Up @@ -578,7 +578,7 @@ ____________________
:p cudecompHandle handle [in]: The initialized cuDecomp library handle
:p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor.
:p T input(*) [in,out]: Device array containing input X-axis aligned pencil data. On successful completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated.
:p T work(*) [in]: Device array to use for halo workspace.
:p T work(*) [in]: Caller-owned device array to use for halo workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management.
:p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation.
:p integer halo_extents(3) [in]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side).
:p logical halo_periods(3) [in]: An array of three boolean values to define halo periodicity of the input data, in global order. If the i-th entry in this array is true, the domain is treated periodically along the i-th global domain axis.
Expand All @@ -603,7 +603,7 @@ ____________________
:p cudecompHandle handle [in]: The initialized cuDecomp library handle
:p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor.
:p T input(*) [in,out]: Device array containing input Y-axis aligned pencil data. On successful completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated.
:p T work(*) [in]: Device array to use for halo workspace.
:p T work(*) [in]: Caller-owned device array to use for halo workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management.
:p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation.
:p integer halo_extents(3) [in]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side).
:p logical halo_periods(3) [in]: An array of three boolean values to define halo periodicity of the input data, in global order. If the i-th entry in this array is true, the domain is treated periodically along the i-th global domain axis.
Expand All @@ -628,7 +628,7 @@ ____________________
:p cudecompHandle handle [in]: The initialized cuDecomp library handle
:p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor.
:p T input(*) [in,out]: Device array containing input Z-axis aligned pencil data. On successful completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated.
:p T work(*) [in]: Device array to use for halo workspace.
:p T work(*) [in]: Caller-owned device array to use for halo workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management.
:p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation.
:p integer halo_extents(3) [in]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side).
:p logical halo_periods(3) [in]: An array of three boolean values to define halo periodicity of the input data, in global order. If the i-th entry in this array is true, the domain is treated periodically along the i-th global domain axis.
Expand Down
36 changes: 34 additions & 2 deletions docs/basic_usage.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -522,7 +522,9 @@ Besides device memory to store pencil data, cuDecomp also requires workspace buf
is used to facilitate local packing/unpacking and transposition operations (which are currently performed
out-of-place). As a result, this workspace buffer will be approximately 2x the size of the largest pencil
assigned to this process. For halo communication, the workspace is used to facilitate local packing of non-contiguous
halo elements. We can query the required workspace sizes, in number of elements, using the
halo elements.

Applications can allocate caller-owned workspaces. Query the required workspace sizes, in number of elements, using the
:ref:`cudecompGetTransposeWorkspaceSize-ref` and :ref:`cudecompGetHaloWorkspaceSize-ref` functions.

.. tabs::
Expand Down Expand Up @@ -582,6 +584,35 @@ is required for NVSHMEM operations (see NVSHMEM documentation for more details).
istat = cudecompMalloc(handle, grid_desc, halo_work_d, halo_work_num_elements)
call CHECK_CUDECOMP_EXIT(istat)

Alternatively, applications may let cuDecomp manage workspace memory by skipping the size queries and allocations
above and passing :code:`CUDECOMP_WORKSPACE_AUTO` directly to each operation. cuDecomp then allocates, grows, and reuses
handle-owned workspace memory and preserves ordering with the stream passed to each operation. All participating ranks
must consistently choose automatic or explicit workspace management for a given operation. Automatic workspace
management cannot be used while the caller's stream is being captured by a CUDA Graph; provide an explicit workspace
in that case.

.. tabs::

.. code-tab:: c++

CHECK_CUDECOMP_EXIT(cudecompTransposeXToY(handle, grid_desc, data_d, data_d,
CUDECOMP_WORKSPACE_AUTO, CUDECOMP_DOUBLE,
pinfo_x.halo_extents, nullptr, nullptr, nullptr, 0));

CHECK_CUDECOMP_EXIT(cudecompUpdateHalosX(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO,
CUDECOMP_DOUBLE, pinfo_x.halo_extents, halo_periods,
0, nullptr, 0));

.. code-tab:: fortran

istat = cudecompTransposeXToY(handle, grid_desc, data_d, data_d, CUDECOMP_WORKSPACE_AUTO, &
CUDECOMP_DOUBLE, pinfo_x%halo_extents, [0,0,0])
call CHECK_CUDECOMP_EXIT(istat)

istat = cudecompUpdateHalosX(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, CUDECOMP_DOUBLE, &
pinfo_x%halo_extents, halo_periods, 1)
call CHECK_CUDECOMP_EXIT(istat)


Transposing the data
--------------------
Expand Down Expand Up @@ -681,7 +712,8 @@ them unspecified in Fortran.
Cleaning up and finalizing the library
--------------------------------------
Finally, we can clean up resources. Note the usage of :ref:`cudecompFree-ref` to deallocate the workspace arrays
allocated with :ref:`cudecompMalloc-ref`.
allocated with :ref:`cudecompMalloc-ref`. Automatically managed workspaces are owned by the handle and are released by
:code:`cudecompFinalize`; applications must not pass them to :ref:`cudecompFree-ref`.

.. tabs::

Expand Down
Loading
Loading