Autotuning

cuDecomp can autotune process grid dimensions and communication backends. See Autotuning process for the search process shared by all bindings.

Autotuning usage

In this section, we will use a modified version of the basic usage example, explaining the changes required to enable the autotuning feature. The complete Python autotuning example can be run with mpirun -n 4 python basic_usage_autotune.py.

Creating a grid descriptor with autotuning enabled

Enabling autotuning requires additional steps prior to the grid descriptor creation.

To enable process grid autotuning, instead of setting pdims entries in the configuration object to fixed values, initialize them to zero. This indicates to cuDecomp that process grid autotuning is desired. Otherwise, if the pdims entries are set to desired process grid dimensions, process grid autotuning is disabled and these will remain fixed during the autotuning process.

config.pdims = (0, 0)  # P_rows, P_cols

In addition to this modification of pdims, an autotuning options object, cudecomp.GridDescAutotuneOptions, must be created, populated, and passed to cudecomp.Handle.create_grid(). Its constructor initializes all fields to defaults.

options = cudecomp.GridDescAutotuneOptions()

First, let’s go over general autotuning options that affect process grid and communication backend autotuning.

The n_warmup_trials and n_trials entries in the options object control the number of warmup and timed trials run for each tested configuration respectively. Here we set them to their default values.

options.n_warmup_trials = 3
options.n_trials = 5

The dtype entry in the options object controls which data type cuDecomp will use for autotuning. This property takes a cudecomp.DataType member.

options.dtype = cudecomp.DataType.DOUBLE

The disable_mpi_backends, disable_nccl_backends, and disable_nvshmem_backends entries are boolean flags controlling whether the autotuner will test transpose and halo communication backends using the MPI, NCCL, or NVSHMEM libraries respectively. These exclusions are applied after any backend environment-variable filter. By default, these flags are set to false and all three backend families are enabled.

options.disable_mpi_backends = False
options.disable_nccl_backends = False
options.disable_nvshmem_backends = False

The skip_threshold entry allows the autotuner to rapidly skip slow performing configurations. In particular, the autotuner will skip testing a configuration if skip_threshold * t > t_best, where t is the duration of the first timed trial for the configuration and t_best is the average trial time of the current best configuration. By default, the threshold is set to zero which disables any skipping. More aggressive skipping can be useful in cases where exhaustive testing of all possible configurations is too expensive.

options.skip_threshold = 0.0

Moving on, these are the options specific to process grid autotuning.

The grid_mode entry controls which type of communication (transpose or halo) to use to autotune the process grid dimensions (see cudecompAutotuneGridMode_t). By default, transpose communication is used.

options.grid_mode = cudecomp.AutotuneGridMode.TRANSPOSE

The allow_uneven_decompositions entry is a boolean flag controlling whether the autotuner will test process grid dimensions that result in uneven distributions of data (i.e. grids where pencil shapes are not identical across ranks). By default, this flag is set to True and uneven distributions are allowed.

options.allow_uneven_decompositions = True

Next, these are the options specific to transpose communication backend autotuning.

The autotune_transpose_backend entry is a boolean flag controlling whether the autotuner will autotune the communication backend used for transposes. By default, this flag is False and the transpose communication backend is fixed to the value set within the configuration object during the autotuning process. In this example, we set it to true to enable transpose backend autotuning.

options.autotune_transpose_backend = True

The transpose_use_inplace_buffers entry is an array of boolean flags that controls whether the transpose communication during autotuning is performed in-place or out-of-place, on a per operation basis. This choice can impact transpose performance due to some optimized paths that skip intermediate local operations in some situations, depending on the input/output buffer locations.

For example, GridDesc.transpose_x_to_y can be a no-op if:

the process grid yields a decomposition with \(XY\)-slabs (i.e. distributed along the \(Z\)-axis only)

the \(X\)- and \(Y\)-pencils are not in the permuted axis_contiguous layout

the transposition is performed in-place

In this configuration, the \(X\)- and \(Y\)-pencil buffers are in identical layouts and contain the same data elements from the global grid. Since the transpose is in-place, the input is already the output buffer, and no operation is performed. In contrast, an out-of-place transpose would require a copy of data between the input and output buffers.

