[STF] Expose geometry-aware allocation and native partition functions in the C API

c/experimental/stf/include/cccl/c/experimental/stf/stf.h

@@ -284,6 +284,19 @@ stf_data_place_handle stf_data_place_current_device(void);
 //! \brief Composite partitioned placement over a grid of execution places.
 stf_data_place_handle stf_data_place_composite(stf_exec_place_handle grid, stf_get_executor_fn mapper);
 
+//! \brief Native blocked partition function for a given dimension,
+//! usable wherever an stf_get_executor_fn is expected without any FFI
+//! callback cost.
+//!
+//! The requested dimension is clamped to the highest axis whose extent is
+//! greater than one: values outside [0, 3] (like -1) always select that
+//! axis, and an in-range \p dim beyond it is clamped down to it (e.g.
+//! \p dim 2 on extents {n, 1, 1, 1} partitions along axis 0).
+stf_get_executor_fn stf_partition_fn_blocked(int dim);
+
+//! \brief Native cyclic (round-robin) partition function.
+stf_get_executor_fn stf_partition_fn_cyclic(void);
+
 //! \brief Create a data_place from green-context helper \p helper and view index \p idx.
 //! Returns NULL on failure or if \p idx is out of range.
 stf_data_place_handle stf_data_place_green_ctx(stf_green_context_helper_handle helper, size_t idx);
