Skip to content

Add rank order option to grid descriptor configuration.#119

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 6 commits into from
Apr 29, 2026
Merged

Add rank order option to grid descriptor configuration. #119

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
13034d6
Add rank order option to grid descriptor configuration.
romerojosh Apr 29, 2026
7c38f1d
Add rank order options to test binaries.
romerojosh Apr 29, 2026
1ce60e2
Set autotuning pdim traversal order based on rank order setting.
romerojosh Apr 29, 2026
6caf44d
Add short rank order test configs.
romerojosh Apr 29, 2026
4d1433c
Fix halo test configuration.
romerojosh Apr 29, 2026
7c40441
Further adjustments to halo rank order tests.
romerojosh Apr 29, 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
8 changes: 8 additions & 0 deletions docs/api/c_api.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -96,6 +96,14 @@ __________________________

------

.. _cudecompRankOrder_t-ref:

cudecompRankOrder_t
__________________
.. doxygenenum :: cudecompRankOrder_t

------

.. _cudecompResult_t-ref:

cudecompResult_t
Expand Down
9 changes: 9 additions & 0 deletions docs/api/f_api.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -47,6 +47,7 @@ ________________________
:f integer gdims(3): dimensions of global data grid
:f integer gdims_dist(3): dimensions of global data grid to use for distribution
:f integer pdims(2): dimensions of process grid
:f cudecompRankOrder rank_order: process grid rank assignment order (default: CUDECOMP_RANK_ORDER_DEFAULT)
:f cudecompTransposeCommType transpose_comm_backend: communication backend to use for transpose communication (default: CUDECOMP_TRANSPOSE_COMM_MPI_P2P)
:f logical transpose_axis_contiguous(3): flag (by axis) indicating if memory should be contiguous along pencil axis (default: [false, false, false])
:f integer transpose_mem_order(3, 3): user-specified memory ordering by axis, overrides transpose_axis_contiguous setting; second index specifies axis,
Expand DownExpand Up@@ -145,6 +146,14 @@ See documention for equivalent C enumerator, :ref:`cudecompAutotuneGridMode_t-re

------

.. _cudecompRankOrder_t-f-ref:

cudecompRankOrder
__________________
See documention for equivalent C enumerator, :ref:`cudecompRankOrder_t-ref`.

------

.. _cudecompResult_t-f-ref:

cudecompResult
Expand Down
13 changes: 13 additions & 0 deletions docs/basic_usage.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -85,6 +85,19 @@ and :code:`pdims[1]` corresponds to :math:`P_{\text{cols}}`. In this example, we

config%pdims = [2, 2] ! [P_rows, P_cols]

By default, ranks are assigned to the process grid in row-major order. To use column-major rank assignment, set
:code:`rank_order` in the configuration structure.

.. tabs::

.. code-tab:: c++

config.rank_order = CUDECOMP_RANK_ORDER_COL_MAJOR;

.. code-tab:: fortran

config%rank_order = CUDECOMP_RANK_ORDER_COL_MAJOR

Next, we set the :code:`gdims` (global grid) entries in the configuration struct. These values correspond to the :math:`X`, :math:`Y`, and :math:`Z`
dimensions of the global grid. In this example, we use a global grid with dimensions :math:`64 \times 64 \times 64`.

Expand Down
7 changes: 5 additions & 2 deletions docs/env_vars.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -95,6 +95,9 @@ CUDECOMP_USE_COL_MAJOR_RANK_ORDER
--------------------------------------
(since v0.6.0)

:code:`CUDECOMP_USE_COL_MAJOR_RANK_ORDER` controls the rank assignment order in the process grid. By default, ranks are assigned in row-major order for consistency with :code:`MPI_Cart_*` routines. When enabled, ranks are assigned in column-major order.
:code:`CUDECOMP_USE_COL_MAJOR_RANK_ORDER` is deprecated. Use the :code:`rank_order` field in
:code:`cudecompGridDescConfig_t`/:code:`cudecompGridDescConfig` instead.

Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable column-major rank assignment.
When :code:`rank_order` is left at :code:`CUDECOMP_RANK_ORDER_DEFAULT`, this environment variable is still honored for
backward compatibility. Setting this variable to :code:`1` enables column-major rank assignment; otherwise rank assignment
is row-major. Explicit :code:`rank_order` settings take precedence over this environment variable.
4 changes: 2 additions & 2 deletions docs/overview.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -34,7 +34,8 @@ the *global grid*. The global grid is decomposed
across :math:`N_{\text{GPU}}` processes in a 2D **process grid** with dimensions :math:`P_{\text{row}} \times P_{\text{col}}`. The processes
are logically grouped by column and row index into :math:`P_{\text{row}}` *row* communicators and :math:`P_{\text{col}}` *column* communicators.
By default, for consistency with :code:`MPI_Cart_*` routines, the ranks are assigned in a row-major ordering (i.e. row communicators
are composed of sequential ranks). This can be changed to column-major ordering using the :code:`CUDECOMP_USE_COL_MAJOR_RANK_ORDER` environment variable (see :ref:`env-var-section-ref`).
are composed of sequential ranks). This can be changed to column-major ordering using the :code:`rank_order` field in
:code:`cudecompGridDescConfig_t`/:code:`cudecompGridDescConfig`.

cuDecomp will distribute the global domain data so that each process is assigned a unique *pencil* of data, with three different
pencil configurations corresponding to different transposed configurations of the global domain. The domain can be
Expand DownExpand Up@@ -62,4 +63,3 @@ shows the :math:`X`-pencil, :math:`Y`-pencil, and :math:`Z`-pencil configuration

.. figure:: images/decomposition.png
:align: center

18 changes: 15 additions & 3 deletions include/cudecomp.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -81,6 +81,16 @@ typedef enum {
CUDECOMP_AUTOTUNE_GRID_HALO = 1 ///< Use halo communication to autotune process grid dimensions
} cudecompAutotuneGridMode_t;

/**
* @brief This enum defines rank assignment order options for the process grid.
*/
typedef enum {
CUDECOMP_RANK_ORDER_DEFAULT = 0, ///< Use the default rank order. This resolves to row-major unless the
///< deprecated CUDECOMP_USE_COL_MAJOR_RANK_ORDER environment variable is set.
CUDECOMP_RANK_ORDER_ROW_MAJOR = 1, ///< Assign ranks in row-major order (default)
CUDECOMP_RANK_ORDER_COL_MAJOR = 2 ///< Assign ranks in column-major order
} cudecompRankOrder_t;

/**
* @brief This enum defines the possible values return values from cuDecomp. Most functions in the cuDecomp library
* will return one of these values to indicate if an operation has completed successfully or an error occured.
Expand DownExpand Up@@ -113,9 +123,11 @@ typedef struct cudecompGridDesc* cudecompGridDesc_t;
*/
typedef struct {
// Grid information
int32_t gdims[3]; ///< dimensions of global data grid
int32_t gdims_dist[3]; ///< dimensions of global data grid to use for distribution
int32_t pdims[2]; ///< dimensions of process grid
int32_t gdims[3]; ///< dimensions of global data grid
int32_t gdims_dist[3]; ///< dimensions of global data grid to use for distribution
int32_t pdims[2]; ///< dimensions of process grid
cudecompRankOrder_t rank_order; ///< process grid rank assignment order
///< (default: CUDECOMP_RANK_ORDER_DEFAULT)

// Transpose settings
cudecompTransposeCommBackend_t transpose_comm_backend; ///< communication backend to use for transpose communication
Expand Down
36 changes: 26 additions & 10 deletions include/internal/common.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -113,10 +113,10 @@ struct cudecompHandle {
""; // directory to write CSV performance reports, empty means no file writing

// Miscellaneous
int32_t device_p2p_ce_count = 0; // number of P2P CEs available
int32_t device_num_sms = 0; // number of SMs on the device
int32_t device_max_threads_per_sm = 0; // maximum threads per SM
bool use_col_major_rank_order = false; // Flag to control whether to use column-major rank order
int32_t device_p2p_ce_count = 0; // number of P2P CEs available
int32_t device_num_sms = 0; // number of SMs on the device
int32_t device_max_threads_per_sm = 0; // maximum threads per SM
bool col_major_rank_order_env_warning_issued = false; // Warn once for deprecated rank order env var
};

// Structure with information about row/column communicator
Expand DownExpand Up@@ -221,15 +221,31 @@ using comm_count_t = int64_t;

enum cudecompCommAxis { CUDECOMP_COMM_COL = 0, CUDECOMP_COMM_ROW = 1 };

static inline void setProcessGridIndex(const cudecompHandle_t handle, cudecompGridDesc_t grid_desc) {
switch (grid_desc->config.rank_order) {
case CUDECOMP_RANK_ORDER_COL_MAJOR:
grid_desc->pidx[0] = handle->rank % grid_desc->config.pdims[0];
grid_desc->pidx[1] = handle->rank / grid_desc->config.pdims[0];
break;
case CUDECOMP_RANK_ORDER_DEFAULT:
case CUDECOMP_RANK_ORDER_ROW_MAJOR:
default:
grid_desc->pidx[0] = handle->rank / grid_desc->config.pdims[1];
grid_desc->pidx[1] = handle->rank % grid_desc->config.pdims[1];
break;
}
}

// Helper function to convert row or column rank to global rank
static inline int getGlobalRank(const cudecompHandle_t handle, const cudecompGridDesc_t grid_desc,
cudecompCommAxis axis, int axis_rank) {
if (handle->use_col_major_rank_order) {
// Column-major rank order
static inline int getGlobalRank(const cudecompHandle_t, const cudecompGridDesc_t grid_desc, cudecompCommAxis axis,
int axis_rank) {
switch (grid_desc->config.rank_order) {
case CUDECOMP_RANK_ORDER_COL_MAJOR:
return (axis == CUDECOMP_COMM_ROW) ? grid_desc->pidx[0] + axis_rank * grid_desc->config.pdims[0]
: grid_desc->config.pdims[0] * grid_desc->pidx[1] + axis_rank;
} else {
// Row-major rank order (default)
case CUDECOMP_RANK_ORDER_DEFAULT:
case CUDECOMP_RANK_ORDER_ROW_MAJOR:
default:
return (axis == CUDECOMP_COMM_ROW) ? grid_desc->config.pdims[1] * grid_desc->pidx[0] + axis_rank
: grid_desc->pidx[1] + axis_rank * grid_desc->config.pdims[1];
}
Expand Down
55 changes: 29 additions & 26 deletions src/autotune.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
*/

#include <algorithm>
#include <array>
#include <numeric>
#include <string>
#include <vector>
Expand DownExpand Up@@ -50,6 +51,20 @@ static std::vector<int> getFactors(int N) {
return factors;
}

static std::vector<std::array<int32_t, 2>> getPdimCandidates(int nranks, cudecompRankOrder_t rank_order) {
std::vector<std::array<int32_t, 2>> pdim_list;
auto factors = getFactors(nranks);
for (auto& factor : factors) {
// Grow the process-grid dimension mapped to contiguous ranks first, preserving the locality-first traversal.
if (rank_order == CUDECOMP_RANK_ORDER_COL_MAJOR) {
pdim_list.push_back({factor, nranks / factor});
} else {
pdim_list.push_back({nranks / factor, factor});
}
}
return pdim_list;
}

template <typename T> static std::vector<T> processTimings(cudecompHandle_t handle, std::vector<T> times, T scale = 1) {
std::sort(times.begin(), times.end());
double t_min = times[0];
Expand DownExpand Up@@ -122,11 +137,11 @@ void autotuneTransposeBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_d
if (!options->transpose_use_inplace_buffers[i]) need_data2 = true;
}

std::vector<int> pdim1_list;
std::vector<std::array<int32_t, 2>> pdim_list;
if (autotune_pdims) {
pdim1_list = getFactors(handle->nranks);
pdim_list = getPdimCandidates(handle->nranks, grid_desc->config.rank_order);
} else {
pdim1_list = {grid_desc->config.pdims[1]};
pdim_list.push_back({grid_desc->config.pdims[0], grid_desc->config.pdims[1]});
}

int32_t pdims_best[2]{grid_desc->config.pdims[0], grid_desc->config.pdims[1]};
Expand All@@ -142,16 +157,10 @@ void autotuneTransposeBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_d
int64_t work_sz = 0;

bool valid = false;
for (auto& pdim1 : pdim1_list) {
grid_desc->config.pdims[0] = handle->nranks / pdim1;
grid_desc->config.pdims[1] = pdim1;
if (handle->use_col_major_rank_order) {
grid_desc->pidx[0] = handle->rank % grid_desc->config.pdims[0];
grid_desc->pidx[1] = handle->rank / grid_desc->config.pdims[0];
} else {
grid_desc->pidx[0] = handle->rank / grid_desc->config.pdims[1];
grid_desc->pidx[1] = handle->rank % grid_desc->config.pdims[1];
}
for (auto& pdims : pdim_list) {
grid_desc->config.pdims[0] = pdims[0];
grid_desc->config.pdims[1] = pdims[1];
setProcessGridIndex(handle, grid_desc);

cudecompPencilInfo_t pinfo_x0, pinfo_x3;
cudecompPencilInfo_t pinfo_y0, pinfo_y1, pinfo_y2, pinfo_y3;
Expand DownExpand Up@@ -604,11 +613,11 @@ void autotuneHaloBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_desc,
#endif
}

std::vector<int> pdim1_list;
std::vector<std::array<int32_t, 2>> pdim_list;
if (autotune_pdims) {
pdim1_list = getFactors(handle->nranks);
pdim_list = getPdimCandidates(handle->nranks, grid_desc->config.rank_order);
} else {
pdim1_list = {grid_desc->config.pdims[1]};
pdim_list.push_back({grid_desc->config.pdims[0], grid_desc->config.pdims[1]});
}

int32_t pdims_best[2]{grid_desc->config.pdims[0], grid_desc->config.pdims[1]};
Expand All@@ -623,16 +632,10 @@ void autotuneHaloBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_desc,
int64_t work_sz = 0;

bool valid = false;
for (auto& pdim1 : pdim1_list) {
grid_desc->config.pdims[0] = handle->nranks / pdim1;
grid_desc->config.pdims[1] = pdim1;
if (handle->use_col_major_rank_order) {
grid_desc->pidx[0] = handle->rank % grid_desc->config.pdims[0];
grid_desc->pidx[1] = handle->rank / grid_desc->config.pdims[0];
} else {
grid_desc->pidx[0] = handle->rank / grid_desc->config.pdims[1];
grid_desc->pidx[1] = handle->rank % grid_desc->config.pdims[1];
}
for (auto& pdims : pdim_list) {
grid_desc->config.pdims[0] = pdims[0];
grid_desc->config.pdims[1] = pdims[1];
setProcessGridIndex(handle, grid_desc);

cudecompPencilInfo_t pinfo;
CHECK_CUDECOMP(cudecompGetPencilInfo(handle, grid_desc, &pinfo, options->halo_axis, options->halo_extents,
Expand Down
46 changes: 30 additions & 16 deletions src/cudecomp.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -142,6 +142,15 @@ static void checkDataType(cudecompDataType_t dtype) {
}
}

static void checkRankOrder(cudecompRankOrder_t rank_order) {
switch (rank_order) {
case CUDECOMP_RANK_ORDER_DEFAULT:
case CUDECOMP_RANK_ORDER_ROW_MAJOR:
case CUDECOMP_RANK_ORDER_COL_MAJOR: return;
default: THROW_INVALID_USAGE("unknown rank order");
}
}

static void checkHandle(cudecompHandle_t handle) {
if (!handle || !handle->initialized) { THROW_INVALID_USAGE("invalid handle"); }
}
Expand All@@ -154,6 +163,7 @@ static void checkConfig(cudecompHandle_t handle, const cudecompGridDescConfig_t*
bool autotune_halos) {
if (!autotune_transpose) { checkTransposeCommBackend(config->transpose_comm_backend); }
if (!autotune_halos) { checkHaloCommBackend(config->halo_comm_backend); }
checkRankOrder(config->rank_order);

int pdims_prod = config->pdims[0] * config->pdims[1];
if (pdims_prod == 0) {
Expand DownExpand Up@@ -368,9 +378,23 @@ static void getCudecompEnvVars(cudecompHandle_t& handle) {
// Check CUDECOMP_PERFORMANCE_REPORT_WRITE_DIR (Directory for CSV performance reports)
const char* performance_write_dir_str = std::getenv("CUDECOMP_PERFORMANCE_REPORT_WRITE_DIR");
if (performance_write_dir_str) { handle->performance_report_write_dir = std::string(performance_write_dir_str); }
}

// Check CUDECOMP_USE_COL_MAJOR_RANK_ORDER (Column-major rank assignment)
handle->use_col_major_rank_order = checkEnvVar("CUDECOMP_USE_COL_MAJOR_RANK_ORDER");
static void resolveRankOrder(cudecompHandle_t handle, cudecompGridDesc_t grid_desc) {
const char* env_val = std::getenv("CUDECOMP_USE_COL_MAJOR_RANK_ORDER");
if (env_val && !handle->col_major_rank_order_env_warning_issued) {
if (handle->rank == 0) {
printf("CUDECOMP:WARN: CUDECOMP_USE_COL_MAJOR_RANK_ORDER is deprecated and will be removed in a future "
"release. Set cudecompGridDescConfig_t::rank_order instead.\n");
}
handle->col_major_rank_order_env_warning_issued = true;
}

if (grid_desc->config.rank_order == CUDECOMP_RANK_ORDER_DEFAULT) {
grid_desc->config.rank_order = (env_val && checkEnvVar("CUDECOMP_USE_COL_MAJOR_RANK_ORDER"))
? CUDECOMP_RANK_ORDER_COL_MAJOR
: CUDECOMP_RANK_ORDER_ROW_MAJOR;
}
}

#ifdef ENABLE_NVSHMEM
Expand DownExpand Up@@ -631,6 +655,7 @@ cudecompResult_t cudecompGridDescCreate(cudecompHandle_t handle, cudecompGridDes
grid_desc = new cudecompGridDesc;
grid_desc->initialized = true;
grid_desc->config = *config;
resolveRankOrder(handle, grid_desc);
auto comm_backend = grid_desc->config.transpose_comm_backend;
auto halo_comm_backend = grid_desc->config.halo_comm_backend;

Expand DownExpand Up@@ -676,13 +701,7 @@ cudecompResult_t cudecompGridDescCreate(cudecompHandle_t handle, cudecompGridDes
if (grid_desc->config.pdims[0] > 0 && grid_desc->config.pdims[1] > 0) {
// If pdims are set, temporarily set up comm info stuctures to determine if we need to create a local NCCL
// communicator
if (handle->use_col_major_rank_order) {
grid_desc->pidx[0] = handle->rank % grid_desc->config.pdims[0];
grid_desc->pidx[1] = handle->rank / grid_desc->config.pdims[0];
} else {
grid_desc->pidx[0] = handle->rank / grid_desc->config.pdims[1];
grid_desc->pidx[1] = handle->rank % grid_desc->config.pdims[1];
}
setProcessGridIndex(handle, grid_desc);
int color_row = grid_desc->pidx[0];
MPI_Comm row_comm;
CHECK_MPI(MPI_Comm_split(handle->mpi_comm, color_row, handle->rank, &row_comm));
Expand DownExpand Up@@ -766,13 +785,7 @@ cudecompResult_t cudecompGridDescCreate(cudecompHandle_t handle, cudecompGridDes
}
}

if (handle->use_col_major_rank_order) {
grid_desc->pidx[0] = handle->rank % grid_desc->config.pdims[0];
grid_desc->pidx[1] = handle->rank / grid_desc->config.pdims[0];
} else {
grid_desc->pidx[0] = handle->rank / grid_desc->config.pdims[1];
grid_desc->pidx[1] = handle->rank % grid_desc->config.pdims[1];
}
setProcessGridIndex(handle, grid_desc);

// Setup final row and column communicators
int color_row = grid_desc->pidx[0];
Expand DownExpand Up@@ -969,6 +982,7 @@ cudecompResult_t cudecompGridDescConfigSetDefaults(cudecompGridDescConfig_t* con
config->gdims[i] = 0;
config->gdims_dist[i] = 0;
}
config->rank_order = CUDECOMP_RANK_ORDER_DEFAULT;

// Transpose Options
config->transpose_comm_backend = CUDECOMP_TRANSPOSE_COMM_MPI_P2P;
Expand Down
8 changes: 8 additions & 0 deletions src/cudecomp_m.cuf
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -48,6 +48,13 @@ module cudecomp
enumerator :: CUDECOMP_AUTOTUNE_GRID_HALO = 1
end enum

! enum for cuDecomp process grid rank order options
enum, bind(c) ! cudecompRankOrder
enumerator :: CUDECOMP_RANK_ORDER_DEFAULT = 0
enumerator :: CUDECOMP_RANK_ORDER_ROW_MAJOR = 1
enumerator :: CUDECOMP_RANK_ORDER_COL_MAJOR = 2
end enum

! enum for cuDecomp supported data types
enum, bind(c) ! cudecompDataType
enumerator :: CUDECOMP_FLOAT = -1
Expand DownExpand Up@@ -88,6 +95,7 @@ module cudecomp
integer(c_int32_t) :: gdims(3) ! dimensions of data grid
integer(c_int32_t) :: gdims_dist(3) ! dimensions of data grid for distribution
integer(c_int32_t) :: pdims(2) ! dimensions of process grid
integer(c_int32_t) :: rank_order ! process grid rank assignment order

! Transpose Options
integer(c_int32_t) :: transpose_comm_backend
Expand Down
Loading
Loading