feat(cutile-rs): add per-operation atomic aggregation

cutile-rs/CHANGELOG.md

@@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ## [Unreleased]
 
+### Added
+
+- Aggregation capabilities for overlapping atomic updates to tensor-shaped
+  storage, one kernel parameter type per operation: `AddTensor`, `AndTensor`,
+  `OrTensor`, `XorTensor`, `MaxTensor` and `MinTensor`, each passed by shared
+  reference — like Rust's own atomics — and each handing out a matching
+  partition (`AddPartition` … `MinPartition`) with one bounds-checked
+  operation (`atomic_add` … `atomic_min`) and its `unsafe` unchecked form,
+  returning the completion `Token`. Element sets are sealed traits
+  (`atomic::AtomicAdd`, `AtomicBitwise`, `AtomicOrd`); comparisons follow the
+  element's signedness, additions pick `AddF`/`Add` by float versus integer.
+  `partition` requires the tile to divide the destination exactly (static
+  shapes decided at compile time, dynamic extents checked before launch).
+  Host binding is through `tensor.atomic()` / `Atomic::owned(tensor)` with
+  exclusive authority; launch-time alias validation accepts same-operation
+  bindings of one allocation, and rejects different operations or any
+  ordinary access to it
+
 ### Changed
 
 - cuTile Rust now lives in the [NVIDIA/cuda-rust](https://github.com/NVIDIA/cuda-rust)

cutile-rs/cutile-book/reference/dsl-api.md

@@ -235,6 +235,103 @@ fn kernel(
 - `tensor.partition_mut(shape)` — Create a mutable `PartitionMut` view (requires `&mut Tensor`; safe because every unproven store through the view is itself `unsafe`)
 - `unsafe tensor.partition_full_mut(shape)` — Create a mutable `PartitionMut` view over the full tensor, without the tile-block offset, for schedule-driven kernels
 
+### Aggregation capabilities
+
+Aggregation capabilities are for kernels in which multiple programs update the
+same output tile. Each capability names the only atomic operation allowed for
+that destination; this prevents two different reductions from being mixed on
+one allocation.
+
+| Tensor | Partition | Checked operation | Elements |
+| --- | --- | --- | --- |
+| `AddTensor` | `AddPartition` | `atomic_add` | `f16`, `bf16`, `f32`, `f64`, `i32`, `u32`, `i64`, `u64` |
+| `AndTensor` | `AndPartition` | `atomic_and` | `i32`, `u32`, `i64`, `u64` |
+| `OrTensor` | `OrPartition` | `atomic_or` | `i32`, `u32`, `i64`, `u64` |
+| `XorTensor` | `XorPartition` | `atomic_xor` | `i32`, `u32`, `i64`, `u64` |
+| `MaxTensor` | `MaxPartition` | `atomic_max` | `i32`, `u32`, `i64`, `u64` |
+| `MinTensor` | `MinPartition` | `atomic_min` | `i32`, `u32`, `i64`, `u64` |
+
+The element sets are sealed. `MaxTensor` and `MinTensor` select signed or
+unsigned comparison from the Rust element type. `AddTensor` uses floating-point
+or integer addition as appropriate. An aggregation parameter does not impose a
+launch-grid bound. A launch may not bind the same allocation as both an
+aggregation destination and an ordinary tensor, or as two different
+aggregation operations. On the host, create the required exclusive binding
+with `tensor.atomic()` or `Atomic::owned(tensor)`.
+
+`partition(tile)` requires the tile shape to divide the destination exactly on
+every axis and returns the matching partition type. Use
+`num_tiles(&partition, axis)` when deriving program indices. The checked
+operation validates its index and returns a completion `Token`; the
+`_unchecked` form is available only when the caller can prove the index is in
+range.
+
+```rust
+#[cutile::entry()]
+fn split_k_atomic(
+    out: &AddTensor<f32, { [-1, -1] }>,
+    lhs: &Tensor<f32, { [-1, -1] }>,
+    rhs: &Tensor<f32, { [-1, -1] }>,
+) {
+    let bid: (i32, i32, i32) = get_tile_block_id();
+    let (m, n, k) = (bid.0, bid.1, bid.2);
+
+    let lhs_part: Partition<f32, { [16, 32] }> = lhs.partition(shape![16, 32]);
+    let rhs_part: Partition<f32, { [32, 16] }> = rhs.partition(shape![32, 16]);
+    let a: Tile<f32, { [16, 32] }> = lhs_part.load([m, k]);
+    let b: Tile<f32, { [32, 16] }> = rhs_part.load([k, n]);
+    let zero: Tile<f32, { [16, 16] }> = constant(0.0f32, shape![16, 16]);
+
+    let acc: AddPartition<f32, { [16, 16] }> = out.partition(shape![16, 16]);
+    let _done: Token = acc.atomic_add([m, n], mma(a, b, zero));
+}
+```
+
+Launch this kernel with an explicit grid whose third axis is the split, for
+example `(tiles_m, tiles_n, k_splits)`. The destination must already contain
+the desired initial value. Common identities are `0` for addition, OR, and
+XOR; all bits set for AND; the minimum representable value for maximum (zero
+for unsigned maximum); and the maximum representable value for minimum.
+
+Do not access the destination through an ordinary tensor binding while the
+atomic launch is in flight. Dynamic extents and tile indices are checked before
+or during launch as appropriate; a statically ragged destination fails during
+kernel compilation. Floating-point addition is not reproducible because the
+interleaving is unspecified. Integer reductions are exact.
+
+Private partials and deterministic merges are ordinary kernels over `Tensor`,
+`Partition`, and `PartitionMut`; the aggregation capabilities are only the
+overlapping-update path.
+
+### Atomic destinations
+
+`Atomic<B>` supplies the exclusive host binding for a shared device-side atomic
+capability:
+
+| Input | Returned | How to build it |
+| --- | --- | --- |
+| `Atomic<&'a mut Tensor<T>>` (borrowed) | `&'a mut Tensor<T>` | `tensor.atomic()` |
+| `Atomic<Tensor<T>>` (owned) | `Tensor<T>` | `Atomic::owned(tensor)` |
+
+A shared `Arc<Tensor<T>>` or `&Tensor<T>` is rejected. The launcher requires
+`KernelAtomic<T>`, implemented only by these two forms.
+
+```rust
+let mut out = api::zeros::<f32>(&[256, 256]).sync_on(&stream)?;
+split_k_atomic(out.atomic(), lhs, rhs)
+    .grid((2, 2, 4))
+    .sync_on(&stream)?;
+```
+
+Because an atomic parameter contributes no grid bound, the grid must come from
+another binding or from an explicit `.grid(..)`.
+
+Within a launch, same-operation atomic bindings may alias. Different
+operations, or an atomic binding with an ordinary tensor access, are rejected.
+The check is allocation-level, so disjoint regions of one allocation are also
+conservatively rejected. Cross-launch coordination remains the caller's
+responsibility.
+
 ### Partition / PartitionMut
 
 `Partition<'a, E, const D: [i32; N]>` — A read-only partitioned view of a tensor, dividing it into tiles indexed by program ID.

cutile-rs/cutile-book/reference/host-api.md

@@ -474,6 +474,7 @@ If any kernel input is `&Tensor<T>` (borrowed), the operation is not
 |---|---|---|
 | `&Tensor<T, S>` | `Tensor<T>`, `Arc<Tensor<T>>`, `&Tensor<T>`, or `&TensorView<T>` | Same as input |
 | `&mut Tensor<T, S>` | `Partition<Tensor<T>>`, `Partition<&mut Tensor<T>>`, or their `MappedLaunchPartition<…>` forms from `.map(..)` | Same as input |
+| `&AddTensor<T, S>` … `&MinTensor<T, S>` | `Atomic<&mut Tensor<T>>` from `tensor.atomic()`, or `Atomic<Tensor<T>>` from `Atomic::owned(tensor)` | The backing `&mut Tensor<T>` / `Tensor<T>` |
 | Scalar (`f32`, `i32`, etc.) | Same scalar | Same scalar |
 | `*mut T`, `*const T` (unsafe only) | `DevicePointer<T>` | `DevicePointer<T>` |
 
@@ -511,6 +512,35 @@ Launchers return arguments in the same ownership form they received.
 | `Partition<Tensor<T>>` (owned) | `Partition<Tensor<T>>` | Yes |
 | `Partition<&'a mut Tensor<T>>` (borrowed) | `Partition<&'a mut Tensor<T>>` | No — tensor is written in place |
 
+### Atomic destinations
+
+`Atomic<B>` supplies the exclusive host binding for a shared device-side atomic
+capability:
+
+| Input | Returned | How to build it |
+|---|---|---|
+| `Atomic<&'a mut Tensor<T>>` (borrowed) | `&'a mut Tensor<T>` | `tensor.atomic()` |
+| `Atomic<Tensor<T>>` (owned) | `Tensor<T>` | `Atomic::owned(tensor)` |
+
+A shared `Arc<Tensor<T>>` or `&Tensor<T>` is rejected. The launcher requires
+`KernelAtomic<T>`, implemented only by these two forms.
+
+```rust
+let mut out = api::zeros::<f32>(&[256, 256]).sync_on(&stream)?;
+split_k_atomic(out.atomic(), lhs, rhs)
+    .grid((2, 2, 4))
+    .sync_on(&stream)?;
+```
+
+Because an atomic parameter contributes no grid bound, the grid must come from
+another binding or from an explicit `.grid(..)`.
+
+Within a launch, same-operation atomic bindings may alias. Different
+operations, or an atomic binding with an ordinary tensor access, are rejected.
+The check is allocation-level, so disjoint regions of one allocation are also
+conservatively rejected. Cross-launch coordination remains the caller's
+responsibility.
+
 The borrowed form is created with `(&mut tensor).partition(shape)`:
 
 ### Owned: `Tensor<T>`

cutile-rs/cutile-compiler/src/compiler/compile_block.rs

@@ -36,7 +36,19 @@ impl<'m> CUDATileFunctionCompiler<'m> {
         matches!(
             ident.to_string().as_str(),
             "Tensor"
+                | "AddTensor"
+                | "AndTensor"
+                | "OrTensor"
+                | "XorTensor"
+                | "MaxTensor"
+                | "MinTensor"
                 | "Partition"
+                | "AddPartition"
+                | "AndPartition"
+                | "OrPartition"
+                | "XorPartition"
+                | "MaxPartition"
+                | "MinPartition"
                 | "BoundedPartition"
                 | "BoundedPartitionMut"
                 | "PartitionMut"

cutile-rs/cutile-compiler/src/compiler/compile_cuda_tile_op.rs

@@ -1307,10 +1307,12 @@ impl<'m> CUDATileFunctionCompiler<'m> {
             }
         };
 
-        let elem_ty_prefix = ptr_value
+        let rust_elem = ptr_value
             .ty
-            .get_cuda_tile_element_type_prefix(self.modules.primitives())?;
-        let atomic_mode = AtomicMode::new(mode.as_str(), elem_ty_prefix)? as i64;
+            .get_instantiated_rust_element_type(self.modules.primitives())
+            .ok_or_else(|| self.jit_error(&call_expr.span(), "unable to determine element type"))?;
+        let atomic_mode = AtomicMode::resolve(mode.as_str(), &rust_elem)
+            .map_err(|e| self.jit_error(&call_expr.span(), &e))? as i64;
 
         let mut operands = vec![ptrs, arg];
         let mut mask_count: i64 = 0;
@@ -1644,10 +1646,12 @@ impl<'m> CUDATileFunctionCompiler<'m> {
                 )
             }
         };
-        let elem_ty_prefix = view_value
-            .ty
-            .get_cuda_tile_element_type_prefix(self.modules.primitives())?;
-        let atomic_mode = AtomicMode::new(mode.as_str(), elem_ty_prefix)? as i64;
+        let atomic_mode = self.resolve_view_atomic_mode(
+            mode.as_str(),
+            &view_value.ty,
+            &value_value.ty,
+            &call_expr.span(),
+        )?;
 
         let mut all_operands = vec![cuda_tile_view_value];
         all_operands.extend_from_slice(&index_values);
@@ -3900,6 +3904,33 @@ impl<'m> CUDATileFunctionCompiler<'m> {
         Ok(Some(result))
     }
 
+    /// Resolves a mode marker from the destination view's Rust element, the
+    /// only record of signedness, and rejects a disagreeing contribution tile.
+    fn resolve_view_atomic_mode(
+        &self,
+        marker: &str,
+        view_ty: &super::tile_rust_type::TileRustType,
+        value_ty: &super::tile_rust_type::TileRustType,
+        span: &proc_macro2::Span,
+    ) -> Result<i64, JITError> {
+        let dest = view_ty
+            .get_instantiated_rust_element_type(self.modules.primitives())
+            .ok_or_else(|| self.jit_error(span, "unable to determine the view's element type"))?;
+        if let Some(contribution) =
+            value_ty.get_instantiated_rust_element_type(self.modules.primitives())
+        {
+            if contribution != dest {
+                return self.jit_error_result(
+                    span,
+                    &format!(
+                        "`{marker}` destination is `{dest}`, but contribution tile is `{contribution}`"
+                    ),
+                );
+            }
+        }
+        Ok(AtomicMode::resolve(marker, &dest).map_err(|e| self.jit_error(span, &e))? as i64)
+    }
+
     fn raw_integer_constant(
         &self,
         module: &mut Module,
@@ -4110,13 +4141,15 @@ impl<'m> CUDATileFunctionCompiler<'m> {
         segments.push(indices.len() as i64);
         if is_atomic {
             let value = compile_value(module, ctx, "value")?;
-            let prefix = value
-                .ty
-                .get_cuda_tile_element_type_prefix(self.modules.primitives())?;
             let mode = super::shared_utils::extract_zst_type_name(argument("mode"), "mode")?;
             builder = builder.attr(
                 "mode",
-                Attribute::i32(AtomicMode::new(&mode, prefix)? as i64),
+                Attribute::i32(self.resolve_view_atomic_mode(
+                    &mode,
+                    &view.ty,
+                    &value.ty,
+                    &call.span(),
+                )?),
             );
             operands.push(
                 value

cutile-rs/cutile-compiler/src/compiler/compile_global.rs

@@ -13,7 +13,7 @@
 use super::_function::CUDATileFunctionCompiler;
 use super::_type;
 use super::_value::{CompilerContext, TileRustValue};
-use super::shared_utils::{AtomicMode, ElementTypePrefix};
+use super::shared_utils::AtomicMode;
 use super::tile_rust_type::TileRustType;
 use crate::error::{JITError, SpannedJITError};
 use crate::generics::GenericVars;
@@ -374,12 +374,9 @@ impl<'m> CUDATileFunctionCompiler<'m> {
             )
         })?;
 
-        let elem_prefix = ElementTypePrefix::new(&info.element_name)?;
-        let mode = if elem_prefix == ElementTypePrefix::Float {
-            AtomicMode::new("AddF", elem_prefix)?
-        } else {
-            AtomicMode::new("Add", elem_prefix)?
-        } as i64;
+        // The same element-selected addition the aggregation types use.
+        let mode = AtomicMode::resolve("Sum", &info.element_name)
+            .map_err(|e| self.jit_error(&method_call.span(), &e))? as i64;
 
         let (op_id, results) =
             OpBuilder::new(Opcode::AtomicRMW, self.ir_location(&method_call.span()))