Validate pencil halo and padding extents

src/cudecomp.cc

@@ -25,6 +25,7 @@
 #include <cstring>
 #include <exception>
 #include <iostream>
+#include <limits>
 #include <memory>
 #include <set>
 #include <string>
@@ -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<int32_t>::max()) {
+    THROW_INVALID_USAGE("computed pencil shape exceeds int32_t limit");
+  }
+  return static_cast<int32_t>(shape);
+}
+
+static int64_t checkedPencilSizeProduct(int64_t size, int64_t shape) {
+  if (size == 0 || shape == 0) { return 0; }
+  if (shape > std::numeric_limits<int64_t>::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<int64_t>(pencil_info->shape[ord]) +
+                                        2 * static_cast<int64_t>(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()

tests/ctest/api_tests.cc

@@ -5,6 +5,7 @@
 
 #include <array>
 #include <cstdint>
+#include <limits>
 #include <memory>
 
 #include <mpi.h>
@@ -1050,15 +1051,60 @@ 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;
   CHECK_CUDECOMP_GLOBAL(active_comm_, cudecompGridDescCreate(handle_, &grid_desc, &config, nullptr));
   cudecomp_test::gridDescGuard grid_desc_guard(handle_, grid_desc);
 
   cudecompPencilInfo_t pinfo;
+  const std::array<int32_t, 3> negative_halo_extents{-1, 0, 0};
+  const std::array<int32_t, 3> negative_padding{0, -1, 0};
+  const std::array<int32_t, 3> oversized_halo_extents{std::numeric_limits<int32_t>::max(), 0, 0};
+  const std::array<int32_t, 3> oversized_padding{std::numeric_limits<int32_t>::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<int32_t>::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) {