autotuneHaloBackend aggressively rejects valid decompositions for flat 3D grids (e.g. NxNx1)

[tategotoazarasi]

When running autotuneHaloBackend on a "flat" 3D grid (where gdims[2] == 1), the autotuner rejects all possible parallel process grid configurations, causing cudecompGridDescCreate to fail with CUDECOMP_RESULT_NOT_SUPPORTED.

This prevents the use of cuDecomp for 2D stencils or simulations mapped to a 3D grid with thickness 1, even when using pencil configurations that should not split the unit dimension (e.g., Z-Pencils).

In src/autotune.cc, the loop for testing grid factors contains the following check (around line 590):

    // Skip any decompositions with empty pencils
    if (std::max(grid_desc->config.pdims[0], grid_desc->config.pdims[1]) >
        std::min(grid_desc->config.gdims[1], grid_desc->config.gdims[2])) {
      continue;
    }

If the user provides a grid with dimensions [N, N, 1] (a 2D plane):

1. grid_desc->config.gdims[2] is 1.
2. std::min(..., 1) results in 1.
3. The check becomes if (std::max(pdims[0], pdims[1]) > 1) continue;.

For any parallel run with nranks > 1, std::max(pdims[0], pdims[1]) will always be greater than 1 (e.g., for 4 ranks: 4x1, 1x4, or 2x2). Consequently, all configurations are skipped.

This check appears to be overly restrictive or assumes a specific pencil orientation (likely X-Pencil) where gdims[1] and gdims[2] are the split dimensions. However, for a Z-Pencil configuration (halo_axis = 2), the decomposition occurs along gdims[0] and gdims[1], and gdims[2] remains contiguous/unsplit. Therefore, the small size of gdims[2] should not invalidate the decomposition.

[romerojosh]

Hi @tategotoazarasi,
Thanks for the report! Indeed, it seems the empty pencil checks are much too aggressive in this case. These checks originate from the transpose autotuning case where by all three pencil orientations are tested by default. For the halo case as you point out, the halo_axis argument to the autotuner limits the testing to a single pencil orientation so the check should be limited to that orientation only. I will put together a fix for this.

[romerojosh]

It looks like this might be a slightly more complicated request than it initially seemed.

It is pretty straightforward to relax this empty pencil check for halos but the issue is that the library currently assumes that if a grid configuration is created successfully, all the cuDecomp APIs will work correctly and I'm not certain this is true for grid configurations with empty pencils since that is not a case we have tried. I will still look continue to look into this.

[romerojosh]

Hi @tategotoazarasi,
Can you see if #102 addresses your use case? As noted in the PR, it might not be the solution I end up going for, but it should unblock you nonetheless.

[romerojosh]

Hi @tategotoazarasi, could you could try out #103 instead of #102? We prefer that one as it is a bit less intrusive a change to enable the functionality you requested. Apologies for the churn.