[STF] pytorch.localized with locality domains

AGENTS.md

@@ -206,6 +206,7 @@ Supported versions: `3.10`, `3.11`, `3.12`, `3.13`
 ### Modules
 
 * **cuda.compute** — Device-level algorithms, iterators, custom GPU types
+* **cuda.stf._experimental** — Sequential Task Flow (CUDASTF) Python bindings in the `cuda-stf` package (Linux only)
 * **cuda.cccl.headers** — Programmatic access to headers
 
 ### Installation
@@ -255,13 +256,15 @@ include_paths = headers.get_include_paths()
 ./ci/test_cuda_compute_python.sh -py-version 3.10
 ./ci/test_cuda_cccl_headers_python.sh -py-version 3.10
 ./ci/test_cuda_cccl_examples_python.sh -py-version 3.10
+./ci/test_cuda_stf_python.sh -py-version 3.10  # Linux only
 ```
 
 Test organization:
 
 * `tests/compute` — Algorithms and iterators
 * `tests/headers` — Header integration
-* `test_examples.py` — Runs compute examples
+* `python/cuda_stf/tests/stf` — Sequential Task Flow (separate `cuda-stf` package, Linux only)
+* `test_examples.py` — Runs compute/coop examples (STF examples live in `python/cuda_stf/tests/test_examples.py`)
 
 ---
 

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

@@ -212,6 +212,22 @@ stf_exec_place_handle stf_exec_place_grid_from_devices(const int* device_ids, si
 stf_exec_place_handle
 stf_exec_place_grid_create(const stf_exec_place_handle* places, size_t count, const stf_dim4* grid_dims);
 
+//! \brief Return a grid with new dimensions and the same linear place order.
+//!
+//! Every extent in \p grid_dims must be positive and their product must equal
+//! the size of \p grid. The returned handle owns an independent grid wrapper;
+//! destroying either handle does not invalidate the other.
+//! \return A new execution-place handle, or NULL if the dimensions are invalid.
+stf_exec_place_handle stf_exec_place_grid_reshape(stf_exec_place_handle grid, const stf_dim4* grid_dims);
+
+//! \brief Collapse a contiguous inclusive range of grid axes.
+//!
+//! Axes in [\p first_axis, \p last_axis] are replaced by one axis whose
+//! extent is their product. Later axes shift left, trailing extents become
+//! one, and linear place order is preserved.
+//! \return A new execution-place handle, or NULL if the axis range is invalid.
+stf_exec_place_handle stf_exec_place_grid_collapse_axes(stf_exec_place_handle grid, size_t first_axis, size_t last_axis);
+
 //! \brief Same as stf_exec_place_destroy (grids are exec_place handles).
 void stf_exec_place_grid_destroy(stf_exec_place_handle grid);
 
@@ -284,6 +300,42 @@ 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 Number of locality domains of a device. Never 0 for a valid
+//! device: without native locality-domain support (pre-13.4 toolkit, or a
+//! driver that cannot answer the query) the device reports a single domain
+//! covering the whole device. Returns 0 only on error (invalid device;
+//! detail on stderr).
+uint32_t stf_locality_domain_count(int dev_id);
+
+//! \brief Execution place pinned to one locality domain of a device (the
+//! whole device with the fallback backend). Ordinals are identity tokens,
+//! validated lazily at use (native backend).
+stf_exec_place_handle stf_exec_place_locality_domain(int dev_id, int domain_id);
+
+//! \brief Grid with one execution place per locality domain of \p dev_id
+//! (a single whole-device place with the fallback backend).
+stf_exec_place_handle stf_exec_place_locality_domain_grid(int dev_id);
+
+//! \brief Data place whose allocations are localized to one locality
+//! domain of a device (plain device memory with the fallback backend).
+stf_data_place_handle stf_data_place_locality_domain(int dev_id, int domain_id);
+
+//! \brief Replicated placement: one copy of the data in the affine memory of
+//! every member of \p grid. Read-only at the place: mutate the data at
+//! another place, the next replicated read re-broadcasts.
+stf_data_place_handle stf_data_place_replicated(stf_exec_place_handle grid);
+
+//! \brief Deferred replicated placement: replicated over the grid of
+//! whichever task the dependency is used with (bound at task acquisition; a
+//! scalar execution place degenerates to its affine data place).
+stf_data_place_handle stf_data_place_replicated_deferred(void);
+
+//! \brief Whether \p h is a replicated data place (concrete or deferred).
+//! Replicated places only support read access; bindings can validate at
+//! dependency construction instead of hitting the C++ exception at task
+//! creation. Returns 1 if replicated, 0 otherwise.
+int stf_data_place_is_replicated(stf_data_place_handle h);
+
 //! \brief Native blocked partition function for a given dimension,
 //! usable wherever an stf_get_executor_fn is expected without any FFI
 //! callback cost.
@@ -379,6 +431,191 @@ void* stf_data_place_allocate_nd(
 
 //! \}
 
+//! \defgroup Placement Tensor placement description and evaluation
+//! \brief Structured partitions (cute_partition) and placement statistics
+//! \{
+
+//! \brief Opaque handle to a structured tensor partition (see
+//! stf_cute_partition_create()). Caller owns the handle; release with
+//! stf_cute_partition_destroy().
+typedef struct stf_cute_partition_opaque_t* stf_cute_partition_handle;
+
+//! \brief Statistics describing how a localized allocation (or a dry-run
+//! evaluation of one) distributes a tensor over data places.
+//! The estimated fraction of block-local bytes ("accuracy") is
+//! matching_samples / total_samples.
+typedef struct stf_placement_stats
+{
+  uint64_t total_bytes; //!< requested payload size in bytes
+  uint64_t vm_bytes; //!< block-rounded virtual reservation size in bytes
+  uint64_t block_size; //!< placement granularity in bytes
+  uint64_t nblocks; //!< number of placement blocks
+  uint64_t nallocs; //!< physical allocations after merging same-owner runs
+  uint64_t total_samples; //!< probes drawn by the block-owner sampler
+  uint64_t matching_samples; //!< probes agreeing with the chosen block owner
+  uint64_t replication_factor; //!< copies of each byte along replicated partition axes (1 = none);
+                               //!< total resident bytes = vm_bytes * replication_factor, and the
+                               //!< bytes_per_grid_index output already counts every copy
+} stf_placement_stats;
+
+//! \brief Per-dimension distribution policy (see stf_partition_dim_spec).
+typedef enum stf_dim_policy
+{
+  STF_DIM_WHOLE        = 0, //!< dimension is not distributed
+  STF_DIM_BLOCKED      = 1, //!< contiguous chunks of ceil(extent / places)
+  STF_DIM_CYCLIC       = 2, //!< round-robin elements
+  STF_DIM_BLOCK_CYCLIC = 3 //!< round-robin blocks of a given size
+} stf_dim_policy;
+
+//! \brief Per-dimension entry of a JAX-like partition specification.
+typedef struct stf_partition_dim_spec
+{
+  int policy; //!< an stf_dim_policy value
+  int mesh_axis; //!< grid axis this dimension distributes over (ignored for STF_DIM_WHOLE)
+  uint64_t block; //!< block size (STF_DIM_BLOCK_CYCLIC only)
+} stf_partition_dim_spec;
+
+//! \brief Evaluate - without allocating - how a localized allocation would
+//! distribute a tensor over the places of a grid.
+//!
+//! Runs the exact same block-owner decision procedure as the allocation path
+//! and returns the resulting statistics, so a candidate mapping can be scored
+//! (and its parameters tuned) before committing memory.
+//!
+//! \param grid      Grid of execution places (must not be NULL)
+//! \param mapper    Partition function mapping element coordinates to a place
+//! \param data_dims Extents of the tensor (dimension 0 fastest; must not be NULL)
+//! \param elemsize  Size of one element in bytes
+//! \param probes    Samples per block for the majority vote (0 = default)
+//! \param block_size Placement granularity in bytes; 0 selects the device
+//!        allocation granularity when a device is present (2 MiB otherwise)
+//! \param out_stats Filled with the resulting statistics (must not be NULL)
+//! \param bytes_per_grid_index Optional array of one entry per grid position
+//!        (length = product of the grid dims), filled with the bytes owned by
+//!        each position; pass NULL to skip
+//! \return 0 on success, non-zero on failure (diagnostic on stderr)
+int stf_placement_evaluate(
+  stf_exec_place_handle grid,
+  stf_get_executor_fn mapper,
+  const stf_dim4* data_dims,
+  uint64_t elemsize,
+  uint64_t probes,
+  uint64_t block_size,
+  stf_placement_stats* out_stats,
+  uint64_t* bytes_per_grid_index);
+
+//! \brief Variant of stf_placement_evaluate() for a structured partition.
+//! The tensor extents are the partition's true extents.
+int stf_placement_evaluate_partition(
+  stf_exec_place_handle grid,
+  stf_cute_partition_handle partition,
+  uint64_t elemsize,
+  uint64_t probes,
+  uint64_t block_size,
+  stf_placement_stats* out_stats,
+  uint64_t* bytes_per_grid_index);
+
+//! \brief Build a structured partition from a JAX-like per-dimension
+//! specification ("dimension 1, blocked over grid axis 0").
+//!
+//! Split dimensions are padded up to divisibility so the underlying layout is
+//! exact; coordinates beyond the true extents own no bytes (predication).
+//!
+//! \param true_dims True tensor extents (dimension 0 fastest; must not be NULL)
+//! \param grid_dims Extents of the grid of places (must not be NULL)
+//! \param spec      One entry per tensor dimension (must not be NULL)
+//! \param rank      Number of entries in \p spec (at most 4)
+//! \return New partition handle, or NULL on invalid input
+//! \param replicated_axes_mask Bitmask of grid axes holding one copy of their
+//! fiber's bytes per coordinate (bit a = native grid axis a; 0 = none). A
+//! replicated axis must not be bound by any spec entry.
+//! stf_placement_evaluate_partition() reports the per-member copies, and a
+//! composite data place built from such a partition is REPLICATED (read-only;
+//! stf_data_place_is_replicated() returns 1): through a logical data it
+//! resolves to one composite allocation per replicated coordinate. Direct
+//! allocation is rejected -- allocate through a logical data, like
+//! stf_data_place_replicated().
+stf_cute_partition_handle stf_cute_partition_create(
+  const stf_dim4* true_dims,
+  const stf_dim4* grid_dims,
+  const stf_partition_dim_spec* spec,
+  size_t rank,
+  uint32_t replicated_axes_mask);
+
+//! \brief Bitmask of replicated grid axes of \p p (native axis numbering)
+uint32_t stf_cute_partition_replicated_axes(stf_cute_partition_handle p);
+
+//! \brief Number of copies the replicated axes of \p p imply (1 = none)
+uint64_t stf_cute_partition_replication_factor(stf_cute_partition_handle p);
+
+//! \brief Build a structured partition directly from flattened
+//! (extent, stride) leaves (expert form; see the C++ cute_partition docs).
+//! Strides are in linear element units over the padded extents, dimension 0
+//! fastest, leaf 0 fastest within each mode.
+//!
+//! \return New partition handle, or NULL if the leaves do not tile the padded
+//!         space exactly
+stf_cute_partition_handle stf_cute_partition_from_leaves(
+  const uint64_t* place_extents,
+  const int64_t* place_strides,
+  const int* place_axes,
+  size_t num_place_leaves,
+  const uint64_t* local_extents,
+  const int64_t* local_strides,
+  size_t num_local_leaves,
+  const stf_dim4* padded_dims,
+  const stf_dim4* true_dims,
+  const stf_dim4* grid_dims);
+
+//! \brief Destroy a partition handle (NULL is ignored).
+void stf_cute_partition_destroy(stf_cute_partition_handle h);
+
+//! \brief Get the true tensor extents of a partition.
+void stf_cute_partition_true_dims(stf_cute_partition_handle h, stf_dim4* out_dims);
+
+//! \brief Get the padded tensor extents of a partition.
+void stf_cute_partition_padded_dims(stf_cute_partition_handle h, stf_dim4* out_dims);
+
+//! \brief Get the grid extents of a partition.
+void stf_cute_partition_grid_dims(stf_cute_partition_handle h, stf_dim4* out_dims);
+
+//! \brief Number of leaves in the place mode.
+size_t stf_cute_partition_num_place_leaves(stf_cute_partition_handle h);
+
+//! \brief Number of leaves in the local mode.
+size_t stf_cute_partition_num_local_leaves(stf_cute_partition_handle h);
+
+//! \brief Fill the place-mode leaves (arrays sized by
+//! stf_cute_partition_num_place_leaves(); any output may be NULL to skip).
+void stf_cute_partition_get_place_leaves(stf_cute_partition_handle h, uint64_t* extents, int64_t* strides, int* axes);
+
+//! \brief Fill the local-mode leaves (arrays sized by
+//! stf_cute_partition_num_local_leaves(); any output may be NULL to skip).
+void stf_cute_partition_get_local_leaves(stf_cute_partition_handle h, uint64_t* extents, int64_t* strides);
+
+//! \brief Linear element offset (in the padded space) of a place's first
+//! element, given the place's linear index in place-mode order.
+//! Returns UINT64_MAX (with a diagnostic on stderr) if the index is out of
+//! range.
+uint64_t stf_cute_partition_place_offset(stf_cute_partition_handle h, uint64_t place_index);
+
+//! \brief Grid position owning the element at the given data coordinates
+//! (closed-form; coordinates must be within the padded extents).
+//! Returns nonzero on failure (with a diagnostic on stderr).
+int stf_cute_partition_owner(stf_cute_partition_handle h, const stf_pos4* data_coords, stf_pos4* out_grid_pos);
+
+//! \brief Create a composite data place backed by a structured partition.
+//!
+//! Such a place is specific to one tensor (the partition's true extents):
+//! allocate with stf_data_place_allocate_nd() using those extents.
+//!
+//! \param grid      Grid of execution places (must not be NULL)
+//! \param partition Structured partition (must not be NULL; copied)
+//! \return New data place handle, or NULL on failure
+stf_data_place_handle stf_data_place_composite_cute(stf_exec_place_handle grid, stf_cute_partition_handle partition);
+
+//! \}
+
 //! \defgroup Handles Opaque Handles
 //! \brief Opaque handle types for STF objects
 //! \{