From f56a25723452668697fa75dea0f574e48fc71efc Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 24 Jun 2026 00:18:12 +0200 Subject: [PATCH] Validate pencil halo and padding extents Signed-off-by: Minh Vu --- src/cudecomp.cc | 55 +++++++++++++++++++++++++++------------- tests/ctest/api_tests.cc | 46 +++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 17 deletions(-) diff --git a/src/cudecomp.cc b/src/cudecomp.cc index 6096ffd..f3981be 100644 --- a/src/cudecomp.cc +++ b/src/cudecomp.cc @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -272,6 +273,31 @@ static void checkConfig(cudecompHandle_t handle, const cudecompGridDescConfig_t* } } +static int32_t getNonnegativePencilExtent(const int32_t extents[], int dim, const char* name) { + int32_t value = extents ? extents[dim] : 0; + if (value < 0) { + std::string message = std::string(name) + " values must be non-negative"; + THROW_INVALID_USAGE(message.c_str()); + } + return value; +} + +static int32_t checkedPencilShape(int64_t shape) { + if (shape < 0) { THROW_INVALID_USAGE("computed pencil shape values must be non-negative"); } + if (shape > std::numeric_limits::max()) { + THROW_INVALID_USAGE("computed pencil shape exceeds int32_t limit"); + } + return static_cast(shape); +} + +static int64_t checkedPencilSizeProduct(int64_t size, int64_t shape) { + if (size == 0 || shape == 0) { return 0; } + if (shape > std::numeric_limits::max() / size) { + THROW_INVALID_USAGE("computed pencil size exceeds int64_t limit"); + } + return size * shape; +} + static void gatherGlobalMPIInfo(cudecompHandle_t& handle) { // Gather hostnames by rank handle->hostnames.resize(handle->nranks); @@ -1093,34 +1119,29 @@ cudecompResult_t cudecompGetPencilInfo(cudecompHandle_t handle, cudecompGridDesc if (i != axis) { int64_t d = grid_desc->config.gdims_dist[i] / grid_desc->config.pdims[j]; int64_t mod = grid_desc->config.gdims_dist[i] % grid_desc->config.pdims[j]; - pencil_info->shape[ord] = d; - pencil_info->shape[ord] += (grid_desc->pidx[j] < mod) ? 1 : 0; + int64_t shape = d; + shape += (grid_desc->pidx[j] < mod) ? 1 : 0; if (grid_desc->pidx[j] == std::min(grid_desc->config.pdims[j], grid_desc->config.gdims_dist[i]) - 1) { // Tack any difference in gdim and gdims_dist to last pencil in gdims_dist decomposition - pencil_info->shape[ord] += (grid_desc->config.gdims[i] - grid_desc->config.gdims_dist[i]); + shape += (grid_desc->config.gdims[i] - grid_desc->config.gdims_dist[i]); } + pencil_info->shape[ord] = checkedPencilShape(shape); pencil_info->lo[ord] = (grid_desc->pidx[j] * d + std::min((int64_t)grid_desc->pidx[j], mod)); j++; } else { - pencil_info->shape[ord] = grid_desc->config.gdims[i]; + pencil_info->shape[ord] = checkedPencilShape(grid_desc->config.gdims[i]); pencil_info->lo[ord] = 0; } pencil_info->hi[ord] = pencil_info->lo[ord] + pencil_info->shape[ord] - 1; - if (halo_extents) { - pencil_info->shape[ord] += 2 * halo_extents[i]; - pencil_info->halo_extents[i] = halo_extents[i]; - } else { - pencil_info->halo_extents[i] = 0; - } + pencil_info->halo_extents[i] = getNonnegativePencilExtent(halo_extents, i, "halo_extents"); + pencil_info->padding[i] = getNonnegativePencilExtent(padding, i, "padding"); - if (padding) { - pencil_info->shape[ord] += padding[i]; - pencil_info->padding[i] = padding[i]; - } else { - pencil_info->padding[i] = 0; - } - pencil_info->size *= pencil_info->shape[ord]; + int64_t shape_with_halo_padding = static_cast(pencil_info->shape[ord]) + + 2 * static_cast(pencil_info->halo_extents[i]) + + pencil_info->padding[i]; + pencil_info->shape[ord] = checkedPencilShape(shape_with_halo_padding); + pencil_info->size = checkedPencilSizeProduct(pencil_info->size, pencil_info->shape[ord]); } } CUDECOMP_CATCH_C_API_ERRORS() diff --git a/tests/ctest/api_tests.cc b/tests/ctest/api_tests.cc index 9ebcadd..a56ba67 100644 --- a/tests/ctest/api_tests.cc +++ b/tests/ctest/api_tests.cc @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -1050,6 +1051,24 @@ TEST_F(ApiGetPencilInfoTest, MatchesExpectedGdimsDistDecomposition) { } } +TEST_F(ApiGetPencilInfoTest, DescribesEmptyPencils) { + auto config = emptyPencilConfig(); + cudecompGridDesc_t grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(active_comm_, cudecompGridDescCreate(handle_, &grid_desc, &config, nullptr)); + cudecomp_test::gridDescGuard grid_desc_guard(handle_, grid_desc); + + for (int axis = 0; axis < 3; ++axis) { + cudecompPencilInfo_t pinfo; + CHECK_CUDECOMP_GLOBAL(active_comm_, cudecompGetPencilInfo(handle_, grid_desc, &pinfo, axis, nullptr, nullptr)); + + bool has_empty_shape = false; + for (int i = 0; i < 3; ++i) { + if (pinfo.shape[i] == 0) { has_empty_shape = true; } + } + if (has_empty_shape) { EXPECT_EQ(0, pinfo.size); } + } +} + TEST_F(ApiGetPencilInfoTest, RejectsInvalidArguments) { auto config = distributedConfig(); cudecompGridDesc_t grid_desc = nullptr; @@ -1057,8 +1076,35 @@ TEST_F(ApiGetPencilInfoTest, RejectsInvalidArguments) { cudecomp_test::gridDescGuard grid_desc_guard(handle_, grid_desc); cudecompPencilInfo_t pinfo; + const std::array negative_halo_extents{-1, 0, 0}; + const std::array negative_padding{0, -1, 0}; + const std::array oversized_halo_extents{std::numeric_limits::max(), 0, 0}; + const std::array oversized_padding{std::numeric_limits::max(), 0, 0}; + EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, cudecompGetPencilInfo(handle_, grid_desc, nullptr, 0, nullptr, nullptr)); EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, cudecompGetPencilInfo(handle_, grid_desc, &pinfo, -1, nullptr, nullptr)); + EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, + cudecompGetPencilInfo(handle_, grid_desc, &pinfo, 0, negative_halo_extents.data(), nullptr)); + EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, + cudecompGetPencilInfo(handle_, grid_desc, &pinfo, 0, nullptr, negative_padding.data())); + EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, + cudecompGetPencilInfo(handle_, grid_desc, &pinfo, 0, oversized_halo_extents.data(), nullptr)); + EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, + cudecompGetPencilInfo(handle_, grid_desc, &pinfo, 0, nullptr, oversized_padding.data())); +} + +TEST_F(ApiGetPencilInfoTest, RejectsSizeOverflow) { + auto config = distributedConfig(); + for (int i = 0; i < 3; ++i) { + config.gdims[i] = std::numeric_limits::max(); + } + + cudecompGridDesc_t grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(active_comm_, cudecompGridDescCreate(handle_, &grid_desc, &config, nullptr)); + cudecomp_test::gridDescGuard grid_desc_guard(handle_, grid_desc); + + cudecompPencilInfo_t pinfo; + EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, cudecompGetPencilInfo(handle_, grid_desc, &pinfo, 0, nullptr, nullptr)); } TEST_F(ApiGetTransposeWorkspaceSizeTest, RejectsInvalidArguments) {