@@ -339,6 +352,31 @@ void stf_data_place_deallocate(stf_data_place_handle h, void* ptr, size_t size,
 //! \return 1 if stream-ordered, 0 otherwise
 int stf_data_place_allocation_is_stream_ordered(stf_data_place_handle h);
 
+//! \brief Allocate memory at a data place for a tensor with the given extents.
+//!
+//! For most places this is equivalent to allocating
+//! prod(data_dims) * elemsize bytes; composite places use the geometry to
+//! back each block of the allocation on the place owning it according to the
+//! partitioner (which is why the plain byte-count stf_data_place_allocate()
+//! fails on them: a byte count alone does not carry the tensor geometry).
+//! Extents follow the dimension-0-fastest linearization convention; row-major
+//! callers should present reversed extents (and a coordinate-reversing
+//! partitioner).
+//!
+//! The product of the extents and \p elemsize is validated before the
+//! allocation is attempted: if it overflows uint64_t (or exceeds
+//! PTRDIFF_MAX), the call fails and returns NULL instead of silently
+//! wrapping to a smaller byte count.
+//!
+//! \param h         Data place handle (must not be NULL)
+//! \param data_dims Extents of the tensor (must not be NULL)
+//! \param elemsize  Size of one element in bytes
+//! \param stream    CUDA stream for stream-ordered allocation (may be NULL)
+//! \return Pointer to allocated memory (release with
+//!         stf_data_place_deallocate()), or NULL on failure
+void* stf_data_place_allocate_nd(
+  stf_data_place_handle h, const stf_dim4* data_dims, uint64_t elemsize, cudaStream_t stream);
+
 //! \}
 
 //! \defgroup Handles Opaque Handles

c/experimental/stf/src/stf.cu

@@ -522,6 +522,28 @@ stf_data_place_handle stf_data_place_composite(stf_exec_place_handle grid, stf_g
   return to_opaque(dp);
 }
 
+stf_get_executor_fn stf_partition_fn_blocked(int dim)
+{
+  switch (dim)
+  {
+    case 0:
+      return reinterpret_cast<stf_get_executor_fn>(&blocked_partition_custom<0>::get_executor);
+    case 1:
+      return reinterpret_cast<stf_get_executor_fn>(&blocked_partition_custom<1>::get_executor);
+    case 2:
+      return reinterpret_cast<stf_get_executor_fn>(&blocked_partition_custom<2>::get_executor);
+    case 3:
+      return reinterpret_cast<stf_get_executor_fn>(&blocked_partition_custom<3>::get_executor);
+    default:
+      return reinterpret_cast<stf_get_executor_fn>(&blocked_partition::get_executor);
+  }
+}
+
+stf_get_executor_fn stf_partition_fn_cyclic(void)
+{
+  return reinterpret_cast<stf_get_executor_fn>(&cyclic_partition::get_executor);
+}
+
 stf_data_place_handle stf_data_place_green_ctx(stf_green_context_helper_handle helper, size_t idx)
 {
 #if _CCCL_CTK_AT_LEAST(12, 4)
@@ -588,6 +610,29 @@ void* stf_data_place_allocate(stf_data_place_handle h, ptrdiff_t size, cudaStrea
   }
 }
 
+void* stf_data_place_allocate_nd(
+  stf_data_place_handle h, const stf_dim4* data_dims, uint64_t elemsize, cudaStream_t stream)
+{
+  _CCCL_ASSERT(h != nullptr, "data_place handle must not be null");
+  _CCCL_ASSERT(data_dims != nullptr, "data_dims must not be null");
+  dim4 dims;
+  ::std::memcpy(&dims, data_dims, sizeof(dims));
+  try
+  {
+    return from_opaque(h)->allocate_nd(dims, elemsize, stream);
+  }
+  catch (const ::std::exception& e)
+  {
+    fprintf(stderr, "stf_data_place_allocate_nd failed: %s\n", e.what());
+    return nullptr;
+  }
+  catch (...)
+  {
+    fprintf(stderr, "stf_data_place_allocate_nd failed: unknown exception\n");
+    return nullptr;
+  }
+}
+
 void stf_data_place_deallocate(stf_data_place_handle h, void* ptr, size_t size, cudaStream_t stream)
 {
   _CCCL_ASSERT(h != nullptr, "data_place handle must not be null");

c/experimental/stf/test/test_allocate_nd.cpp

@@ -0,0 +1,163 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cstdint>
+#include <vector>
+
+#include <cuda_runtime.h>
+
+#include <c2h/catch2_test_helper.h>
+#include <cccl/c/experimental/stf/stf.h>
+
+namespace
+{
+inline constexpr uint64_t one_mib = 1024 * 1024;
+
+stf_exec_place_handle make_dev0_grid(size_t nplaces)
+{
+  std::vector<stf_exec_place_handle> places(nplaces);
+  for (auto& place : places)
+  {
+    place = stf_exec_place_device(0);
+    REQUIRE(place != nullptr);
+  }
+  stf_exec_place_handle const grid = stf_exec_place_grid_create(places.data(), nplaces, nullptr);
+  REQUIRE(grid != nullptr);
+  for (const auto& place : places)
+  {
+    stf_exec_place_destroy(place);
+  }
+  return grid;
+}
+
+void check_device_round_trip(void* ptr, uint64_t n)
+{
+  const std::vector<int> host(n, 42);
+  REQUIRE(cudaMemcpy(ptr, host.data(), n * sizeof(int), cudaMemcpyHostToDevice) == cudaSuccess);
+  std::vector<int> back(n, 0);
+  REQUIRE(cudaMemcpy(back.data(), ptr, n * sizeof(int), cudaMemcpyDeviceToHost) == cudaSuccess);
+  REQUIRE(back[0] == 42);
+  REQUIRE(back[n - 1] == 42);
+}
+} // namespace
+
+C2H_TEST("shaped allocation on an ordinary data place", "[places][allocate]")
+{
+  constexpr uint64_t n = one_mib; // ints
+  constexpr stf_dim4 dims{n, 1, 1, 1};
+
+  stf_data_place_handle const dp = stf_data_place_device(0);
+  REQUIRE(dp != nullptr);
+
+  // On a non-composite place the geometry degenerates to a byte count
+  void* const ptr = stf_data_place_allocate_nd(dp, &dims, sizeof(int), nullptr);
+  REQUIRE(ptr != nullptr);
+  check_device_round_trip(ptr, n);
+
+  stf_data_place_deallocate(dp, ptr, n * sizeof(int), nullptr);
+  stf_data_place_destroy(dp);
+}
+
+C2H_TEST("shaped allocation on composite data places", "[places][allocate]")
+{
+  stf_exec_place_handle const grid = make_dev0_grid(2);
+
+  constexpr uint64_t n = one_mib; // ints
+  constexpr stf_dim4 dims{n, 1, 1, 1};
+
+  stf_data_place_handle const dp = stf_data_place_composite(grid, stf_partition_fn_blocked(0));
+  REQUIRE(dp != nullptr);
+
+  // A byte count alone cannot carry the tensor geometry: must fail cleanly
+  void* const bad = stf_data_place_allocate(dp, static_cast<ptrdiff_t>(n * sizeof(int)), nullptr);
+  REQUIRE(bad == nullptr);
+
+  void* const ptr = stf_data_place_allocate_nd(dp, &dims, sizeof(int), nullptr);
+  REQUIRE(ptr != nullptr);
+
+  // Memory must be usable from the device
+  check_device_round_trip(ptr, n);
+
+  stf_data_place_deallocate(dp, ptr, n * sizeof(int), nullptr);
+  stf_data_place_destroy(dp);
+
+  // Same flow through the native cyclic partition function
+  stf_data_place_handle const dpc = stf_data_place_composite(grid, stf_partition_fn_cyclic());
+  REQUIRE(dpc != nullptr);
+
+  void* const ptr2 = stf_data_place_allocate_nd(dpc, &dims, sizeof(int), nullptr);
+  REQUIRE(ptr2 != nullptr);
+  check_device_round_trip(ptr2, n);
+
+  stf_data_place_deallocate(dpc, ptr2, n * sizeof(int), nullptr);
+  stf_data_place_destroy(dpc);
+  stf_exec_place_destroy(grid);
+}
+
+C2H_TEST("blocked partition function covers every dimension selector", "[places][allocate]")
+{
+  stf_exec_place_handle const grid = make_dev0_grid(2);
+
+  // 64 * 64 * 16 * 4 ints = 1 MiB: every dimension is divisible by the grid
+  constexpr stf_dim4 dims{64, 64, 16, 4};
+  constexpr uint64_t n = dims.x * dims.y * dims.z * dims.t;
+
+  // Dimensions 0-3 select that axis; out-of-range values (like -1) select the
+  // highest axis whose extent is greater than one. All must yield a usable
+  // native mapper.
+  for (const int dim : {0, 1, 2, 3, -1, 4})
+  {
+    const stf_get_executor_fn mapper = stf_partition_fn_blocked(dim);
+    REQUIRE(mapper != nullptr);
+
+    stf_data_place_handle const dp = stf_data_place_composite(grid, mapper);
+    REQUIRE(dp != nullptr);
+
+    void* const ptr = stf_data_place_allocate_nd(dp, &dims, sizeof(int), nullptr);
+    REQUIRE(ptr != nullptr);
+    check_device_round_trip(ptr, n);
+
+    stf_data_place_deallocate(dp, ptr, n * sizeof(int), nullptr);
+    stf_data_place_destroy(dp);
+  }
+
+  stf_exec_place_destroy(grid);
+}
+
+C2H_TEST("shaped allocation rejects overflowing geometries", "[places][allocate]")
+{
+  // (2^64-1)^2 wraps to 1: an unchecked size computation would hand back a
+  // one-byte allocation for an astronomically large tensor
+  constexpr stf_dim4 huge{UINT64_MAX, UINT64_MAX, 1, 1};
+
+  stf_data_place_handle const dp = stf_data_place_device(0);
+  REQUIRE(dp != nullptr);
+  REQUIRE(stf_data_place_allocate_nd(dp, &huge, 1, nullptr) == nullptr);
+
+  // elemsize participates in the product too
+  constexpr stf_dim4 max_1d{UINT64_MAX, 1, 1, 1};
+  REQUIRE(stf_data_place_allocate_nd(dp, &max_1d, 2, nullptr) == nullptr);
+
+  // A representable product that exceeds PTRDIFF_MAX must also be rejected
+  constexpr stf_dim4 above_ptrdiff{uint64_t{1} << 62, 2, 1, 1};
+  REQUIRE(stf_data_place_allocate_nd(dp, &above_ptrdiff, 1, nullptr) == nullptr);
+
+  stf_data_place_destroy(dp);
+
+  // On a composite place the wrapped geometry used to reach the blocked
+  // partitioner with a zero part_size and kill the process with SIGFPE
+  stf_exec_place_handle const grid = make_dev0_grid(2);
+  stf_data_place_handle const dpc  = stf_data_place_composite(grid, stf_partition_fn_blocked(1));
+  REQUIRE(dpc != nullptr);
+  REQUIRE(stf_data_place_allocate_nd(dpc, &huge, 1, nullptr) == nullptr);
+
+  stf_data_place_destroy(dpc);
+  stf_exec_place_destroy(grid);
+}

cudax/include/cuda/experimental/__places/partitions/blocked_partition.cuh

@@ -114,8 +114,14 @@ public:
 
     size_t extent = data_dims.get(target_dim);
 
-    size_t nplaces   = grid_dims.x;
+    size_t nplaces = grid_dims.x;
+    _CCCL_ASSERT(nplaces > 0, "blocked partition requires a non-empty grid");
+
     size_t part_size = (extent + nplaces - 1) / nplaces;
+    // A zero part_size (empty extent, or extent + nplaces - 1 wrapping) would
+    // make the division below SIGFPE; allocate_nd() rejects such geometries
+    // before the mapper runs
+    _CCCL_ASSERT(part_size > 0, "blocked partition applied to an empty or wrapping extent");
 
     // Get the coordinate in the selected dimension
     size_t c = data_coords.get(target_dim);
@@ -130,6 +136,12 @@ public:
 //! across execution places. By default, partitioning occurs along the last dimension, but a
 //! specific dimension can be selected using the template parameter. This approach ensures
 //! good spatial locality and is particularly effective for regular data access patterns.
+//!
+//! When mapping element coordinates (get_executor), the selected dimension is
+//! clamped to the highest axis whose extent is greater than one: -1 always
+//! selects that axis, and a larger explicit dimension is clamped down to it
+//! (e.g. blocked_partition_custom<2> on extents {n, 1, 1, 1} partitions along
+//! axis 0).
 using blocked_partition = blocked_partition_custom<>;
 
 #ifdef UNITTESTED_FILE

cudax/include/cuda/experimental/__places/places.cuh

@@ -34,6 +34,7 @@
 #include <cuda/experimental/__places/exec_place_resources.cuh>
 #include <cuda/experimental/__stf/utility/core.cuh>
 
+#include <limits>
 #include <stdexcept>
 #include <typeinfo>
 
@@ -373,9 +374,29 @@ public:
    * composite places use the geometry to back each block of the allocation on
    * the place that owns it according to the partitioner. Extents follow the
    * dimension-0-fastest convention of dim4::get_index().
+   *
+   * @throws std::invalid_argument if the product of the extents and elemsize
+   * overflows size_t or exceeds PTRDIFF_MAX
    */
   void* allocate_nd(dim4 data_dims, size_t elemsize, cudaStream_t stream = nullptr) const
   {
+    // An unchecked x*y*z*t*elemsize product can wrap: an ordinary place would
+    // then silently return a tiny allocation for an astronomically large
+    // tensor, and a composite place would feed the degenerate size to its
+    // partitioner. Reject any geometry whose byte count is not representable.
+    size_t total_bytes = elemsize;
+    for (size_t extent : {data_dims.x, data_dims.y, data_dims.z, data_dims.t})
+    {
+      if (extent != 0 && total_bytes > ::std::numeric_limits<size_t>::max() / extent)
+      {
+        throw ::std::invalid_argument("allocate_nd: extents and element size overflow the addressable byte count");
+      }
+      total_bytes *= extent;
+    }
+    if (total_bytes > static_cast<size_t>(::std::numeric_limits<::std::ptrdiff_t>::max()))
+    {
+      throw ::std::invalid_argument("allocate_nd: allocation size exceeds PTRDIFF_MAX");
+    }
     return pimpl_->allocate_nd(data_dims, elemsize, stream);
   }
 

cudax/test/places/placement.cu

@@ -247,6 +247,44 @@ void test_shaped_alloc_cute_composite(int ndevs)
   printf("  shaped allocation (cute composite) test PASSED\n");
 }
 
+void test_shaped_alloc_overflow(int ndevs)
+{
+  printf("Testing shaped allocation with overflowing geometries...\n");
+
+  const size_t huge = ::std::numeric_limits<size_t>::max();
+  // (2^64-1)^2 wraps to 1: an unchecked size computation would hand back a
+  // one-byte allocation for an astronomically large tensor
+  const dim4 wrapping_dims(huge, huge, 1, 1);
+
+  auto expect_invalid = [](const data_place& dp, dim4 dims, size_t elemsize) {
+    bool thrown = false;
+    try
+    {
+      dp.allocate_nd(dims, elemsize);
+    }
+    catch (const ::std::invalid_argument&)
+    {
+      thrown = true;
+    }
+    EXPECT(thrown, "overflowing geometry must throw invalid_argument");
+  };
+
+  data_place dev = data_place::device(0);
+  expect_invalid(dev, wrapping_dims, 1);
+  // elemsize participates in the product too
+  expect_invalid(dev, dim4(huge, 1, 1, 1), 2);
+  // A representable product that exceeds PTRDIFF_MAX must also be rejected
+  expect_invalid(dev, dim4(size_t{1} << 62, 2, 1, 1), 1);
+
+  // On a composite place the wrapped geometry used to reach the blocked
+  // partitioner with a zero part_size and kill the process with SIGFPE
+  auto grid    = make_device_grid(ndevs, 2);
+  data_place c = data_place::composite(blocked_partition_custom<1>{}, grid);
+  expect_invalid(c, wrapping_dims, 1);
+
+  printf("  shaped allocation overflow test PASSED\n");
+}
+
 void test_multi_gpu_residency(int ndevs)
 {
   if (ndevs < 2)
@@ -336,6 +374,7 @@ int main()
   test_evaluate_cute_matches_mapper();
   test_shaped_alloc_callback_composite(ndevs);
   test_shaped_alloc_cute_composite(ndevs);
+  test_shaped_alloc_overflow(ndevs);
   test_multi_gpu_residency(ndevs);
 
   printf("\n=== All placement tests PASSED ===\n");