[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,15 @@ 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
+//! (values outside [0, 3] select the highest-rank dimension, like -1),
+//! usable wherever an stf_get_executor_fn is expected without any FFI
+//! callback cost.
+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 +348,26 @@ 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).
+//!
+//! \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,130 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 <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-rank dimension. 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);
+}