In this example, we use in-place buffers for all transpose operations so we can set all elements of transpose_use_inplace_buffers to True. By default, the entries are set to False and out-of-place buffers are used during autotuning.

options.transpose_use_inplace_buffers[:] = True

The transpose_op_weights entry is an array of floating point weights that enable adjusting the contribution of the different transpose operations to the trial timings used by the autotuner. By default, the trial timings used by the autotuner are an unweighted sum of the X-to-Y, Y-to-Z, Z-to-Y, and Y-to-X transpose timings. The entries in transpose_op_weights are multiplicative weights that are applied to the contribution of each transpose operation to the total trial timing. This option is meant for programs that may invoke the different transpose operations an unequal number of times and may want the autotuner to emphasize the more frequently invoked transpose operations when measuring the performance of a backend and process grid configuration. For example, setting the weight to 0.0 for one of the transpose operations will indicate to the autotuner that the timing of that operation should not contribute to the trial time sum. On a related note, the autotuner will skip running any transpose operation with a weight of 0.0 for efficiency. In this example, we autotune using the full set of transpose operations, and therefore set all elements of transpose_op_weights to 1.0. We should note that this is the default behavior, and thus there is no need to explicitly set the elements to 1.0 generally.

options.transpose_op_weights[:] = 1.0

There are additional entries to supply per transpose operation input_halo_extents, output_halo_extents, input_padding, and output_padding arguments for use during autotuning: transpose_input_halo_extents, transpose_output_halo_extents, transpose_input_padding, and transpose_output_padding respectively. These entries default to supplying zero halo extent and padding arguments to the transpose operations. In this example, we want the autotuner to supply non-zero input_halo_extents and output_halo_extents arguments to the transpose operations involving the \(X\)-pencil.

options.transpose_input_halo_extents[0] = [1, 1, 1]   # X-to-Y input.
options.transpose_output_halo_extents[3] = [1, 1, 1]  # Y-to-X output.

All other transpose operations do not have non-zero halo extents or padding arguments, so we leave those entries as their default values.

Lastly, these are the options specific to halo communication backend autotuning.

The autotune_halo_backend entry is a boolean flag controlling whether the autotuner will autotune the communication backend used for halo exchanges. By default, this flag is False and the halo communication backend is fixed to the value set within the configuration object during the autotuning process. In this example, we set it to true to enable halo backend autotuning.

options.autotune_halo_backend = True

The halo_extents, halo_periods, and halo_axis define the halo configuration to use during halo autotuning. See documentation on the halo communication routines, like cudecompUpdateHalosX for descriptions of these fields. In this example, we autotune for \(X\)-pencil halo exchanges with one halo element in each direction with periodic boundaries.

options.halo_axis = 0
options.halo_extents = (1, 1, 1)
options.halo_periods = (True, True, True)

There is an additional entry, halo_padding, to supply the padding argument to use during halo autotuning. This entry defaults to supplying zero padding. In this example, we do not supply non-zero padding for the halo operations so we leave this entry default valued.

With the grid descriptor configuration and autotuning options objects created and populated, we can now create the grid descriptor with autotuning.

grid = handle.create_grid(config, autotune_options=options)

Autotuner output and querying results

When autotuning is enabled, cuDecomp reports trial timings and the selected configuration. The output is the same for all bindings; see Autotuner output and querying results for an example and explanation of the reported timings and selections.

After the autotuning process is complete, the grid descriptor is created and ready to use. Entries in the configuration object provided to Handle.create_grid corresponding to autotuned fields (pdims, transpose_comm_backend, and halo_comm_backend) are updated to reflect the autotuning selections. Thus, one can also run the following code like the following to inspect and report the final configuration used:

if rank == 0:
    print(f"running on {config.pdims[0]} x {config.pdims[1]} process grid...")
    print(f"running using {cudecomp.transpose_comm_backend_to_string(config.transpose_comm_backend)} transpose backend...")
    print(f"running using {cudecomp.halo_comm_backend_to_string(config.halo_comm_backend)} halo backend...")

As autotuning only impacts grid descriptor creation, the rest of the usage of the library is unchanged from that illustrated in the basic usage section.