From 6fdc67ecb2debfd7b2e1b4ca96216e6a5054f899 Mon Sep 17 00:00:00 2001 From: Robert Ream <154010+robertream@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:03:43 -0700 Subject: [PATCH] feat(cutile-rs): add per-operation atomic aggregation Adds typed atomic aggregation capabilities with compile-time operation and element validation, launch-time alias checks, and compile-only and GPU test coverage. Depends on #1557 (fix/cutile-compiler: preserve unsigned integer literal ranges). Storage identity is opaque: `Tensor::storage_id`, the stored-binding traits, and `validate_binding_aliases` carry a `StorageId` (equality is its whole public surface) rather than a bare `usize` address. Library kernels built on these capabilities live in `cutile-kernels`, generic over the element wherever the added value comes from an input (`scatter_add`, `embedding_backward_scatter_add`, `segment_sum`, `weighted_histogram` over `AtomicAdd`; `max_reduce` over `AtomicOrd` with the element selecting signed vs unsigned for both the intra-tile reduction and the atomic; `bitmask_or` over `AtomicBitwise`). Compile-only tests drive each kernel across its whole element set asserting the emitted mode per pair, host references cover the edge cases, and device tests compare against them. GPU correctness of the unsigned-comparison and bf16/64-bit-constant cells depends on the literal-encoding fixes (PR #1557); the combined tree passed 20/20 required tests on an H100 (sm_90). `ElementType` gains a `ONE` constant, so every kernel is generic over its element, `histogram` included (its increment is `E::ONE` materialized on-device). Signed-off-by: Robert Ream <154010+robertream@users.noreply.github.com> --- cutile-rs/CHANGELOG.md | 18 + cutile-rs/Cargo.lock | 1 + cutile-rs/cutile-book/reference/dsl-api.md | 97 + cutile-rs/cutile-book/reference/host-api.md | 30 + .../src/compiler/compile_block.rs | 12 + .../src/compiler/compile_cuda_tile_op.rs | 55 +- .../src/compiler/compile_global.rs | 11 +- .../src/compiler/compile_intrinsic.rs | 132 +- .../src/compiler/shared_utils.rs | 37 + .../src/compiler/tile_rust_type.rs | 15 - .../src/kernel_entry_generator.rs | 69 +- cutile-rs/cutile-kernels/Cargo.toml | 1 + cutile-rs/cutile-kernels/src/aggregation.rs | 208 ++ cutile-rs/cutile-kernels/src/lib.rs | 5 + cutile-rs/cutile-kernels/tests/aggregation.rs | 369 ++++ cutile-rs/cutile-kernels/tests/gpu.rs | 219 +++ cutile-rs/cutile-macro/src/_module.rs | 8 +- .../src/kernel_launcher_generator.rs | 330 +++- .../cutile-macro/src/validate_dsl_syntax.rs | 15 + cutile-rs/cutile/src/_core.rs | 1354 ++++++++++++- cutile-rs/cutile/src/_tileir.rs | 4 +- cutile-rs/cutile/src/prelude.rs | 6 +- cutile-rs/cutile/src/tensor.rs | 213 +++ cutile-rs/cutile/src/tile_kernel.rs | 212 +++ cutile-rs/cutile/tests/aggregation_tensor.rs | 1671 +++++++++++++++++ cutile-rs/cutile/tests/gpu.rs | 3 + .../cutile/tests/gpu/aggregation_tensor.rs | 1003 ++++++++++ cutile-rs/cutile/tests/ui.rs | 5 + ...egation_binding_cannot_be_reimplemented.rs | 12 + ...ion_binding_cannot_be_reimplemented.stderr | 51 + ...n_tensor_needs_exclusive_host_authority.rs | 24 + ...nsor_needs_exclusive_host_authority.stderr | 20 + ...gation_tensor_rejects_mutable_reference.rs | 16 + ...on_tensor_rejects_mutable_reference.stderr | 5 + ...ion_tensor_requires_a_supported_element.rs | 18 + ...tensor_requires_a_supported_element.stderr | 77 + ...gation_tensor_withholds_ordinary_access.rs | 33 + ...on_tensor_withholds_ordinary_access.stderr | 58 + .../aggregation_withholds_other_operations.rs | 17 + ...regation_withholds_other_operations.stderr | 11 + .../max_tensor_requires_an_ordered_element.rs | 17 + ..._tensor_requires_an_ordered_element.stderr | 85 + cutile-rs/cutile/tests/ui_aggregation.rs | 39 + cutile-rs/scripts/run_cpu_tests.sh | 2 + 44 files changed, 6510 insertions(+), 78 deletions(-) create mode 100644 cutile-rs/cutile-kernels/src/aggregation.rs create mode 100644 cutile-rs/cutile-kernels/tests/aggregation.rs create mode 100644 cutile-rs/cutile-kernels/tests/gpu.rs create mode 100644 cutile-rs/cutile/tests/aggregation_tensor.rs create mode 100644 cutile-rs/cutile/tests/gpu/aggregation_tensor.rs create mode 100644 cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.rs create mode 100644 cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.stderr create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.rs create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.stderr create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.rs create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.stderr create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.rs create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.stderr create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.rs create mode 100644 cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.stderr create mode 100644 cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.rs create mode 100644 cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.stderr create mode 100644 cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.rs create mode 100644 cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.stderr create mode 100644 cutile-rs/cutile/tests/ui_aggregation.rs diff --git a/cutile-rs/CHANGELOG.md b/cutile-rs/CHANGELOG.md index 4b7ec6eff9..03e5d97448 100644 --- a/cutile-rs/CHANGELOG.md +++ b/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) diff --git a/cutile-rs/Cargo.lock b/cutile-rs/Cargo.lock index d95ed3effd..7396e02d67 100644 --- a/cutile-rs/Cargo.lock +++ b/cutile-rs/Cargo.lock @@ -556,6 +556,7 @@ version = "0.4.0" dependencies = [ "cutile", "cutile-compiler", + "cutile-ir", ] [[package]] diff --git a/cutile-rs/cutile-book/reference/dsl-api.md b/cutile-rs/cutile-book/reference/dsl-api.md index 23d5abdf5b..98a176a353 100644 --- a/cutile-rs/cutile-book/reference/dsl-api.md +++ b/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, + lhs: &Tensor, + rhs: &Tensor, +) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let (m, n, k) = (bid.0, bid.1, bid.2); + + let lhs_part: Partition = lhs.partition(shape![16, 32]); + let rhs_part: Partition = rhs.partition(shape![32, 16]); + let a: Tile = lhs_part.load([m, k]); + let b: Tile = rhs_part.load([k, n]); + let zero: Tile = constant(0.0f32, shape![16, 16]); + + let acc: AddPartition = 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` supplies the exclusive host binding for a shared device-side atomic +capability: + +| Input | Returned | How to build it | +| --- | --- | --- | +| `Atomic<&'a mut Tensor>` (borrowed) | `&'a mut Tensor` | `tensor.atomic()` | +| `Atomic>` (owned) | `Tensor` | `Atomic::owned(tensor)` | + +A shared `Arc>` or `&Tensor` is rejected. The launcher requires +`KernelAtomic`, implemented only by these two forms. + +```rust +let mut out = api::zeros::(&[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. diff --git a/cutile-rs/cutile-book/reference/host-api.md b/cutile-rs/cutile-book/reference/host-api.md index d56cfada89..427835fa0d 100644 --- a/cutile-rs/cutile-book/reference/host-api.md +++ b/cutile-rs/cutile-book/reference/host-api.md @@ -474,6 +474,7 @@ If any kernel input is `&Tensor` (borrowed), the operation is not |---|---|---| | `&Tensor` | `Tensor`, `Arc>`, `&Tensor`, or `&TensorView` | Same as input | | `&mut Tensor` | `Partition>`, `Partition<&mut Tensor>`, or their `MappedLaunchPartition<…>` forms from `.map(..)` | Same as input | +| `&AddTensor` … `&MinTensor` | `Atomic<&mut Tensor>` from `tensor.atomic()`, or `Atomic>` from `Atomic::owned(tensor)` | The backing `&mut Tensor` / `Tensor` | | Scalar (`f32`, `i32`, etc.) | Same scalar | Same scalar | | `*mut T`, `*const T` (unsafe only) | `DevicePointer` | `DevicePointer` | @@ -511,6 +512,35 @@ Launchers return arguments in the same ownership form they received. | `Partition>` (owned) | `Partition>` | Yes | | `Partition<&'a mut Tensor>` (borrowed) | `Partition<&'a mut Tensor>` | No — tensor is written in place | +### Atomic destinations + +`Atomic` supplies the exclusive host binding for a shared device-side atomic +capability: + +| Input | Returned | How to build it | +|---|---|---| +| `Atomic<&'a mut Tensor>` (borrowed) | `&'a mut Tensor` | `tensor.atomic()` | +| `Atomic>` (owned) | `Tensor` | `Atomic::owned(tensor)` | + +A shared `Arc>` or `&Tensor` is rejected. The launcher requires +`KernelAtomic`, implemented only by these two forms. + +```rust +let mut out = api::zeros::(&[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` diff --git a/cutile-rs/cutile-compiler/src/compiler/compile_block.rs b/cutile-rs/cutile-compiler/src/compiler/compile_block.rs index a1de161a30..aade41a3cd 100644 --- a/cutile-rs/cutile-compiler/src/compiler/compile_block.rs +++ b/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" diff --git a/cutile-rs/cutile-compiler/src/compiler/compile_cuda_tile_op.rs b/cutile-rs/cutile-compiler/src/compiler/compile_cuda_tile_op.rs index 3591abaf8a..6bb0625196 100644 --- a/cutile-rs/cutile-compiler/src/compiler/compile_cuda_tile_op.rs +++ b/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 { + 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 diff --git a/cutile-rs/cutile-compiler/src/compiler/compile_global.rs b/cutile-rs/cutile-compiler/src/compiler/compile_global.rs index 0609b429c1..5d39875bbc 100644 --- a/cutile-rs/cutile-compiler/src/compiler/compile_global.rs +++ b/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())) diff --git a/cutile-rs/cutile-compiler/src/compiler/compile_intrinsic.rs b/cutile-rs/cutile-compiler/src/compiler/compile_intrinsic.rs index 26f633e7cb..796a4ae542 100644 --- a/cutile-rs/cutile-compiler/src/compiler/compile_intrinsic.rs +++ b/cutile-rs/cutile-compiler/src/compiler/compile_intrinsic.rs @@ -1081,6 +1081,116 @@ impl<'m> CUDATileFunctionCompiler<'m> { ) .map(Some) } + // Atomic views require exact, unpadded tiling. + "require_tiles_divide_exactly" => { + if call_expr.args.len() != 2 { + return self.jit_error_result( + &call_expr.span(), + &format!( + "`require_tiles_divide_exactly` expects 2 arguments, got {}", + call_expr.args.len() + ), + ); + } + let tensor = self + .compile_expression( + module, + block_id, + &call_expr.args[0], + generic_vars, + ctx, + None, + )? + .ok_or_else(|| { + self.jit_error( + &call_expr.args[0].span(), + "failed to compile the aggregation tensor", + ) + })?; + let tile = self + .compile_expression( + module, + block_id, + &call_expr.args[1], + generic_vars, + ctx, + None, + )? + .ok_or_else(|| { + self.jit_error( + &call_expr.args[1].span(), + "failed to compile the tile shape", + ) + })?; + let tile_shape = + self.static_shape_from_value(&tile, generic_vars, &call_expr.args[1].span())?; + // The destination must be traceable to a kernel parameter. + let Some(origin) = tensor.tensor_origin.as_ref() else { + return self.jit_error_result( + &call_expr.args[0].span(), + "an aggregation partition must come from a kernel parameter, so its \ + tiling can be checked against the destination's extents at launch", + ); + }; + let Some(¶m) = self.param_index.get(origin.as_str()) else { + return self.jit_error_result( + &call_expr.args[0].span(), + &format!("`{origin}` is not a kernel parameter"), + ); + }; + for (axis, &extent) in tile_shape.iter().enumerate() { + if extent < 1 { + return self.jit_error_result( + &call_expr.args[1].span(), + &format!("tile extent on axis {axis} must be positive, got {extent}"), + ); + } + let declared = match self.validator.params.get(param) { + Some(cuda_async::device_context::ValidParamType::Tensor(tensor)) => { + tensor.shape.get(axis).copied().filter(|&e| e >= 0) + } + _ => None, + }; + if let Some(dim) = declared { + if dim % extent != 0 { + return self.jit_error_result( + &call_expr.args[0].span(), + &format!( + "`{origin}` axis {axis} has declared extent {dim}, which is \ + not a multiple of the aggregation tile extent {extent}: an \ + atomic reduction addresses an unpadded view, so a partial \ + edge tile would aggregate outside the tensor" + ), + ); + } + continue; + } + let predicate = cuda_async::predicate::Predicate::DivisibleBy { + term: cuda_async::predicate::Term::atom(cuda_async::predicate::Atom::Dim { + param, + axis, + }), + divisor: i64::from(extent), + }; + if !self.lower_obligation( + predicate, + format!( + "`{origin}` axis {axis} must be a multiple of the aggregation tile \ + extent {extent}: an atomic reduction addresses an unpadded view, so \ + a partial edge tile would aggregate outside the tensor" + ), + ) { + return self.jit_error_result( + &call_expr.args[1].span(), + &format!( + "cannot establish that `{origin}` axis {axis} is a multiple of \ + {extent} at launch time" + ), + ); + } + } + Ok(None) + } "check_bounded_partition_access" => self.compile_check_bounded_partition_access( module, block_id, @@ -2437,14 +2547,20 @@ impl<'m> CUDATileFunctionCompiler<'m> { // Both view types walk the same ladder over the same // goals; they differ only in the Rust type of the // receiver, which the checker never inspects. - "check_partition_access" | "check_partition_access_mut" => self - .compile_check_partition_access( - module, - block_id, - call_expr, - generic_vars, - ctx, - ), + "check_partition_access" + | "check_partition_access_mut" + | "check_add_partition_access" + | "check_and_partition_access" + | "check_or_partition_access" + | "check_xor_partition_access" + | "check_max_partition_access" + | "check_min_partition_access" => self.compile_check_partition_access( + module, + block_id, + call_expr, + generic_vars, + ctx, + ), _ => self.jit_error_result( &call_expr.span(), &format!("Unexpected compiler_op call {}", call_expr_func_str), diff --git a/cutile-rs/cutile-compiler/src/compiler/shared_utils.rs b/cutile-rs/cutile-compiler/src/compiler/shared_utils.rs index e974c0d17f..f721da1ce7 100644 --- a/cutile-rs/cutile-compiler/src/compiler/shared_utils.rs +++ b/cutile-rs/cutile-compiler/src/compiler/shared_utils.rs @@ -113,6 +113,43 @@ impl AtomicMode { } Ok(result) } + + /// Resolves a mode marker against a Rust element type. + pub fn resolve(marker: &str, rust_elem: &str) -> Result { + #[derive(Clone, Copy)] + enum Class { + Float, + Signed, + Unsigned, + } + use Class::{Float, Signed, Unsigned}; + let class = match rust_elem { + "f16" | "bf16" | "f32" | "f64" => Float, + "i32" | "i64" => Signed, + "u32" | "u64" => Unsigned, + _ if matches!(marker, "Sum" | "Maximum" | "Minimum") => { + return Err(format!( + "`{rust_elem}` has no atomic reduction modes, so `{marker}` cannot be resolved" + )) + } + _ => Float, + }; + match (marker, class) { + ("Sum", Float) => Ok(AtomicMode::AddF), + ("Sum", Signed | Unsigned) => Ok(AtomicMode::Add), + ("Maximum", Signed) => Ok(AtomicMode::Max), + ("Maximum", Unsigned) => Ok(AtomicMode::UMax), + ("Minimum", Signed) => Ok(AtomicMode::Min), + ("Minimum", Unsigned) => Ok(AtomicMode::UMin), + ("Maximum" | "Minimum", Float) => Err(format!( + "`{marker}` has no atomic reduction mode for `{rust_elem}`" + )), + (concrete, _) => { + let prefix = ElementTypePrefix::new(rust_elem).map_err(|e| e.to_string())?; + AtomicMode::new(concrete, prefix).map_err(|e| e.to_string()) + } + } + } } // Re-export from bounds.rs (canonical location shared with old compiler). diff --git a/cutile-rs/cutile-compiler/src/compiler/tile_rust_type.rs b/cutile-rs/cutile-compiler/src/compiler/tile_rust_type.rs index 6eb770fe17..263051ed52 100644 --- a/cutile-rs/cutile-compiler/src/compiler/tile_rust_type.rs +++ b/cutile-rs/cutile-compiler/src/compiler/tile_rust_type.rs @@ -14,7 +14,6 @@ use crate::error::JITError; use crate::generics::{GenericVars, TypeInstance}; use crate::types::*; use cutile_ir::ir::Type as TileIrType; -use quote::ToTokens; use std::collections::HashMap; use syn::ItemImpl; @@ -102,20 +101,6 @@ impl TileRustType { )), } } - pub(crate) fn get_cuda_tile_element_type_prefix( - &self, - primitives: &HashMap<(String, String), ItemImpl>, - ) -> Result { - let cuda_elem_ty_str = self - .get_cuda_tile_element_type(primitives)? - .ok_or_else(|| { - JITError::generic_err(&format!( - "unable to determine element type for `{}`", - self.rust_ty.to_token_stream() - )) - })?; - super::shared_utils::ElementTypePrefix::new(&cuda_elem_ty_str) - } pub(crate) fn get_cuda_tile_type_str(&self) -> Option { self.cuda_tile_ty_str.clone() } diff --git a/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs b/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs index 33ce9dce04..6e3564fe62 100644 --- a/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs +++ b/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs @@ -25,10 +25,66 @@ use syn::punctuated::Punctuated; use syn::visit_mut::VisitMut; use syn::{Expr, FnArg, GenericArgument, ItemFn, ItemImpl, Lit, Stmt, Token}; +/// The six aggregation operation families, with the names each one +/// contributes to generated code. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum AggregationKind { + Add, + And, + Or, + Xor, + Max, + Min, +} + +impl AggregationKind { + fn from_type_name(name: &str) -> Option { + Some(match name { + "AddTensor" => AggregationKind::Add, + "AndTensor" => AggregationKind::And, + "OrTensor" => AggregationKind::Or, + "XorTensor" => AggregationKind::Xor, + "MaxTensor" => AggregationKind::Max, + "MinTensor" => AggregationKind::Min, + _ => return None, + }) + } + /// The capability's Rust type name, as spelled in the entry wrapper. + fn tensor_type(self) -> &'static str { + match self { + AggregationKind::Add => "AddTensor", + AggregationKind::And => "AndTensor", + AggregationKind::Or => "OrTensor", + AggregationKind::Xor => "XorTensor", + AggregationKind::Max => "MaxTensor", + AggregationKind::Min => "MinTensor", + } + } + /// The DSL view constructor the entry wrapper calls for this family. + fn make_view(self) -> &'static str { + match self { + AggregationKind::Add => "make_add_tensor_view", + AggregationKind::And => "make_and_tensor_view", + AggregationKind::Or => "make_or_tensor_view", + AggregationKind::Xor => "make_xor_tensor_view", + AggregationKind::Max => "make_max_tensor_view", + AggregationKind::Min => "make_min_tensor_view", + } + } +} + #[derive(Clone, Debug, PartialEq, Eq)] enum TensorParamKind { Tensor, - MappedPartitionMut { map_shape: String }, + /// One of the aggregation capabilities (`&AddTensor<_, _>` through + /// `&MinTensor<_, _>`): the same whole-tensor Tile IR view and the same + /// launcher ABI as an immutable `&Tensor`, wrapped in the atomic-only + /// Rust capability. Never partitioned by the entry wrapper — the atomic + /// operation builds its own view from the capability. + Aggregation(AggregationKind), + MappedPartitionMut { + map_shape: String, + }, } struct TensorInput { @@ -383,6 +439,14 @@ impl TensorInput { ).unwrap(); statements.push(tensor_stmnt); } + TensorParamKind::Aggregation(kind) => { + let tensor = kind.tensor_type(); + let make_view = kind.make_view(); + let tensor_stmnt = syn::parse2::( + format!("let {var_name}: {tensor}<{element_type}, {shape_param}> = unsafe {{ {make_view}({final_ptr_var}, {final_shape_var}, {strides_array_var}, {token_var}) }};").parse().unwrap() + ).unwrap(); + statements.push(tensor_stmnt); + } TensorParamKind::MappedPartitionMut { map_shape } => { let tensor_var = format!("{var_name}_tensor"); let tensor_stmnt = syn::parse2::( @@ -430,6 +494,9 @@ fn get_tensor_param_kind(ty: &syn::Type) -> Result { }; match type_ident.to_string().as_str() { "Tensor" => Ok(TensorParamKind::Tensor), + name if AggregationKind::from_type_name(name).is_some() => Ok( + TensorParamKind::Aggregation(AggregationKind::from_type_name(name).unwrap()), + ), "MappedPartitionMut" => { let Some(map_shape) = type_generic_args.args.iter().nth(2) else { return SourceLocation::unknown().jit_error_result( diff --git a/cutile-rs/cutile-kernels/Cargo.toml b/cutile-rs/cutile-kernels/Cargo.toml index 9360902bb0..a38472db05 100644 --- a/cutile-rs/cutile-kernels/Cargo.toml +++ b/cutile-rs/cutile-kernels/Cargo.toml @@ -14,6 +14,7 @@ readme = "README.md" [dependencies] cutile = { workspace = true } cutile-compiler = { workspace = true } +cutile-ir = { workspace = true } [lints] workspace = true diff --git a/cutile-rs/cutile-kernels/src/aggregation.rs b/cutile-rs/cutile-kernels/src/aggregation.rs new file mode 100644 index 0000000000..bf246ac9c1 --- /dev/null +++ b/cutile-rs/cutile-kernels/src/aggregation.rs @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Scatter/reduce kernels over the typed aggregation capabilities. +//! +//! Each kernel writes through an aggregation parameter (`AddTensor`, +//! `MaxTensor`, or `OrTensor`), so overlapping destinations are the intended +//! case: repeated indices, duplicate bins, and shared segments compose +//! atomically. Kernels are generic over the element where the added value +//! comes from an input tensor; the element's sealed trait bounds the set +//! (`AtomicAdd` admits the floats and 32/64-bit integers, `AtomicOrd` and +//! `AtomicBitwise` the 32/64-bit integers), and for the comparisons the +//! element also selects the signed or unsigned Tile IR mode. Common +//! contract, stated once: +//! +//! - **Indices** are `u32` row/bin ids loaded from memory and must lie in +//! range for the destination; each access is bounds-checked on device. +//! - **Initialization** is the caller's: seed the destination with the +//! operation's identity for a pure reduction (`0`/`0.0` for add and OR, +//! the signedness-correct extreme for max), or any other seed for seeded +//! accumulation. +//! - **Aliasing**: the destination may not also be bound as an ordinary +//! input or output of the same launch, nor as a different aggregation +//! operation; the launcher rejects both before launch. +//! - **Float order**: atomic float addition fixes no order, so float results +//! may differ between runs; integer kernels are exact. +//! +//! Unsupported elements fail at the type level; there is no float maximum +//! and no atomic reduction for 8- or 16-bit integers: +//! +//! ```compile_fail +//! #[cutile::module] +//! mod rejected { +//! use cutile::core::*; +//! #[cutile::entry()] +//! fn max_f32(out: &MaxTensor) {} +//! } +//! ``` + +#[cutile::module] +pub mod scatter_add_module { + use cutile::core::*; + + /// `out[indices[i]] += values[i]` for every element `i`, atomically. + /// Repeated indices accumulate. Launch with grid `(values.len(), 1, 1)`. + #[cutile::entry()] + fn scatter_add( + indices: &Tensor, + values: &Tensor, + out: &AddTensor, + ) { + let pid: (i32, i32, i32) = get_tile_block_id(); + let id_tile: Tile = indices.partition(const_shape![1]).load([pid.0]); + let idc: Tile = bitcast(id_tile); + let idx: i32 = tile_to_scalar(idc.reshape(const_shape![])); + let v: Tile = values.partition(const_shape![1]).load([pid.0]); + let acc: AddPartition = out.partition(const_shape![1]); + let _t: Token = acc.atomic_add([idx], v); + } +} + +#[cutile::module] +pub mod embedding_backward_scatter_add_module { + use cutile::core::*; + + /// Accumulates row gradients into an embedding-table gradient: + /// `table_grad[token_ids[i], :] += grads[i, :]`, atomically, so repeated + /// token ids compose. `D` is the embedding width; grid `(rows, 1, 1)`. + #[cutile::entry()] + fn embedding_backward_scatter_add( + token_ids: &Tensor, + grads: &Tensor, + table_grad: &AddTensor, + ) { + let pid: (i32, i32, i32) = get_tile_block_id(); + let id_tile: Tile = token_ids.partition(const_shape![1]).load([pid.0]); + let idc: Tile = bitcast(id_tile); + let row: i32 = tile_to_scalar(idc.reshape(const_shape![])); + let g: Tile = grads.partition(const_shape![1, D]).load([pid.0, 0i32]); + let acc: AddPartition = table_grad.partition(const_shape![1, D]); + let _t: Token = acc.atomic_add([row, 0i32], g); + } +} + +#[cutile::module] +pub mod segment_sum_module { + use cutile::core::*; + + /// `out[segment_ids[i]] += values[i]`, atomically. An empty segment keeps + /// its seed. Launch with grid `(values.len(), 1, 1)`. + #[cutile::entry()] + fn segment_sum( + segment_ids: &Tensor, + values: &Tensor, + out: &AddTensor, + ) { + let pid: (i32, i32, i32) = get_tile_block_id(); + let seg_tile: Tile = segment_ids.partition(const_shape![1]).load([pid.0]); + let segc: Tile = bitcast(seg_tile); + let seg: i32 = tile_to_scalar(segc.reshape(const_shape![])); + let v: Tile = values.partition(const_shape![1]).load([pid.0]); + let acc: AddPartition = out.partition(const_shape![1]); + let _t: Token = acc.atomic_add([seg], v); + } +} + +#[cutile::module] +pub mod histogram_module { + use cutile::core::*; + + /// `bins[bin_ids[i]] += 1` for every element, atomically. The increment + /// is `E::ONE` materialized on-device, so no weights buffer is read; + /// `weighted_histogram` is the variant that reads a weight per element. + /// Launch with grid `(bin_ids.len(), 1, 1)`. + #[cutile::entry()] + fn histogram( + bin_ids: &Tensor, + bins: &AddTensor, + ) { + let pid: (i32, i32, i32) = get_tile_block_id(); + let bin_tile: Tile = bin_ids.partition(const_shape![1]).load([pid.0]); + let binc: Tile = bitcast(bin_tile); + let bin: i32 = tile_to_scalar(binc.reshape(const_shape![])); + let one: Tile = constant(E::ONE, const_shape![1]); + let acc: AddPartition = bins.partition(const_shape![1]); + let _t: Token = acc.atomic_add([bin], one); + } +} + +#[cutile::module] +pub mod weighted_histogram_module { + use cutile::core::*; + + /// `bins[bin_ids[i]] += weights[i]`, atomically. The element-generic + /// histogram. Launch with grid `(bin_ids.len(), 1, 1)`. + #[cutile::entry()] + fn weighted_histogram( + bin_ids: &Tensor, + weights: &Tensor, + bins: &AddTensor, + ) { + let pid: (i32, i32, i32) = get_tile_block_id(); + let bin_tile: Tile = bin_ids.partition(const_shape![1]).load([pid.0]); + let binc: Tile = bitcast(bin_tile); + let bin: i32 = tile_to_scalar(binc.reshape(const_shape![])); + let w: Tile = weights.partition(const_shape![1]).load([pid.0]); + let acc: AddPartition = bins.partition(const_shape![1]); + let _t: Token = acc.atomic_add([bin], w); + } +} + +#[cutile::module] +pub mod max_reduce_module { + use cutile::core::*; + + /// Folds each row of `B` values into `out[0]` by atomic maximum. The + /// element selects the comparison for both the intra-tile reduction and + /// the atomic: signed `Max` for `i32`/`i64`, unsigned `Umax` for + /// `u32`/`u64`. Seed `out[0]` with the signedness-correct minimum + /// (`i32::MIN`, or `0` for unsigned) for a pure maximum. + /// + /// `values` is viewed as `[rows, B]`; `B` is inferred at launch. Grid + /// `(rows, 1, 1)`. + #[cutile::entry()] + fn max_reduce( + values: &Tensor, + out: &MaxTensor, + ) { + let pid: (i32, i32, i32) = get_tile_block_id(); + let row: Tile = values.partition(const_shape![1, B]).load([pid.0, 0i32]); + let m: Tile = reduce_max(row, 1i32); + let acc: MaxPartition = out.partition(const_shape![1]); + let _t: Token = acc.atomic_max([0i32], m); + } +} + +#[cutile::module] +pub mod bitmask_or_module { + use cutile::core::*; + + /// `out[dest_ids[i]] |= masks[i]`, atomically. OR is idempotent, so + /// repeated identical masks are harmless. Launch with grid + /// `(masks.len(), 1, 1)`. + #[cutile::entry()] + fn bitmask_or( + dest_ids: &Tensor, + masks: &Tensor, + out: &OrTensor, + ) { + let pid: (i32, i32, i32) = get_tile_block_id(); + let id_tile: Tile = dest_ids.partition(const_shape![1]).load([pid.0]); + let idc: Tile = bitcast(id_tile); + let dest: i32 = tile_to_scalar(idc.reshape(const_shape![])); + let mask: Tile = masks.partition(const_shape![1]).load([pid.0]); + let acc: OrPartition = out.partition(const_shape![1]); + let _t: Token = acc.atomic_or([dest], mask); + } +} + +pub use bitmask_or_module::bitmask_or; +pub use embedding_backward_scatter_add_module::embedding_backward_scatter_add; +pub use histogram_module::histogram; +pub use max_reduce_module::max_reduce; +pub use scatter_add_module::scatter_add; +pub use segment_sum_module::segment_sum; +pub use weighted_histogram_module::weighted_histogram; diff --git a/cutile-rs/cutile-kernels/src/lib.rs b/cutile-rs/cutile-kernels/src/lib.rs index 51a520a75a..a61ceac699 100644 --- a/cutile-rs/cutile-kernels/src/lib.rs +++ b/cutile-rs/cutile-kernels/src/lib.rs @@ -12,6 +12,7 @@ //! transformer utilities. Raw-pointer and benchmark-oriented kernels live under //! `experimental`. +pub mod aggregation; pub mod argmax; pub mod attention; pub mod embeddings; @@ -21,6 +22,10 @@ pub mod norms; pub mod pointwise; pub mod positional; +pub use aggregation::{ + bitmask_or, embedding_backward_scatter_add, histogram, max_reduce, scatter_add, segment_sum, + weighted_histogram, +}; pub use argmax::{argmax_blocks_f16, argmax_reduce_blocks_to_u32, lm_head_argmax_blocks_f16}; pub use attention::{ flash_attn_causal_seq_dynpos_f16, flash_attn_causal_seq_f16, fmha_causal, diff --git a/cutile-rs/cutile-kernels/tests/aggregation.rs b/cutile-rs/cutile-kernels/tests/aggregation.rs new file mode 100644 index 0000000000..8aa081216e --- /dev/null +++ b/cutile-rs/cutile-kernels/tests/aggregation.rs @@ -0,0 +1,369 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-only coverage for the aggregation kernels, plus host reference +//! implementations the GPU tests compare against. + +use cutile::compile_api::KernelCompiler; +use cutile_compiler::ast::Module as ModuleAst; +use cutile_ir::bytecode::Opcode; +use cutile_ir::ir::{Attribute, Module, OpId, Type}; + +const TEST_STACK_SIZE: usize = 32_000_000; +const TARGET: &str = "sm_120"; + +const MODE_OR: i64 = 1; +const MODE_ADD: i64 = 3; +const MODE_ADD_F: i64 = 4; +const MODE_MAX: i64 = 5; +const MODE_UMAX: i64 = 7; +const ORDERING_RELAXED: i64 = 1; +const SCOPE_DEVICE: i64 = 1; + +struct Case { + name: &'static str, + module_ast: fn() -> ModuleAst, + module_name: &'static str, + function_name: &'static str, + generics: Vec, + strides: Vec<(&'static str, &'static [i32])>, + mode: i64, +} + +const UNIT: &[i32] = &[1]; +const ROW: &[i32] = &[-1, 1]; + +fn modes_for(elem: &str) -> (i64, i64) { + // (AddF/Add, Max/UMax) by element. + let add = if matches!(elem, "f16" | "bf16" | "f32" | "f64") { + MODE_ADD_F + } else { + MODE_ADD + }; + let max = if matches!(elem, "u32" | "u64") { + MODE_UMAX + } else { + MODE_MAX + }; + (add, max) +} + +const ADD_ELEMS: &[&str] = &["f16", "bf16", "f32", "f64", "i32", "u32", "i64", "u64"]; +const INT_ELEMS: &[&str] = &["i32", "u32", "i64", "u64"]; + +/// Every generic kernel, driven across every element of its trait, with the +/// mode each (kernel, element) pair must emit. One `Case` per instantiation. +fn cases() -> Vec { + use cutile_kernels::aggregation::*; + let mut v = Vec::new(); + for &e in ADD_ELEMS { + let (add, _) = modes_for(e); + v.push(Case { + name: "scatter_add", + module_ast: scatter_add_module::__module_ast_self, + module_name: "scatter_add_module", + function_name: "scatter_add", + generics: vec![e.to_string()], + strides: vec![("indices", UNIT), ("values", UNIT), ("out", UNIT)], + mode: add, + }); + v.push(Case { + name: "segment_sum", + module_ast: segment_sum_module::__module_ast_self, + module_name: "segment_sum_module", + function_name: "segment_sum", + generics: vec![e.to_string()], + strides: vec![("segment_ids", UNIT), ("values", UNIT), ("out", UNIT)], + mode: add, + }); + v.push(Case { + name: "histogram", + module_ast: histogram_module::__module_ast_self, + module_name: "histogram_module", + function_name: "histogram", + generics: vec![e.to_string()], + strides: vec![("bin_ids", UNIT), ("bins", UNIT)], + mode: add, + }); + v.push(Case { + name: "weighted_histogram", + module_ast: weighted_histogram_module::__module_ast_self, + module_name: "weighted_histogram_module", + function_name: "weighted_histogram", + generics: vec![e.to_string()], + strides: vec![("bin_ids", UNIT), ("weights", UNIT), ("bins", UNIT)], + mode: add, + }); + v.push(Case { + name: "embedding_backward_scatter_add", + module_ast: embedding_backward_scatter_add_module::__module_ast_self, + module_name: "embedding_backward_scatter_add_module", + function_name: "embedding_backward_scatter_add", + generics: vec![e.to_string(), "8".to_string()], + strides: vec![("token_ids", UNIT), ("grads", ROW), ("table_grad", ROW)], + mode: add, + }); + } + for &e in INT_ELEMS { + let (_, max) = modes_for(e); + v.push(Case { + name: "max_reduce", + module_ast: max_reduce_module::__module_ast_self, + module_name: "max_reduce_module", + function_name: "max_reduce", + generics: vec![e.to_string(), "64".to_string()], + strides: vec![("values", ROW), ("out", UNIT)], + mode: max, + }); + v.push(Case { + name: "bitmask_or", + module_ast: bitmask_or_module::__module_ast_self, + module_name: "bitmask_or_module", + function_name: "bitmask_or", + generics: vec![e.to_string()], + strides: vec![("dest_ids", UNIT), ("masks", UNIT), ("out", UNIT)], + mode: MODE_OR, + }); + } + v +} + +fn ops(module: &Module) -> Vec { + fn walk(module: &Module, op: OpId, out: &mut Vec) { + out.push(op); + for ®ion in &module.op(op).regions { + for &block in &module.region(region).blocks { + for &inner in &module.block(block).ops { + walk(module, inner, out); + } + } + } + } + let mut out = Vec::new(); + for &function in &module.functions { + walk(module, function, &mut out); + } + out +} + +fn int_attr(module: &Module, op: OpId, name: &str) -> i64 { + module + .op(op) + .attributes + .iter() + .find_map(|(n, attr)| match attr { + Attribute::Integer(v, _) if n == name => Some(*v), + _ => None, + }) + .unwrap_or_else(|| panic!("missing integer attribute `{name}`")) +} + +/// Every kernel lowers to exactly one atomic reduction with its intended +/// mode (signedness included), relaxed ordering, device scope, and a token +/// as the only result; indices loaded from memory keep a check, recorded in +/// the placement counts rather than silently dropped. +#[test] +fn every_kernel_reaches_its_aggregation_mode() { + std::thread::Builder::new() + .stack_size(TEST_STACK_SIZE) + .spawn(|| { + for case in cases() { + let artifacts = + KernelCompiler::new(case.module_ast, case.module_name, case.function_name) + .target(TARGET) + .generics( + case.generics + .iter() + .map(|s| s.to_string()) + .collect::>(), + ) + .strides(&case.strides) + .compile() + .unwrap_or_else(|e| panic!("{}: {e}", case.name)); + let counts = artifacts.check_counts(); + let module = artifacts.into_module(); + let reductions: Vec = ops(&module) + .into_iter() + .filter(|&id| module.op(id).opcode == Opcode::AtomicRedViewTko) + .collect(); + assert_eq!(reductions.len(), 1, "{}: one reduction", case.name); + let op = reductions[0]; + assert_eq!( + int_attr(&module, op, "mode"), + case.mode, + "{}: mode", + case.name + ); + assert_eq!( + int_attr(&module, op, "memory_ordering_semantics"), + ORDERING_RELAXED, + "{}: ordering", + case.name + ); + assert_eq!( + int_attr(&module, op, "memory_scope"), + SCOPE_DEVICE, + "{}: scope", + case.name + ); + assert_eq!( + module.op(op).result_types, + vec![Type::Token], + "{}: token result", + case.name + ); + let placed = counts.discharged + counts.hoisted + counts.in_place; + assert!( + placed > 0, + "{}: no access check was placed at all", + case.name + ); + } + }) + .expect("spawn") + .join() + .expect("test thread"); +} + +/// The signed and unsigned maxima must not collapse to one comparison. +#[test] +fn signed_and_unsigned_max_select_different_modes() { + assert_ne!(MODE_MAX, MODE_UMAX); + // Enforced per kernel by `every_kernel_reaches_its_aggregation_mode`; + // this test exists so a change collapsing the modes fails by name. + let modes: Vec = cases() + .iter() + .filter(|c| c.name.starts_with("max_reduce")) + .map(|c| c.mode) + .collect(); + assert_eq!(modes, vec![MODE_MAX, MODE_UMAX, MODE_MAX, MODE_UMAX]); +} + +// ── Host references, shared with the GPU tests ────────────────────────────── + +pub fn scatter_add_ref(indices: &[u32], values: &[f32], out: &mut [f32]) { + for (&i, &v) in indices.iter().zip(values) { + out[i as usize] += v; + } +} + +pub fn histogram_ref(bin_ids: &[u32], bins: &mut [u32]) { + for &b in bin_ids { + bins[b as usize] += 1; + } +} + +pub fn weighted_histogram_ref(bin_ids: &[u32], weights: &[f32], bins: &mut [f32]) { + for (&b, &w) in bin_ids.iter().zip(weights) { + bins[b as usize] += w; + } +} + +pub fn max_reduce_i32_ref(values: &[i32], seed: i32) -> i32 { + values.iter().copied().fold(seed, i32::max) +} + +pub fn max_reduce_u32_ref(values: &[u32], seed: u32) -> u32 { + values.iter().copied().fold(seed, u32::max) +} + +pub fn bitmask_or_ref(dest_ids: &[u32], masks: &[u32], out: &mut [u32]) { + for (&d, &m) in dest_ids.iter().zip(masks) { + out[d as usize] |= m; + } +} + +/// Small deterministic inputs: repeated indices, an empty segment, duplicate +/// bins, zeros, negatives for signed max, and values straddling the signed +/// midpoint for unsigned max. +#[test] +fn host_references_handle_the_edge_cases() { + // Repeated indices accumulate; untouched slots keep the identity seed. + let mut out = vec![0.0f32; 4]; + scatter_add_ref(&[1, 1, 3, 1], &[1.0, 2.0, 5.0, 0.0], &mut out); + assert_eq!(out, vec![0.0, 3.0, 0.0, 5.0]); + + // Segment 2 is empty and keeps its seed; segment sums share the scatter + // reference, which is the same contraction. + let mut segs = vec![0.0f32; 3]; + scatter_add_ref(&[0, 0, 1], &[2.0, -2.0, 7.0], &mut segs); + assert_eq!(segs, vec![0.0, 7.0, 0.0]); + + let mut bins = vec![0u32; 3]; + histogram_ref(&[2, 2, 2, 0], &mut bins); + assert_eq!(bins, vec![1, 0, 3]); + + let mut wbins = vec![0.0f32; 2]; + weighted_histogram_ref(&[1, 1, 0], &[0.5, 0.25, 0.0], &mut wbins); + assert_eq!(wbins, vec![0.0, 0.75]); + + // Negative values with the signed identity. + assert_eq!(max_reduce_i32_ref(&[-5, -2, -9], i32::MIN), -2); + assert_eq!(max_reduce_i32_ref(&[], i32::MIN), i32::MIN); + // Above the signed midpoint: the unsigned comparison must pick the + // high-bit value a signed one would call negative. + assert_eq!( + max_reduce_u32_ref(&[0x7FFF_FFF0, 0x8000_0010], 0), + 0x8000_0010 + ); + + let mut words = vec![0u32; 2]; + bitmask_or_ref(&[0, 0, 1], &[0b01, 0b10, 0b100], &mut words); + assert_eq!(words, vec![0b11, 0b100]); +} + +#[test] +fn reduce_max_uses_the_elements_signedness() { + std::thread::Builder::new() + .stack_size(TEST_STACK_SIZE) + .spawn(|| { + use cutile_kernels::aggregation::max_reduce_module; + let func = "max_reduce"; + for (elem, want_signed) in [ + ("i32", "signed"), + ("u32", "unsigned"), + ("i64", "signed"), + ("u64", "unsigned"), + ] { + let module = KernelCompiler::new( + max_reduce_module::__module_ast_self, + "max_reduce_module", + func, + ) + .target(TARGET) + .generics(vec![elem.to_string(), "64".to_string()]) + .strides(&[("values", ROW), ("out", UNIT)]) + .compile() + .unwrap_or_else(|e| panic!("{func}<{elem}>: {e}")) + .into_module(); + // Find the MaxI op (the reduce combinator) and read its + // signedness attribute. + let mut seen = None; + for id in ops(&module) { + let op = module.op(id); + if format!("{:?}", op.opcode).contains("MaxI") { + for (n, a) in &op.attributes { + if n == "signedness" { + if let Attribute::Integer(v, _) = a { + seen = Some( + if *v == 0 { "unsigned" } else { "signed" }.to_string(), + ); + } + } + } + } + } + let got = seen.unwrap_or_else(|| panic!("{func}<{elem}>: no MaxI signedness attr")); + eprintln!("{func}<{elem}>: MaxI signedness = {got} (want {want_signed})"); + assert!( + got.to_lowercase().contains(want_signed), + "{func}<{elem}>: MaxI signedness {got}, expected {want_signed}" + ); + } + }) + .expect("spawn") + .join() + .expect("probe"); +} diff --git a/cutile-rs/cutile-kernels/tests/gpu.rs b/cutile-rs/cutile-kernels/tests/gpu.rs new file mode 100644 index 0000000000..62a9b8217b --- /dev/null +++ b/cutile-rs/cutile-kernels/tests/gpu.rs @@ -0,0 +1,219 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Device tests for the aggregation kernels, each against the host reference +//! from `tests/aggregation.rs`'s edge-case inputs: repeated indices, an empty +//! segment, duplicate bins, negatives for signed max, and values above the +//! signed midpoint for unsigned max. +//! +//! Skips print a `SKIP Tile IR execution:` line, so an evidence harness can +//! refuse a silently green run. + +use cutile::prelude::*; +use std::sync::Arc; + +/// Runtime prerequisites only; mirrors the cutile gpu suite's gate. +fn supports_tile_ir() -> bool { + use cutile_compiler::cuda_tile_runtime_utils::ToolkitCapabilities; + use cutile_ir::ir::{Location, Module}; + let caps = ToolkitCapabilities::for_device(0).expect("selected toolkit capabilities"); + let result = caps + .target + .validate_module(&Module::new("test_prerequisites")) + .and_then(|()| { + cutile_ir::requirements::Feature::AtomicRedView.check(&caps.target, &Location::Unknown) + }); + if let Err(error) = result { + eprintln!("SKIP Tile IR execution: {error}"); + return false; + } + true +} + +fn device_u32(host: Vec) -> Arc> { + let n = host.len(); + api::copy_host_vec_to_device(&Arc::new(host)) + .reshape(&[n]) + .sync() + .expect("to device") + .into() +} + +fn device_f32(host: Vec) -> Arc> { + let n = host.len(); + api::copy_host_vec_to_device(&Arc::new(host)) + .reshape(&[n]) + .sync() + .expect("to device") + .into() +} + +#[test] +fn scatter_add_matches_the_reference() { + if !supports_tile_ir() { + return; + } + let indices = vec![1u32, 1, 3, 1, 0, 3]; + let values = vec![1.0f32, 2.0, 5.0, 0.5, -1.0, 1.5]; + let mut expected = vec![0.0f32; 4]; + for (&i, &v) in indices.iter().zip(&values) { + expected[i as usize] += v; + } + let n = indices.len() as u32; + let mut out = api::zeros::(&[4]).sync().expect("zeros"); + let _ = cutile_kernels::scatter_add(device_u32(indices), device_f32(values), out.atomic()) + .grid((n, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, expected); +} + +#[test] +fn segment_sum_keeps_an_empty_segments_seed() { + if !supports_tile_ir() { + return; + } + // Segment 1 is empty and must keep the identity seed. + let segment_ids = vec![0u32, 0, 2, 2, 2]; + let values = vec![2.0f32, -2.0, 1.0, 1.0, 1.0]; + let n = segment_ids.len() as u32; + let mut out = api::zeros::(&[3]).sync().expect("zeros"); + let _ = cutile_kernels::segment_sum(device_u32(segment_ids), device_f32(values), out.atomic()) + .grid((n, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, vec![0.0, 0.0, 3.0]); +} + +#[test] +fn histogram_counts_duplicate_bins_exactly() { + if !supports_tile_ir() { + return; + } + let bin_ids = vec![2u32, 2, 2, 0, 1, 2]; + let n = bin_ids.len() as u32; + let mut bins = api::zeros::(&[3]).sync().expect("zeros"); + let _ = cutile_kernels::histogram(device_u32(bin_ids), bins.atomic()) + .grid((n, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = bins.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, vec![1, 1, 4]); +} + +#[test] +fn weighted_histogram_accumulates_weights() { + if !supports_tile_ir() { + return; + } + let bin_ids = vec![1u32, 1, 0, 1]; + let weights = vec![0.5f32, 0.25, 0.0, 0.25]; + let n = bin_ids.len() as u32; + let mut bins = api::zeros::(&[2]).sync().expect("zeros"); + let _ = + cutile_kernels::weighted_histogram(device_u32(bin_ids), device_f32(weights), bins.atomic()) + .grid((n, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = bins.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, vec![0.0, 1.0]); +} + +#[test] +fn signed_max_reduce_handles_negatives() { + if !supports_tile_ir() { + return; + } + // All negative, seeded with the signed identity: the answer is the + // largest (least negative) value, which an unsigned comparison would not + // pick. + let values: Vec = (0..64).map(|i| -1 - (i % 7)).collect(); + let dev: Arc> = api::copy_host_vec_to_device(&Arc::new(values)) + .reshape(&[1, 64]) + .sync() + .expect("to device") + .into(); + let mut out = api::full::(i32::MIN, &[1]).sync().expect("seed"); + let _ = cutile_kernels::max_reduce(dev, out.atomic()) + .grid((1, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, vec![-1]); +} + +#[test] +fn unsigned_max_reduce_crosses_the_signed_midpoint() { + if !supports_tile_ir() { + return; + } + let mut values = vec![3u32; 64]; + values[17] = 0x7FFF_FFF0; + values[41] = 0x8000_0010; // negative under a signed comparison + let mut out = api::zeros::(&[1]).sync().expect("seed"); + let values_dev: Arc> = api::copy_host_vec_to_device(&Arc::new(values)) + .reshape(&[1, 64]) + .sync() + .expect("to device") + .into(); + let _ = cutile_kernels::max_reduce(values_dev, out.atomic()) + .grid((1, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, vec![0x8000_0010]); +} + +#[test] +fn bitmask_or_merges_duplicate_destinations() { + if !supports_tile_ir() { + return; + } + let dest_ids = vec![0u32, 0, 1, 0]; + let masks = vec![0b01u32, 0b10, 0b100, 0b01]; + let n = masks.len() as u32; + let mut out = api::zeros::(&[2]).sync().expect("zeros"); + let _ = cutile_kernels::bitmask_or(device_u32(dest_ids), device_u32(masks), out.atomic()) + .grid((n, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, vec![0b11, 0b100]); +} + +#[test] +fn embedding_backward_accumulates_repeated_token_ids() { + if !supports_tile_ir() { + return; + } + const D: usize = 8; + let token_ids = vec![1u32, 1, 0, 1]; + let rows = token_ids.len(); + let grads: Vec = (0..rows * D).map(|i| (i % 5) as f32 - 2.0).collect(); + let mut expected = vec![0.0f32; 3 * D]; + for (r, &t) in token_ids.iter().enumerate() { + for d in 0..D { + expected[t as usize * D + d] += grads[r * D + d]; + } + } + let grads_dev: Arc> = api::copy_host_vec_to_device(&Arc::new(grads)) + .reshape(&[rows, D]) + .sync() + .expect("grads") + .into(); + let mut table = api::zeros::(&[3, D]).sync().expect("zeros"); + let _ = cutile_kernels::embedding_backward_scatter_add( + device_u32(token_ids), + grads_dev, + table.atomic(), + ) + .grid((rows as u32, 1, 1)) + .sync() + .expect("launch"); + let host: Vec = table.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host, expected); +} diff --git a/cutile-rs/cutile-macro/src/_module.rs b/cutile-rs/cutile-macro/src/_module.rs index a601b6cda3..612a22b4ce 100644 --- a/cutile-rs/cutile-macro/src/_module.rs +++ b/cutile-rs/cutile-macro/src/_module.rs @@ -399,7 +399,7 @@ fn module_inner( use #tile_rust_crate_root::error::{*}; use #tile_rust_crate_root::DType; use #tile_rust_crate_root::{tensor}; - use #tile_rust_crate_root::tensor::{KernelInput, KernelInputStored, KernelOutput, KernelOutputStored, SpecializationBits}; + use #tile_rust_crate_root::tensor::{KernelAtomic, KernelInput, KernelInputStored, KernelOutput, KernelOutputStored, SpecializationBits}; use #tile_rust_crate_root::tile_kernel::{*}; use #tile_rust_crate_root::cuda_async::error::{*}; use #tile_rust_crate_root::cuda_async::scheduling_policies::SchedulingPolicy; @@ -880,7 +880,11 @@ pub fn kernel_launcher( for (ki_idx, ki_name) in kernel_input_info.type_param_names.iter().enumerate() { let elem = &kernel_input_info.element_type_names[ki_idx]; struct_generics.params.push( - syn::parse_str::(&format!("{ki_name}: KernelInput<{elem}>")).unwrap(), + syn::parse_str::(&format!( + "{ki_name}: {}<{elem}>", + kernel_input_info.bound_names[ki_idx] + )) + .unwrap(), ); } for (ko_idx, ko_name) in kernel_input_info.ko_type_param_names.iter().enumerate() { diff --git a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs index 27dac6ab6c..3e2705aa23 100644 --- a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs +++ b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs @@ -80,6 +80,10 @@ //! //! - `&mut Tensor` → `impl KernelOutput` (accepts `Partition>` or `Partition<&mut Tensor>`) //! - `&Tensor` → `impl KernelInput` (accepts `Tensor`, `Arc>`, or `&Tensor`) +//! - `&AddTensor` (and the other aggregation capabilities) → `impl KernelAtomic` (accepts only `Atomic<&mut Tensor>` +//! or `Atomic>`: exclusive host authority). Launches with the +//! `&Tensor` ABI and contributes no grid bound, so several programs may +//! target one output tile. //! - `*mut T` → `DevicePointer` (unsafe only) //! - Scalars remain unchanged //! @@ -89,7 +93,9 @@ //! ## Grid Inference //! //! Launch grid dimensions are automatically inferred from partitioned tensors. -//! If multiple partitions exist, their grids must match. +//! If multiple partitions exist, their grids must match. Atomic parameters +//! bound nothing, because overlap is the point. A kernel whose only +//! destination is atomic needs an explicit `.grid(..)`. use cutile_compiler::kernel_naming::KernelNaming; use cutile_compiler::syn_utils::*; @@ -106,6 +112,58 @@ use syn::{ use crate::error::{Error, SpannedError}; +/// The six aggregation operation families, by capability type name. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum AggregationKind { + Add, + And, + Or, + Xor, + Max, + Min, +} + +impl AggregationKind { + const ALL: [AggregationKind; 6] = [ + AggregationKind::Add, + AggregationKind::And, + AggregationKind::Or, + AggregationKind::Xor, + AggregationKind::Max, + AggregationKind::Min, + ]; + fn tensor_type(self) -> &'static str { + match self { + AggregationKind::Add => "AddTensor", + AggregationKind::And => "AndTensor", + AggregationKind::Or => "OrTensor", + AggregationKind::Xor => "XorTensor", + AggregationKind::Max => "MaxTensor", + AggregationKind::Min => "MinTensor", + } + } + fn op_variant(self) -> &'static str { + match self { + AggregationKind::Add => "Add", + AggregationKind::And => "And", + AggregationKind::Or => "Or", + AggregationKind::Xor => "Xor", + AggregationKind::Max => "Max", + AggregationKind::Min => "Min", + } + } +} + +/// The aggregation family of `&AddTensor<..>` through `&MinTensor<..>`, or +/// `None` for any other reference type. +fn aggregation_tensor_kind(ref_ty: &TypeReference) -> Option { + let (ident, _) = get_ident_generic_args(&Type::Reference(ref_ty.clone())); + let ident = ident?; + AggregationKind::ALL + .into_iter() + .find(|kind| ident == kind.tensor_type()) +} + /// Classification of generic parameter types. /// /// Used to determine how to handle and infer different kinds of generic parameters. @@ -540,6 +598,10 @@ pub fn generate_kernel_launcher( // Track element type per param (Some for &Tensor params, None for others). let mut param_element_types: Vec> = vec![]; let mut tensor_param_kinds: HashMap = HashMap::new(); + // One `(name, allocation identity, access)` triple per tensor parameter, + // in signature order. Only emitted into the launcher when some parameter + // is atomic — ordinary kernels keep byte-identical generated code. + let mut alias_binding_exprs: Vec = vec![]; let mut required_generics: RequiredGenerics = RequiredGenerics::new(&item.sig.generics); for (i, ty) in input_types.iter().enumerate() { @@ -549,10 +611,13 @@ pub fn generate_kernel_launcher( let res = get_tensor_code(i, var_name, ref_ty, &mut required_generics)?; let tensor_kind = if ref_ty.mutability.is_some() { TensorParamKind::Output + } else if let Some(kind) = aggregation_tensor_kind(ref_ty) { + TensorParamKind::Aggregation { kind } } else { TensorParamKind::Input }; tensor_param_kinds.insert(var_name.clone(), tensor_kind); + alias_binding_exprs.push(alias_binding_expr(var_name, tensor_kind)); arg_types.push(res.fn_arg.ty.as_ref().clone()); param_element_types.push(res.element_type_name); stride_args.push(res.stride_expr_str); @@ -569,6 +634,7 @@ pub fn generate_kernel_launcher( let ref_ty: TypeReference = parse_quote!(&mut #path_ty); let res = get_tensor_code(i, var_name, &ref_ty, &mut required_generics)?; tensor_param_kinds.insert(var_name.clone(), TensorParamKind::Output); + alias_binding_exprs.push(alias_binding_expr(var_name, TensorParamKind::Output)); arg_types.push(res.fn_arg.ty.as_ref().clone()); param_element_types.push(res.element_type_name); stride_args.push(res.stride_expr_str); @@ -678,15 +744,33 @@ pub fn generate_kernel_launcher( let mut ki_type_param_names: Vec = vec![]; let mut ki_element_type_names: Vec = vec![]; let mut ki_param_idx: Vec> = vec![]; + // The trait each KernelInput-family param is bounded by: `KernelInput` + // for an ordinary `&Tensor`, `KernelAtomic` for an aggregation capability. The + // plumbing (`Stored`, `Returned`, `prepare`, `recover`, `PrepareInput`) + // is shared through the supertrait; only the bound differs, which is + // what keeps a plain `Arc>` out of an atomic parameter. + let mut ki_bound_names: Vec = vec![]; // Build KernelOutput metadata: which params are &mut Tensor (Partition). let mut ko_type_param_names: Vec = vec![]; let mut ko_element_type_names: Vec = vec![]; let mut ko_param_idx: Vec> = vec![]; for (i, ty) in arg_types.iter().enumerate() { let ty_str = ty.to_token_stream().to_string(); - if ty_str.starts_with("Arc <") { + let is_atomic = matches!( + tensor_param_kinds.get(¶m_names[i]), + Some(TensorParamKind::Aggregation { .. }) + ); + if ty_str.starts_with("Arc <") || is_atomic { let idx = ki_type_param_names.len(); ki_type_param_names.push(format!("_K{}", idx)); + ki_bound_names.push( + if is_atomic { + "KernelAtomic" + } else { + "KernelInput" + } + .to_string(), + ); ki_element_type_names.push( param_element_types[i] .clone() @@ -735,6 +819,7 @@ pub fn generate_kernel_launcher( let kernel_input_info = KernelInputInfo { type_param_names: ki_type_param_names, element_type_names: ki_element_type_names, + bound_names: ki_bound_names, param_kernel_input_idx: ki_param_idx.clone(), ko_type_param_names: ko_type_param_names.clone(), ko_element_type_names: ko_element_type_names.clone(), @@ -798,9 +883,10 @@ pub fn generate_kernel_launcher( let mut struct_generics = generic_params.clone(); for (ki_idx, ki_name) in ki_info.type_param_names.iter().enumerate() { let elem = &ki_info.element_type_names[ki_idx]; - struct_generics.params.push( - syn::parse_str::(&format!("{ki_name}: KernelInput<{elem}>")).unwrap(), - ); + let bound = &ki_info.bound_names[ki_idx]; + struct_generics + .params + .push(syn::parse_str::(&format!("{ki_name}: {bound}<{elem}>")).unwrap()); } for (ko_idx, ko_name) in ki_info.ko_type_param_names.iter().enumerate() { let elem = &ki_info.ko_element_type_names[ko_idx]; @@ -1037,6 +1123,19 @@ pub fn generate_kernel_launcher( )); launcher_method.block.stmts.extend(validator_statements); + // One allocation bound both atomically and through an ordinary access in + // the same launch is a data race the kernel cannot order away; reject it + // before any GPU work. Ordinary kernels emit nothing here. + if tensor_param_kinds + .values() + .any(|kind| matches!(kind, TensorParamKind::Aggregation { .. })) + { + launcher_method.block.stmts.push(parse_stmt(format!( + "validate_binding_aliases(&[{}])?;", + alias_binding_exprs.join(",") + ))); + } + // Above the `!_compile_only` gate so warmup validates the grid too. launcher_method.block.stmts.push(parse_stmt( "if self._programmatic_dependent_launch { validate_programmatic_dependent_launch(ctx.get_device_id())?; }".to_string(), @@ -1166,7 +1265,13 @@ pub fn generate_kernel_launcher( let function_param = format!("arg{}", i); let type_param_name = format!("_A{}", i); let arg_type_str = &arg_aliases[i]; - let is_arc = arg_type_str.starts_with("Arc <"); + // Both families take the `PrepareInput` path; `KernelAtomic` only + // narrows which host bindings the where clause accepts. + let is_arc = arg_type_str.starts_with("Arc <") + || matches!( + tensor_param_kinds.get(¶m_names[i]), + Some(TensorParamKind::Aggregation { .. }) + ); is_arc_param.push(is_arc); launcher_direct.sig.inputs.push(FnArg::Typed( @@ -1208,9 +1313,10 @@ pub fn generate_kernel_launcher( ); let ki_idx = ki_param_idx[i].unwrap(); let elem = &ki_info.element_type_names[ki_idx]; + let bound = &ki_info.bound_names[ki_idx]; where_clause.predicates.push( syn::parse2::( - format!("{}: KernelInput<{}>", intermediate_type, elem) + format!("{}: {}<{}>", intermediate_type, bound, elem) .parse() .unwrap(), ) @@ -1345,7 +1451,7 @@ pub fn generate_kernel_launcher( /// /// A parsed `Stmt` AST node fn parse_stmt(s: String) -> Stmt { - syn::parse::(s.parse().unwrap()).unwrap() + syn::parse_str::(&s).unwrap() } /// Parses a string into a Rust expression AST node. @@ -1366,7 +1472,7 @@ fn parse_stmt(s: String) -> Stmt { /// Currently unused (marked with `#[allow(dead_code)]`). #[allow(dead_code)] fn parse_expr(s: String) -> Expr { - syn::parse::(s.parse().unwrap()).unwrap() + syn::parse_str::(&s).unwrap() } /// Metadata about KernelInput params for the struct definition in _module.rs. @@ -1375,6 +1481,9 @@ pub struct KernelInputInfo { pub type_param_names: Vec, /// Element type name for each KernelInput param, e.g., ["T", "SrcType"]. pub element_type_names: Vec, + /// The trait each KernelInput-family param is bounded by: `"KernelInput"` + /// for `&Tensor`, `"KernelAtomic"` for an aggregation capability. + pub bound_names: Vec, /// For each param in the args tuple, the index into type_param_names (if it's /// a KernelInput param), or None (if it's a partition/scalar). pub param_kernel_input_idx: Vec>, @@ -1416,6 +1525,12 @@ struct TensorLaunchCode { enum TensorParamKind { Input, Output, + /// An aggregation capability: bound with exclusive host authority, launched + /// with the whole-tensor input ABI, and constraining the launch grid not + /// at all — several programs may target one output tile on purpose. + Aggregation { + kind: AggregationKind, + }, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -1600,10 +1715,32 @@ fn parse_precondition_axis_arg(expr: &Expr) -> Result { .map_err(|err| crate::error::syn_err(lit_int.span(), &err.to_string())) } +/// `(name, allocation identity, access)` for one tensor parameter, as the +/// generated launcher spells it. +fn alias_binding_expr(name: &str, kind: TensorParamKind) -> String { + match kind { + TensorParamKind::Input => { + format!(r#"("{name}", KernelInputStored::storage_id(&{name}), BindingAccess::Read)"#) + } + TensorParamKind::Aggregation { kind } => { + let op = kind.op_variant(); + format!( + r#"("{name}", KernelInputStored::storage_id(&{name}), BindingAccess::Atomic(AggregationOp::{op}))"# + ) + } + TensorParamKind::Output => { + format!(r#"("{name}", KernelOutputStored::storage_id(&{name}), BindingAccess::Write)"#) + } + } +} + fn precondition_shape_expr(name: &str, kind: TensorParamKind) -> TokenStream2 { let ident = Ident::new(name, Span::call_site()); match kind { - TensorParamKind::Input => quote! { #ident.shape().to_vec() }, + // Both bind the whole tensor; `shape()` is the same trait method. + TensorParamKind::Input | TensorParamKind::Aggregation { .. } => { + quote! { #ident.shape().to_vec() } + } TensorParamKind::Output => quote! { KernelOutputStored::shape_as_i32(&#ident) }, } } @@ -1743,9 +1880,26 @@ fn get_tensor_code( return ty.err("Expected a named type identifier for tensor parameter."); }; let type_ident = type_ident.to_string(); - if type_ident != "Tensor" && type_ident != "MappedPartitionMut" { + // An aggregation capability takes the immutable launch path throughout: + // the same pointer/shape/strides ABI, the same whole-tensor validation, + // and — deliberately — no launch-grid bound. Only the host binding type + // differs. + let is_atomic = AggregationKind::ALL + .into_iter() + .any(|kind| type_ident == kind.tensor_type()); + if type_ident != "Tensor" && type_ident != "MappedPartitionMut" && !is_atomic { return ty.err(&format!("Expected Tensor type, got {}.", type_ident)); } + // The capability is shared by construction: every program in the launch + // holds it and the atomics resolve the collisions. `&mut` would promise + // the disjointness the operation deliberately does not have. + if is_atomic && ty.mutability.is_some() { + return ty.err(&format!( + "`&mut {type_ident}<..>` is not a kernel parameter type. Atomic \ + updates go through a shared `&{type_ident}<..>`, like Rust's own \ + atomics.", + )); + } if type_ident == "MappedPartitionMut" && ty.mutability.is_none() { return ty.err("internal error: MappedPartitionMut launcher lowering must use mutable output metadata."); } @@ -1785,11 +1939,12 @@ fn get_tensor_code( .to_string(); let tensor_type = if ty.mutability.is_some() { format!("tensor::Partition>") + } else if is_atomic { + format!("tensor::Atomic>") } else { format!("Arc>") }; - let fn_arg = - syn::parse::(format!("{var_name}: {tensor_type}").parse().unwrap()).unwrap(); + let fn_arg = syn::parse_str::(&format!("{var_name}: {tensor_type}")).unwrap(); // Stride expr. let stride_expr_str = if ty.mutability.is_some() { format!( @@ -2156,3 +2311,152 @@ pub fn infer_shape_params_from_tensor_type( } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + /// The generated launcher body for one kernel signature, as a string. + fn launcher_for(item: ItemFn) -> String { + let root = Ident::new("cutile", Span::call_site()); + let (_generics, _arg_types, launcher, _info) = + generate_kernel_launcher(&item, "m", "k", "k_entry", "K", "KArgs", &root) + .expect("launcher generation failed"); + launcher.to_string() + } + + /// An `&mut Tensor` output binds a partition, and a partition defines the + /// launch grid exactly. This is the behaviour the atomic parameter must + /// *not* inherit; asserted here so the contrast below is meaningful. + #[test] + fn a_mutable_output_contributes_an_exact_grid_bound() { + let launcher = launcher_for(parse_quote! { + #[cutile::entry()] + fn k(out: &mut Tensor) {} + }); + assert!( + launcher.contains("grid_bound"), + "a `&mut Tensor` output must bound the launch grid: {launcher}" + ); + } + + /// An atomic kernel may put several programs on one output tile on + /// purpose, so the capability contributes no grid bound: the generated + /// `infer_launch_grid` call gets an empty bound list, and the grid comes + /// from elsewhere (another binding, or an explicit `.grid(..)`). + #[test] + fn an_atomic_parameter_contributes_no_grid_bound() { + let launcher = launcher_for(parse_quote! { + #[cutile::entry()] + fn k(out: &AddTensor, src: &Tensor) {} + }); + assert!( + launcher.contains("infer_launch_grid (& [])"), + "the atomic capability must impose no grid bound: {launcher}" + ); + assert!( + !launcher.contains("grid_bound"), + "no binding in this kernel may bound the grid: {launcher}" + ); + } + + /// The atomic parameter is bound by `KernelAtomic`, not `KernelInput`: + /// that is what keeps a shared `Arc>` out of it while the + /// ordinary input next to it still accepts one. + #[test] + fn an_atomic_parameter_is_bounded_by_kernel_atomic() { + let launcher = launcher_for(parse_quote! { + #[cutile::entry()] + fn k(out: &AddTensor, src: &Tensor) {} + }); + assert!( + launcher.contains("_K0 : KernelAtomic < f32 >"), + "the atomic parameter must carry the exclusive-authority bound: {launcher}" + ); + assert!( + launcher.contains("_K1 : KernelInput < f32 >"), + "the ordinary input must keep its own bound: {launcher}" + ); + assert!( + launcher.contains("tensor :: Atomic < tensor :: Tensor < f32 > >"), + "the atomic parameter must bind through `Atomic<..>`: {launcher}" + ); + } + + /// Alias validation is emitted only for kernels that have an atomic + /// parameter, and it records every tensor parameter's allocation + /// identity and access kind. + #[test] + fn alias_validation_is_emitted_only_for_atomic_kernels() { + let ordinary = launcher_for(parse_quote! { + #[cutile::entry()] + fn k(out: &mut Tensor, src: &Tensor) {} + }); + assert!( + !ordinary.contains("validate_binding_aliases"), + "an ordinary kernel must emit no alias validation: {ordinary}" + ); + + let atomic = launcher_for(parse_quote! { + #[cutile::entry()] + fn k( + out: &AddTensor, + dst: &mut Tensor, + src: &Tensor, + ) {} + }); + assert!( + atomic.contains("BindingAccess :: Atomic (AggregationOp :: Add)") + && atomic.contains("BindingAccess :: Write") + && atomic.contains("BindingAccess :: Read"), + "every tensor parameter's access kind and operation must be recorded: {atomic}" + ); + assert!( + atomic.contains("KernelInputStored :: storage_id (& out)") + && atomic.contains("KernelOutputStored :: storage_id (& dst)") + && atomic.contains("KernelInputStored :: storage_id (& src)"), + "every tensor parameter's allocation identity must be recorded: {atomic}" + ); + } + + /// Every capability family's binding carries its own operation identity, + /// so the launch-time alias check can tell two different aggregations of + /// one allocation apart. + #[test] + fn every_aggregation_family_records_its_operation() { + for kind in AggregationKind::ALL { + let (tensor, op) = (kind.tensor_type(), kind.op_variant()); + let elem = if tensor == "AddTensor" { "f32" } else { "i32" }; + let item: ItemFn = syn::parse_str(&format!( + "#[cutile::entry()] fn k(out: &{tensor}<{elem}, {{ [-1, -1] }}>) {{}}" + )) + .unwrap(); + let launcher = launcher_for(item); + let expected = format!("BindingAccess :: Atomic (AggregationOp :: {op})"); + assert!( + launcher.contains(&expected), + "{tensor}: the binding must carry {op}: {launcher}" + ); + assert!( + launcher.contains("KernelAtomic"), + "{tensor}: the exclusive-authority bound must be kept: {launcher}" + ); + } + } + + /// The capability is shared by construction; `&mut` would promise the + /// disjointness the atomic reduction deliberately does not have. + #[test] + fn a_mutable_atomic_reference_is_rejected() { + let item: ItemFn = parse_quote! { + #[cutile::entry()] + fn k(out: &mut AddTensor) {} + }; + let root = Ident::new("cutile", Span::call_site()); + let Err(err) = generate_kernel_launcher(&item, "m", "k", "k_entry", "K", "KArgs", &root) + else { + panic!("expected `&mut AddTensor<..>` to be rejected"); + }; + assert!(err.to_string().contains("`&mut AddTensor<..>`"), "{err}"); + } +} diff --git a/cutile-rs/cutile-macro/src/validate_dsl_syntax.rs b/cutile-rs/cutile-macro/src/validate_dsl_syntax.rs index 4ff5ad49c4..db4a836fd5 100644 --- a/cutile-rs/cutile-macro/src/validate_dsl_syntax.rs +++ b/cutile-rs/cutile-macro/src/validate_dsl_syntax.rs @@ -115,6 +115,21 @@ pub fn validate_entry_point_parameters(item: &ItemFn) -> Result<(), Error> { if type_name == "MappedPartitionMut" { ty.err("MappedPartitionMut parameters are passed by value; use `mut z: MappedPartitionMut<...>`, not `&mut MappedPartitionMut<...>`.")?; } + if matches!( + type_name.as_str(), + "AddTensor" + | "AndTensor" + | "OrTensor" + | "XorTensor" + | "MaxTensor" + | "MinTensor" + ) { + // A supported reference parameter. The capability is + // shared by construction, and a `&mut` spelling is + // rejected by launcher generation, which `process_items` + // runs before this pass. + continue; + } if type_name != "Tensor" { ty.err(&format!( "References to `{}` as parameters are not supported. \ diff --git a/cutile-rs/cutile/src/_core.rs b/cutile-rs/cutile/src/_core.rs index 19d3a0d00f..728ada80b6 100644 --- a/cutile-rs/cutile/src/_core.rs +++ b/cutile-rs/cutile/src/_core.rs @@ -109,6 +109,21 @@ pub mod atomic { pub struct Umax; pub struct Umin; pub struct Xchg; + + /// Addition, with the Tile IR mode selected by the destination's element + /// type: `AddF` for floats, `Add` for integers. + pub struct Sum; + + /// Maximum, with the comparison selected by the destination's element + /// type: `Max` for signed integers, `Umax` for unsigned. The Tile IR + /// scalar erases signedness, so the Rust element is the record of which + /// comparison was asked for. No float mode exists: Tile IR has no atomic + /// floating-point maximum. + pub struct Maximum; + + /// Minimum, with the comparison selected by the destination's element + /// type: `Min` for signed integers, `Umin` for unsigned. See [`Maximum`]. + pub struct Minimum; impl Mode for Add {} impl Mode for AddF {} impl Mode for And {} @@ -119,6 +134,59 @@ pub mod atomic { impl Mode for Umax {} impl Mode for Umin {} impl Mode for Xchg {} + impl Mode for Sum {} + impl Mode for Maximum {} + impl Mode for Minimum {} + + mod sealed_element { + pub trait Element {} + } + + /// Elements Tile IR can add atomically. + /// + /// Narrows the element parameter of the add-aggregation types, so an + /// illegal element fails at the kernel signature rather than at the call + /// site. Sealed: an outside impl would emit IR the target verifier rejects. + pub trait AtomicAdd: super::core::ElementType + sealed_element::Element {} + + impl sealed_element::Element for half::f16 {} + impl sealed_element::Element for half::bf16 {} + impl sealed_element::Element for f32 {} + impl sealed_element::Element for f64 {} + + impl sealed_element::Element for i32 {} + impl sealed_element::Element for u32 {} + impl sealed_element::Element for i64 {} + impl sealed_element::Element for u64 {} + + impl AtomicAdd for half::f16 {} + impl AtomicAdd for half::bf16 {} + impl AtomicAdd for f32 {} + impl AtomicAdd for f64 {} + impl AtomicAdd for i32 {} + impl AtomicAdd for u32 {} + impl AtomicAdd for i64 {} + impl AtomicAdd for u64 {} + + /// Elements Tile IR can combine bitwise atomically. Integers only; the + /// operations are bit-pattern transforms, so signedness is irrelevant and + /// both spellings of each width are admitted. + pub trait AtomicBitwise: super::core::ElementType + sealed_element::Element {} + + impl AtomicBitwise for i32 {} + impl AtomicBitwise for u32 {} + impl AtomicBitwise for i64 {} + impl AtomicBitwise for u64 {} + + /// Elements Tile IR can compare atomically. Integers only; the element's + /// signedness selects the comparison (`Max` or `Umax`, `Min` or `Umin`), + /// which is why no unsigned-named method exists anywhere in the safe API. + pub trait AtomicOrd: super::core::ElementType + sealed_element::Element {} + + impl AtomicOrd for i32 {} + impl AtomicOrd for u32 {} + impl AtomicOrd for i64 {} + impl AtomicOrd for u64 {} } /// Memory ordering, with per-op-family sub-traits restricting which variants @@ -140,6 +208,10 @@ pub mod ordering { pub trait AtomicStoreMode: StoreMode + sealed::AtomicOrdering {} /// Atomic global read-modify-write orderings. pub trait GlobalAtomicMode: AtomicMode + sealed::AtomicOrdering {} + /// `atomic_red_view_tko`: `Relaxed` and nothing else. The IR verifier + /// requires `order == 1` for a view atomic, so the other three + /// [`AtomicMode`] variants are typeable today and always fail. + pub trait RedViewMode: AtomicMode {} pub struct Weak; pub struct Relaxed; @@ -179,6 +251,7 @@ pub mod ordering { impl GlobalAtomicMode for Acquire {} impl GlobalAtomicMode for Release {} impl GlobalAtomicMode for AcqRel {} + impl RedViewMode for Relaxed {} } /// Memory scope for atomics and load/store ordering. Mirrors Tile IR @@ -200,6 +273,13 @@ pub mod scope { impl sealed::GlobalScope for System {} impl GlobalScope for Device {} impl GlobalScope for System {} + /// `atomic_red_view_tko`: `TileBlock` or `Device`; the verifier rejects + /// `System`. Note that only `Device` makes *cross-program* collisions + /// coherent — `TileBlock` is legal and almost never what a shared + /// aggregation destination wants. + pub trait RedViewScope: Mode {} + impl RedViewScope for TileBlock {} + impl RedViewScope for Device {} } mod global_sealed { @@ -357,6 +437,7 @@ pub mod core { /// Marker trait for valid `Tile` / `Tensor` element types. pub trait ElementType: Copy + Clone { const ZERO: Self; + const ONE: Self; } /// Rust byte marker for packed Tile IR sub-byte values. @@ -368,58 +449,71 @@ pub mod core { #[cuda_tile::ty(name = "bf16")] impl ElementType for bf16 { const ZERO: Self = bf16::ZERO; + const ONE: Self = bf16::ONE; } #[cuda_tile::ty(name = "f16")] impl ElementType for f16 { const ZERO: Self = f16::ZERO; + const ONE: Self = f16::ONE; } #[cuda_tile::ty(name = "f32")] impl ElementType for f32 { const ZERO: Self = 0.0; + const ONE: Self = 1.0; } #[cuda_tile::ty(name = "i8")] impl ElementType for i8 { const ZERO: Self = 0; + const ONE: Self = 1; } impl ByteElement for i8 {} #[cuda_tile::ty(name = "i8")] impl ElementType for u8 { const ZERO: Self = 0; + const ONE: Self = 1; } impl ByteElement for u8 {} #[cuda_tile::ty(name = "i32")] impl ElementType for i32 { const ZERO: Self = 0; + const ONE: Self = 1; } #[cuda_tile::ty(name = "i32")] impl ElementType for u32 { const ZERO: Self = 0; + const ONE: Self = 1; } #[cuda_tile::ty(name = "i64")] impl ElementType for i64 { const ZERO: Self = 0; + const ONE: Self = 1; } #[cuda_tile::ty(name = "i64")] impl ElementType for u64 { const ZERO: Self = 0; + const ONE: Self = 1; } #[cuda_tile::ty(name = "f64")] impl ElementType for f64 { const ZERO: Self = 0.0; + const ONE: Self = 1.0; } #[cuda_tile::ty(name = "i16")] impl ElementType for i16 { const ZERO: Self = 0; + const ONE: Self = 1; } #[cuda_tile::ty(name = "i16")] impl ElementType for u16 { const ZERO: Self = 0; + const ONE: Self = 1; } #[cuda_tile::ty(name = "i1")] impl ElementType for bool { const ZERO: Self = false; + const ONE: Self = true; } // GPU-specific types: re-exported from cuda-core. @@ -435,35 +529,43 @@ pub mod core { #[cuda_tile::ty(name = "tf32")] impl ElementType for tf32 { const ZERO: Self = tf32(0); + const ONE: Self = tf32(0x3F800000); } #[cuda_tile::ty(name = "f8e4m3fn")] impl ElementType for f8e4m3fn { const ZERO: Self = f8e4m3fn(0); + const ONE: Self = f8e4m3fn(0x38); } #[cuda_tile::ty(name = "f8e5m2")] impl ElementType for f8e5m2 { const ZERO: Self = f8e5m2(0); + const ONE: Self = f8e5m2(0x3C); } #[cuda_tile::ty(name = "f8e8m0fnu")] impl ElementType for f8e8m0fnu { const ZERO: Self = f8e8m0fnu(0); + const ONE: Self = f8e8m0fnu(0x7F); } #[cuda_tile::ty(name = "f8e5m3fnu")] impl ElementType for f8e5m3fnu { const ZERO: Self = f8e5m3fnu(0); + const ONE: Self = f8e5m3fnu(0x78); } #[cuda_tile::ty(name = "f4e2m1fn")] impl ElementType for f4e2m1fn { const ZERO: Self = f4e2m1fn(0); + const ONE: Self = f4e2m1fn(0x2); } #[cuda_tile::ty(name = "i8")] impl ElementType for f4e2m1fnx2 { const ZERO: Self = f4e2m1fnx2::from_bits(0); + const ONE: Self = f4e2m1fnx2::from_nibbles(0x2, 0x2); } impl ByteElement for f4e2m1fnx2 {} #[cuda_tile::ty(name = "i4")] impl ElementType for i4 { const ZERO: Self = i4(0); + const ONE: Self = i4(1); } /// Marker trait for scalar values that can be broadcast to tiles. @@ -896,6 +998,476 @@ pub mod core { } } + /// Shared-reference capability for overlapping atomic addition. + /// The destination is initialized by the caller; ordinary aliases are not + /// allowed while the launch is in flight. + #[cuda_tile::ty(name="!cuda_tile.tensor_view", + type_params=["{D}xE", "strides"], + type_meta=["base", "shape", "strides", "token"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct AddTensor { + _type: PhantomData, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl AddTensor { + fn shape<'b>(&self) -> Shape<'b, S> { + get_add_tensor_shape_meta(self) + } + + /// Creates an unpadded atomic partition. Tile extents must divide the + /// destination exactly on every axis. + pub fn partition<'a, const R: [i32; N]>( + &'a self, + tile: Shape, + ) -> AddPartition<'a, E, R> { + require_add_tiles_divide_exactly(self, tile); + let token: Token = get_add_tensor_token(self); + // SAFETY: built from this capability's own base, shape and + // strides, starting from its current ordering token; unpadded, + // as the Tile IR atomic reduction requires. + let view: AddPartition = + unsafe { make_add_partition_view(self, tile, padding::None, token) }; + view + } + } + + /// Tiled view returned by [`AddTensor::partition`]. + #[cuda_tile::ty(name="!cuda_tile.partition_view", + type_params=["tile"], + type_params_optional=["padding_value", "tensor_view"], + type_meta=["token", "tensor_view.shape()"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct AddPartition<'a, E: atomic::AtomicAdd, const D: [i32; N]> { + _type: PhantomData, + _tensor: PhantomData<&'a ()>, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl<'a, E: atomic::AtomicAdd, const D: [i32; N]> AddPartition<'a, E, D> { + /// Atomically adds `value` at a checked tile `index`. + pub fn atomic_add(&self, index: [i32; N], value: Tile) -> Token { + check_add_partition_access(self, index); + // SAFETY: ordered after this view's current token, bounds-checked + // above, and an aggregating mode this element type admits. + let done: Token = unsafe { self.atomic_add_unchecked(index, value) }; + done + } + + /// Atomically adds without checking `index`. + /// + /// ## Safety + /// + /// `index` must be in the partition grid, and the tile extent must + /// divide the destination extent on every axis. + pub unsafe fn atomic_add_unchecked(&self, index: [i32; N], value: Tile) -> Token { + let token: Token = get_add_partition_token(self); + // SAFETY: the view token, bounds contract, and host alias contract + // are upheld by this capability. + let done: Token = unsafe { + atomic_red_view_raw( + self, + value, + index, + atomic::Sum, + ordering::Relaxed, + scope::Device, + Some(token), + ) + }; + done + } + } + + /// Shared-reference capability for overlapping atomic bitwise AND. + #[cuda_tile::ty(name="!cuda_tile.tensor_view", + type_params=["{D}xE", "strides"], + type_meta=["base", "shape", "strides", "token"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct AndTensor { + _type: PhantomData, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl AndTensor { + fn shape<'b>(&self) -> Shape<'b, S> { + get_and_tensor_shape_meta(self) + } + + /// Creates an unpadded atomic partition. Tile extents must divide the + /// destination exactly on every axis. + pub fn partition<'a, const R: [i32; N]>( + &'a self, + tile: Shape, + ) -> AndPartition<'a, E, R> { + require_and_tiles_divide_exactly(self, tile); + let token: Token = get_and_tensor_token(self); + // SAFETY: built from this capability's own base, shape and + // strides, starting from its current ordering token; unpadded, + // as the Tile IR atomic reduction requires. + let view: AndPartition = + unsafe { make_and_partition_view(self, tile, padding::None, token) }; + view + } + } + + /// Tiled view returned by [`AndTensor::partition`]. + #[cuda_tile::ty(name="!cuda_tile.partition_view", + type_params=["tile"], + type_params_optional=["padding_value", "tensor_view"], + type_meta=["token", "tensor_view.shape()"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct AndPartition<'a, E: atomic::AtomicBitwise, const D: [i32; N]> { + _type: PhantomData, + _tensor: PhantomData<&'a ()>, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl<'a, E: atomic::AtomicBitwise, const D: [i32; N]> AndPartition<'a, E, D> { + /// Atomically ANDs `value` at a checked tile `index`. + pub fn atomic_and(&self, index: [i32; N], value: Tile) -> Token { + check_and_partition_access(self, index); + // SAFETY: ordered after this view's current token, bounds-checked + // above, and an aggregating mode this element type admits. + let done: Token = unsafe { self.atomic_and_unchecked(index, value) }; + done + } + + /// Atomically ANDs without checking `index`. + /// + /// # Safety + /// `index` must be in the partition grid and the tile must divide the + /// destination extent exactly. + pub unsafe fn atomic_and_unchecked(&self, index: [i32; N], value: Tile) -> Token { + let token: Token = get_and_partition_token(self); + // SAFETY: the view token, bounds contract, and host alias contract + // are upheld by this capability. + let done: Token = unsafe { + atomic_red_view_raw( + self, + value, + index, + atomic::And, + ordering::Relaxed, + scope::Device, + Some(token), + ) + }; + done + } + } + + /// Shared-reference capability for overlapping atomic bitwise OR. + #[cuda_tile::ty(name="!cuda_tile.tensor_view", + type_params=["{D}xE", "strides"], + type_meta=["base", "shape", "strides", "token"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct OrTensor { + _type: PhantomData, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl OrTensor { + fn shape<'b>(&self) -> Shape<'b, S> { + get_or_tensor_shape_meta(self) + } + + /// Creates an unpadded atomic partition. Tile extents must divide the + /// destination exactly on every axis. + pub fn partition<'a, const R: [i32; N]>(&'a self, tile: Shape) -> OrPartition<'a, E, R> { + require_or_tiles_divide_exactly(self, tile); + let token: Token = get_or_tensor_token(self); + // SAFETY: built from this capability's own base, shape and + // strides, starting from its current ordering token; unpadded, + // as the Tile IR atomic reduction requires. + let view: OrPartition = + unsafe { make_or_partition_view(self, tile, padding::None, token) }; + view + } + } + + /// Tiled view returned by [`OrTensor::partition`]. + #[cuda_tile::ty(name="!cuda_tile.partition_view", + type_params=["tile"], + type_params_optional=["padding_value", "tensor_view"], + type_meta=["token", "tensor_view.shape()"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct OrPartition<'a, E: atomic::AtomicBitwise, const D: [i32; N]> { + _type: PhantomData, + _tensor: PhantomData<&'a ()>, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl<'a, E: atomic::AtomicBitwise, const D: [i32; N]> OrPartition<'a, E, D> { + /// Atomically ORs `value` at a checked tile `index`. + pub fn atomic_or(&self, index: [i32; N], value: Tile) -> Token { + check_or_partition_access(self, index); + // SAFETY: ordered after this view's current token, bounds-checked + // above, and an aggregating mode this element type admits. + let done: Token = unsafe { self.atomic_or_unchecked(index, value) }; + done + } + + /// Atomically ORs without checking `index`. + /// + /// # Safety + /// `index` must be in the partition grid and the tile must divide the + /// destination extent exactly. + pub unsafe fn atomic_or_unchecked(&self, index: [i32; N], value: Tile) -> Token { + let token: Token = get_or_partition_token(self); + // SAFETY: the view token, bounds contract, and host alias contract + // are upheld by this capability. + let done: Token = unsafe { + atomic_red_view_raw( + self, + value, + index, + atomic::Or, + ordering::Relaxed, + scope::Device, + Some(token), + ) + }; + done + } + } + + /// Shared-reference capability for overlapping atomic bitwise XOR. + #[cuda_tile::ty(name="!cuda_tile.tensor_view", + type_params=["{D}xE", "strides"], + type_meta=["base", "shape", "strides", "token"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct XorTensor { + _type: PhantomData, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl XorTensor { + fn shape<'b>(&self) -> Shape<'b, S> { + get_xor_tensor_shape_meta(self) + } + + /// Creates an unpadded atomic partition. Tile extents must divide the + /// destination exactly on every axis. + pub fn partition<'a, const R: [i32; N]>( + &'a self, + tile: Shape, + ) -> XorPartition<'a, E, R> { + require_xor_tiles_divide_exactly(self, tile); + let token: Token = get_xor_tensor_token(self); + // SAFETY: built from this capability's own base, shape and + // strides, starting from its current ordering token; unpadded, + // as the Tile IR atomic reduction requires. + let view: XorPartition = + unsafe { make_xor_partition_view(self, tile, padding::None, token) }; + view + } + } + + /// Tiled view returned by [`XorTensor::partition`]. + #[cuda_tile::ty(name="!cuda_tile.partition_view", + type_params=["tile"], + type_params_optional=["padding_value", "tensor_view"], + type_meta=["token", "tensor_view.shape()"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct XorPartition<'a, E: atomic::AtomicBitwise, const D: [i32; N]> { + _type: PhantomData, + _tensor: PhantomData<&'a ()>, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl<'a, E: atomic::AtomicBitwise, const D: [i32; N]> XorPartition<'a, E, D> { + /// Atomically XORs `value` at a checked tile `index`. + pub fn atomic_xor(&self, index: [i32; N], value: Tile) -> Token { + check_xor_partition_access(self, index); + // SAFETY: ordered after this view's current token, bounds-checked + // above, and an aggregating mode this element type admits. + let done: Token = unsafe { self.atomic_xor_unchecked(index, value) }; + done + } + + /// Atomically XORs without checking `index`. + /// + /// # Safety + /// `index` must be in the partition grid and the tile must divide the + /// destination extent exactly. + pub unsafe fn atomic_xor_unchecked(&self, index: [i32; N], value: Tile) -> Token { + let token: Token = get_xor_partition_token(self); + // SAFETY: the view token, bounds contract, and host alias contract + // are upheld by this capability. + let done: Token = unsafe { + atomic_red_view_raw( + self, + value, + index, + atomic::Xor, + ordering::Relaxed, + scope::Device, + Some(token), + ) + }; + done + } + } + + /// Shared-reference capability for overlapping atomic maximum. Signedness + /// follows the Rust element type; floating-point maximum is unsupported. + #[cuda_tile::ty(name="!cuda_tile.tensor_view", + type_params=["{D}xE", "strides"], + type_meta=["base", "shape", "strides", "token"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct MaxTensor { + _type: PhantomData, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl MaxTensor { + fn shape<'b>(&self) -> Shape<'b, S> { + get_max_tensor_shape_meta(self) + } + + /// Creates an unpadded atomic partition. Tile extents must divide the + /// destination exactly on every axis. + pub fn partition<'a, const R: [i32; N]>( + &'a self, + tile: Shape, + ) -> MaxPartition<'a, E, R> { + require_max_tiles_divide_exactly(self, tile); + let token: Token = get_max_tensor_token(self); + // SAFETY: built from this capability's own base, shape and + // strides, starting from its current ordering token; unpadded, + // as the Tile IR atomic reduction requires. + let view: MaxPartition = + unsafe { make_max_partition_view(self, tile, padding::None, token) }; + view + } + } + + /// Tiled view returned by [`MaxTensor::partition`]. + #[cuda_tile::ty(name="!cuda_tile.partition_view", + type_params=["tile"], + type_params_optional=["padding_value", "tensor_view"], + type_meta=["token", "tensor_view.shape()"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct MaxPartition<'a, E: atomic::AtomicOrd, const D: [i32; N]> { + _type: PhantomData, + _tensor: PhantomData<&'a ()>, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl<'a, E: atomic::AtomicOrd, const D: [i32; N]> MaxPartition<'a, E, D> { + /// Atomically takes the maximum with `value` at a checked tile `index`. + pub fn atomic_max(&self, index: [i32; N], value: Tile) -> Token { + check_max_partition_access(self, index); + // SAFETY: ordered after this view's current token, bounds-checked + // above, and an aggregating mode this element type admits. + let done: Token = unsafe { self.atomic_max_unchecked(index, value) }; + done + } + + /// Atomically takes the maximum without checking `index`. + /// + /// # Safety + /// `index` must be in the partition grid and the tile must divide the + /// destination extent exactly. + pub unsafe fn atomic_max_unchecked(&self, index: [i32; N], value: Tile) -> Token { + let token: Token = get_max_partition_token(self); + // SAFETY: the view token, bounds contract, and host alias contract + // are upheld by this capability. + let done: Token = unsafe { + atomic_red_view_raw( + self, + value, + index, + atomic::Maximum, + ordering::Relaxed, + scope::Device, + Some(token), + ) + }; + done + } + } + + /// Shared-reference capability for overlapping atomic minimum. Signedness + /// follows the Rust element type; floating-point minimum is unsupported. + #[cuda_tile::ty(name="!cuda_tile.tensor_view", + type_params=["{D}xE", "strides"], + type_meta=["base", "shape", "strides", "token"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct MinTensor { + _type: PhantomData, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl MinTensor { + fn shape<'b>(&self) -> Shape<'b, S> { + get_min_tensor_shape_meta(self) + } + + /// Creates an unpadded atomic partition. Tile extents must divide the + /// destination exactly on every axis. + pub fn partition<'a, const R: [i32; N]>( + &'a self, + tile: Shape, + ) -> MinPartition<'a, E, R> { + require_min_tiles_divide_exactly(self, tile); + let token: Token = get_min_tensor_token(self); + // SAFETY: built from this capability's own base, shape and + // strides, starting from its current ordering token; unpadded, + // as the Tile IR atomic reduction requires. + let view: MinPartition = + unsafe { make_min_partition_view(self, tile, padding::None, token) }; + view + } + } + + /// Tiled view returned by [`MinTensor::partition`]. + #[cuda_tile::ty(name="!cuda_tile.partition_view", + type_params=["tile"], + type_params_optional=["padding_value", "tensor_view"], + type_meta=["token", "tensor_view.shape()"])] + #[cuda_tile::variadic_struct(N = 6)] + pub struct MinPartition<'a, E: atomic::AtomicOrd, const D: [i32; N]> { + _type: PhantomData, + _tensor: PhantomData<&'a ()>, + } + + #[cuda_tile::variadic_impl(N = 6)] + impl<'a, E: atomic::AtomicOrd, const D: [i32; N]> MinPartition<'a, E, D> { + /// Atomically takes the minimum with `value` at a checked tile `index`. + pub fn atomic_min(&self, index: [i32; N], value: Tile) -> Token { + check_min_partition_access(self, index); + // SAFETY: ordered after this view's current token, bounds-checked + // above, and an aggregating mode this element type admits. + let done: Token = unsafe { self.atomic_min_unchecked(index, value) }; + done + } + + /// Atomically takes the minimum without checking `index`. + /// + /// # Safety + /// `index` must be in the partition grid and the tile must divide the + /// destination extent exactly. + pub unsafe fn atomic_min_unchecked(&self, index: [i32; N], value: Tile) -> Token { + let token: Token = get_min_partition_token(self); + // SAFETY: the view token, bounds contract, and host alias contract + // are upheld by this capability. + let done: Token = unsafe { + atomic_red_view_raw( + self, + value, + index, + atomic::Minimum, + ordering::Relaxed, + scope::Device, + Some(token), + ) + }; + done + } + } + /// Private logical index into a tensor partition grid. /// /// `PartitionIndex` values are produced by cutile schedule helpers. Safe @@ -1811,32 +2383,230 @@ pub mod core { unreachable!() } - #[cuda_tile::compiler_op(name = "dim_new")] - pub fn dim_new(size: i32) -> Dim { - unreachable!() - } - - #[cuda_tile::compiler_op(name = "dim_from_i32")] - pub fn dim_from_i32(size: i32) -> Dim { + /// Require that `tile` divides the `AddTensor` exactly on every axis. + /// + /// Emits one launch-time obligation per axis, + /// `dim(tensor, axis) % tile[axis] == 0`; obligations the JIT can + /// already discharge (a static shape that divides) cost nothing, and a + /// static shape that does not divide fails compilation. A safety + /// requirement, not a convenience: the partition index space is a + /// ceiling division and the atomic reduction's view is unpadded, with no + /// documented masking for partial-edge lanes. + #[cuda_tile::compiler_op(name = "require_tiles_divide_exactly")] + #[cuda_tile::variadic_op(N = 6)] + fn require_add_tiles_divide_exactly< + E: atomic::AtomicAdd, + const S: [i32; N], + const R: [i32; N], + >( + tensor: &AddTensor, + tile: Shape, + ) { unreachable!() } - #[cuda_tile::compiler_op(name = "dim_value")] - pub fn dim_value(dim: Dim) -> i32 { + /// Require that `tile` divides the `AndTensor` exactly on every axis. + /// + /// Emits one launch-time obligation per axis, + /// `dim(tensor, axis) % tile[axis] == 0`; obligations the JIT can + /// already discharge (a static shape that divides) cost nothing, and a + /// static shape that does not divide fails compilation. A safety + /// requirement, not a convenience: the partition index space is a + /// ceiling division and the atomic reduction's view is unpadded, with no + /// documented masking for partial-edge lanes. + #[cuda_tile::compiler_op(name = "require_tiles_divide_exactly")] + #[cuda_tile::variadic_op(N = 6)] + fn require_and_tiles_divide_exactly< + E: atomic::AtomicBitwise, + const S: [i32; N], + const R: [i32; N], + >( + tensor: &AndTensor, + tile: Shape, + ) { unreachable!() } - #[cuda_tile::compiler_op(name = "coord_as_array")] - pub fn coord2_as_array(index: Coord2) -> [i32; 2] { + /// Require that `tile` divides the `OrTensor` exactly on every axis. + /// + /// Emits one launch-time obligation per axis, + /// `dim(tensor, axis) % tile[axis] == 0`; obligations the JIT can + /// already discharge (a static shape that divides) cost nothing, and a + /// static shape that does not divide fails compilation. A safety + /// requirement, not a convenience: the partition index space is a + /// ceiling division and the atomic reduction's view is unpadded, with no + /// documented masking for partial-edge lanes. + #[cuda_tile::compiler_op(name = "require_tiles_divide_exactly")] + #[cuda_tile::variadic_op(N = 6)] + fn require_or_tiles_divide_exactly< + E: atomic::AtomicBitwise, + const S: [i32; N], + const R: [i32; N], + >( + tensor: &OrTensor, + tile: Shape, + ) { unreachable!() } - #[cuda_tile::compiler_op(name = "coord_as_array")] - pub fn coord3_as_array(index: Coord3) -> [i32; 3] { + /// Require that `tile` divides the `XorTensor` exactly on every axis. + /// + /// Emits one launch-time obligation per axis, + /// `dim(tensor, axis) % tile[axis] == 0`; obligations the JIT can + /// already discharge (a static shape that divides) cost nothing, and a + /// static shape that does not divide fails compilation. A safety + /// requirement, not a convenience: the partition index space is a + /// ceiling division and the atomic reduction's view is unpadded, with no + /// documented masking for partial-edge lanes. + #[cuda_tile::compiler_op(name = "require_tiles_divide_exactly")] + #[cuda_tile::variadic_op(N = 6)] + fn require_xor_tiles_divide_exactly< + E: atomic::AtomicBitwise, + const S: [i32; N], + const R: [i32; N], + >( + tensor: &XorTensor, + tile: Shape, + ) { unreachable!() } - #[cuda_tile::compiler_op(name = "partition_with_bounds")] + /// Require that `tile` divides the `MaxTensor` exactly on every axis. + /// + /// Emits one launch-time obligation per axis, + /// `dim(tensor, axis) % tile[axis] == 0`; obligations the JIT can + /// already discharge (a static shape that divides) cost nothing, and a + /// static shape that does not divide fails compilation. A safety + /// requirement, not a convenience: the partition index space is a + /// ceiling division and the atomic reduction's view is unpadded, with no + /// documented masking for partial-edge lanes. + #[cuda_tile::compiler_op(name = "require_tiles_divide_exactly")] + #[cuda_tile::variadic_op(N = 6)] + fn require_max_tiles_divide_exactly< + E: atomic::AtomicOrd, + const S: [i32; N], + const R: [i32; N], + >( + tensor: &MaxTensor, + tile: Shape, + ) { + unreachable!() + } + + /// Require that `tile` divides the `MinTensor` exactly on every axis. + /// + /// Emits one launch-time obligation per axis, + /// `dim(tensor, axis) % tile[axis] == 0`; obligations the JIT can + /// already discharge (a static shape that divides) cost nothing, and a + /// static shape that does not divide fails compilation. A safety + /// requirement, not a convenience: the partition index space is a + /// ceiling division and the atomic reduction's view is unpadded, with no + /// documented masking for partial-edge lanes. + #[cuda_tile::compiler_op(name = "require_tiles_divide_exactly")] + #[cuda_tile::variadic_op(N = 6)] + fn require_min_tiles_divide_exactly< + E: atomic::AtomicOrd, + const S: [i32; N], + const R: [i32; N], + >( + tensor: &MinTensor, + tile: Shape, + ) { + unreachable!() + } + + /// [`check_partition_access`] for an `AddPartition`. Same compiler + /// handler; a separate declaration only because the view type differs. + #[cuda_tile::compiler_op(name = "check")] + #[cuda_tile::variadic_op(N = 6)] + pub fn check_add_partition_access( + part: &AddPartition<'_, E, S>, + index: [i32; N], + ) { + unreachable!() + } + + /// [`check_partition_access`] for an `AndPartition`. Same compiler + /// handler; a separate declaration only because the view type differs. + #[cuda_tile::compiler_op(name = "check")] + #[cuda_tile::variadic_op(N = 6)] + pub fn check_and_partition_access( + part: &AndPartition<'_, E, S>, + index: [i32; N], + ) { + unreachable!() + } + + /// [`check_partition_access`] for an `OrPartition`. Same compiler + /// handler; a separate declaration only because the view type differs. + #[cuda_tile::compiler_op(name = "check")] + #[cuda_tile::variadic_op(N = 6)] + pub fn check_or_partition_access( + part: &OrPartition<'_, E, S>, + index: [i32; N], + ) { + unreachable!() + } + + /// [`check_partition_access`] for a `XorPartition`. Same compiler + /// handler; a separate declaration only because the view type differs. + #[cuda_tile::compiler_op(name = "check")] + #[cuda_tile::variadic_op(N = 6)] + pub fn check_xor_partition_access( + part: &XorPartition<'_, E, S>, + index: [i32; N], + ) { + unreachable!() + } + + /// [`check_partition_access`] for a `MaxPartition`. Same compiler + /// handler; a separate declaration only because the view type differs. + #[cuda_tile::compiler_op(name = "check")] + #[cuda_tile::variadic_op(N = 6)] + pub fn check_max_partition_access( + part: &MaxPartition<'_, E, S>, + index: [i32; N], + ) { + unreachable!() + } + + /// [`check_partition_access`] for a `MinPartition`. Same compiler + /// handler; a separate declaration only because the view type differs. + #[cuda_tile::compiler_op(name = "check")] + #[cuda_tile::variadic_op(N = 6)] + pub fn check_min_partition_access( + part: &MinPartition<'_, E, S>, + index: [i32; N], + ) { + unreachable!() + } + + #[cuda_tile::compiler_op(name = "dim_new")] + pub fn dim_new(size: i32) -> Dim { + unreachable!() + } + + #[cuda_tile::compiler_op(name = "dim_from_i32")] + pub fn dim_from_i32(size: i32) -> Dim { + unreachable!() + } + + #[cuda_tile::compiler_op(name = "dim_value")] + pub fn dim_value(dim: Dim) -> i32 { + unreachable!() + } + + #[cuda_tile::compiler_op(name = "coord_as_array")] + pub fn coord2_as_array(index: Coord2) -> [i32; 2] { + unreachable!() + } + + #[cuda_tile::compiler_op(name = "coord_as_array")] + pub fn coord3_as_array(index: Coord3) -> [i32; 3] { + unreachable!() + } + + #[cuda_tile::compiler_op(name = "partition_with_bounds")] pub fn partition_with_bounds<'a, E: ElementType, const S: [i32; 2]>( part: Partition<'a, E, S>, bounds: (Dim, Dim), @@ -2219,6 +2989,66 @@ pub mod core { unreachable!() } + /// Extract an `AddTensor`'s shape from its type metadata. + /// Internal to its `shape`. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "shape")] + fn get_add_tensor_shape_meta<'s, E: atomic::AtomicAdd, const S: [i32; N]>( + tensor: &AddTensor, + ) -> Shape<'s, S> { + unreachable!() + } + + /// Extract an `AndTensor`'s shape from its type metadata. + /// Internal to its `shape`. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "shape")] + fn get_and_tensor_shape_meta<'s, E: atomic::AtomicBitwise, const S: [i32; N]>( + tensor: &AndTensor, + ) -> Shape<'s, S> { + unreachable!() + } + + /// Extract an `OrTensor`'s shape from its type metadata. + /// Internal to its `shape`. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "shape")] + fn get_or_tensor_shape_meta<'s, E: atomic::AtomicBitwise, const S: [i32; N]>( + tensor: &OrTensor, + ) -> Shape<'s, S> { + unreachable!() + } + + /// Extract a `XorTensor`'s shape from its type metadata. + /// Internal to its `shape`. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "shape")] + fn get_xor_tensor_shape_meta<'s, E: atomic::AtomicBitwise, const S: [i32; N]>( + tensor: &XorTensor, + ) -> Shape<'s, S> { + unreachable!() + } + + /// Extract a `MaxTensor`'s shape from its type metadata. + /// Internal to its `shape`. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "shape")] + fn get_max_tensor_shape_meta<'s, E: atomic::AtomicOrd, const S: [i32; N]>( + tensor: &MaxTensor, + ) -> Shape<'s, S> { + unreachable!() + } + + /// Extract a `MinTensor`'s shape from its type metadata. + /// Internal to its `shape`. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "shape")] + fn get_min_tensor_shape_meta<'s, E: atomic::AtomicOrd, const S: [i32; N]>( + tensor: &MinTensor, + ) -> Shape<'s, S> { + unreachable!() + } + /// Extract a `Tensor`'s ordering token. #[cuda_tile::variadic_op(N = 6)] #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] @@ -2226,6 +3056,72 @@ pub mod core { unreachable!() } + /// Extract an `AddTensor`'s ordering token. Internal to its + /// `partition`, which seeds the view it returns with the capability's + /// current token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_add_tensor_token( + tensor: &AddTensor, + ) -> Token { + unreachable!() + } + + /// Extract an `AndTensor`'s ordering token. Internal to its + /// `partition`, which seeds the view it returns with the capability's + /// current token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_and_tensor_token( + tensor: &AndTensor, + ) -> Token { + unreachable!() + } + + /// Extract an `OrTensor`'s ordering token. Internal to its + /// `partition`, which seeds the view it returns with the capability's + /// current token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_or_tensor_token( + tensor: &OrTensor, + ) -> Token { + unreachable!() + } + + /// Extract a `XorTensor`'s ordering token. Internal to its + /// `partition`, which seeds the view it returns with the capability's + /// current token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_xor_tensor_token( + tensor: &XorTensor, + ) -> Token { + unreachable!() + } + + /// Extract a `MaxTensor`'s ordering token. Internal to its + /// `partition`, which seeds the view it returns with the capability's + /// current token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_max_tensor_token( + tensor: &MaxTensor, + ) -> Token { + unreachable!() + } + + /// Extract a `MinTensor`'s ordering token. Internal to its + /// `partition`, which seeds the view it returns with the capability's + /// current token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_min_tensor_token( + tensor: &MinTensor, + ) -> Token { + unreachable!() + } + /// Extract a `Tensor`'s base pointer from its type metadata. #[cuda_tile::variadic_op(N = 6)] #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "base")] @@ -2303,6 +3199,60 @@ pub mod core { unreachable!() } + /// Extract an `AddPartition`'s ordering token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_add_partition_token( + view: &AddPartition<'_, E, D>, + ) -> Token { + unreachable!() + } + + /// Extract an `AndPartition`'s ordering token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_and_partition_token( + view: &AndPartition<'_, E, D>, + ) -> Token { + unreachable!() + } + + /// Extract an `OrPartition`'s ordering token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_or_partition_token( + view: &OrPartition<'_, E, D>, + ) -> Token { + unreachable!() + } + + /// Extract a `XorPartition`'s ordering token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_xor_partition_token( + view: &XorPartition<'_, E, D>, + ) -> Token { + unreachable!() + } + + /// Extract a `MaxPartition`'s ordering token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_max_partition_token( + view: &MaxPartition<'_, E, D>, + ) -> Token { + unreachable!() + } + + /// Extract a `MinPartition`'s ordering token. + #[cuda_tile::variadic_op(N = 6)] + #[cuda_tile::compiler_op(name = "return_type_meta_field", type_meta_field = "token")] + fn get_min_partition_token( + view: &MinPartition<'_, E, D>, + ) -> Token { + unreachable!() + } + /// Number of tiles along `axis` (`cdiv(tensor_dim, tile_dim)`). /// `axis` must be a compile-time constant in `0..N`. Lowers to /// `cuda_tile.get_index_space_shape` with axis-th result extracted. @@ -3366,8 +4316,8 @@ pub mod core { V, I, M: atomic::Mode, - O: ordering::AtomicMode, - Sc: scope::Mode, + O: ordering::RedViewMode, + Sc: scope::RedViewScope, >( view: &V, value: Tile, @@ -3546,6 +4496,186 @@ pub mod core { unreachable!() } + /// Build an `AddTensor` capability from a base pointer, shape, + /// strides, and a token. Emitted by generated entry wrappers for + /// `&AddTensor<_, _>` parameters; the same Tile IR tensor view as + /// `make_tensor_view`, with the atomic-only Rust capability on top. + /// + /// # Safety + /// As for `make_tensor_view`: the layout must be valid and the storage + /// live for every access through the capability. In addition, the caller + /// must hold exclusive mutable authority over that storage — see + /// `AddTensor` for the full precondition list. + #[cuda_tile::op(name="cuda_tile.make_tensor_view", since="V13_2", + params=["base", "shape.dims", "strides.dims"], + has_variadic_params=true, + output_type_params=["strides"], + output_type_meta=["base", "shape", "strides", "token"] + )] + #[cuda_tile::variadic_op(N = 6)] + pub unsafe fn make_add_tensor_view< + E: atomic::AtomicAdd, + const D: [i32; N], + const C: [i32; N], + >( + base: PointerTile<*mut E, { [] }>, + shape: Shape, + strides: Array, + token: Token, + ) -> AddTensor { + unreachable!() + } + + /// Build an `AndTensor` capability from a base pointer, shape, + /// strides, and a token. Emitted by generated entry wrappers for + /// `&AndTensor<_, _>` parameters; the same Tile IR tensor view as + /// `make_tensor_view`, with the atomic-only Rust capability on top. + /// + /// # Safety + /// As for `make_tensor_view`: the layout must be valid and the storage + /// live for every access through the capability. In addition, the caller + /// must hold exclusive mutable authority over that storage — see + /// `AndTensor` for the full precondition list. + #[cuda_tile::op(name="cuda_tile.make_tensor_view", since="V13_2", + params=["base", "shape.dims", "strides.dims"], + has_variadic_params=true, + output_type_params=["strides"], + output_type_meta=["base", "shape", "strides", "token"] + )] + #[cuda_tile::variadic_op(N = 6)] + pub unsafe fn make_and_tensor_view< + E: atomic::AtomicBitwise, + const D: [i32; N], + const C: [i32; N], + >( + base: PointerTile<*mut E, { [] }>, + shape: Shape, + strides: Array, + token: Token, + ) -> AndTensor { + unreachable!() + } + + /// Build an `OrTensor` capability from a base pointer, shape, + /// strides, and a token. Emitted by generated entry wrappers for + /// `&OrTensor<_, _>` parameters; the same Tile IR tensor view as + /// `make_tensor_view`, with the atomic-only Rust capability on top. + /// + /// # Safety + /// As for `make_tensor_view`: the layout must be valid and the storage + /// live for every access through the capability. In addition, the caller + /// must hold exclusive mutable authority over that storage — see + /// `OrTensor` for the full precondition list. + #[cuda_tile::op(name="cuda_tile.make_tensor_view", since="V13_2", + params=["base", "shape.dims", "strides.dims"], + has_variadic_params=true, + output_type_params=["strides"], + output_type_meta=["base", "shape", "strides", "token"] + )] + #[cuda_tile::variadic_op(N = 6)] + pub unsafe fn make_or_tensor_view< + E: atomic::AtomicBitwise, + const D: [i32; N], + const C: [i32; N], + >( + base: PointerTile<*mut E, { [] }>, + shape: Shape, + strides: Array, + token: Token, + ) -> OrTensor { + unreachable!() + } + + /// Build a `XorTensor` capability from a base pointer, shape, + /// strides, and a token. Emitted by generated entry wrappers for + /// `&XorTensor<_, _>` parameters; the same Tile IR tensor view as + /// `make_tensor_view`, with the atomic-only Rust capability on top. + /// + /// # Safety + /// As for `make_tensor_view`: the layout must be valid and the storage + /// live for every access through the capability. In addition, the caller + /// must hold exclusive mutable authority over that storage — see + /// `XorTensor` for the full precondition list. + #[cuda_tile::op(name="cuda_tile.make_tensor_view", since="V13_2", + params=["base", "shape.dims", "strides.dims"], + has_variadic_params=true, + output_type_params=["strides"], + output_type_meta=["base", "shape", "strides", "token"] + )] + #[cuda_tile::variadic_op(N = 6)] + pub unsafe fn make_xor_tensor_view< + E: atomic::AtomicBitwise, + const D: [i32; N], + const C: [i32; N], + >( + base: PointerTile<*mut E, { [] }>, + shape: Shape, + strides: Array, + token: Token, + ) -> XorTensor { + unreachable!() + } + + /// Build a `MaxTensor` capability from a base pointer, shape, + /// strides, and a token. Emitted by generated entry wrappers for + /// `&MaxTensor<_, _>` parameters; the same Tile IR tensor view as + /// `make_tensor_view`, with the atomic-only Rust capability on top. + /// + /// # Safety + /// As for `make_tensor_view`: the layout must be valid and the storage + /// live for every access through the capability. In addition, the caller + /// must hold exclusive mutable authority over that storage — see + /// `MaxTensor` for the full precondition list. + #[cuda_tile::op(name="cuda_tile.make_tensor_view", since="V13_2", + params=["base", "shape.dims", "strides.dims"], + has_variadic_params=true, + output_type_params=["strides"], + output_type_meta=["base", "shape", "strides", "token"] + )] + #[cuda_tile::variadic_op(N = 6)] + pub unsafe fn make_max_tensor_view< + E: atomic::AtomicOrd, + const D: [i32; N], + const C: [i32; N], + >( + base: PointerTile<*mut E, { [] }>, + shape: Shape, + strides: Array, + token: Token, + ) -> MaxTensor { + unreachable!() + } + + /// Build a `MinTensor` capability from a base pointer, shape, + /// strides, and a token. Emitted by generated entry wrappers for + /// `&MinTensor<_, _>` parameters; the same Tile IR tensor view as + /// `make_tensor_view`, with the atomic-only Rust capability on top. + /// + /// # Safety + /// As for `make_tensor_view`: the layout must be valid and the storage + /// live for every access through the capability. In addition, the caller + /// must hold exclusive mutable authority over that storage — see + /// `MinTensor` for the full precondition list. + #[cuda_tile::op(name="cuda_tile.make_tensor_view", since="V13_2", + params=["base", "shape.dims", "strides.dims"], + has_variadic_params=true, + output_type_params=["strides"], + output_type_meta=["base", "shape", "strides", "token"] + )] + #[cuda_tile::variadic_op(N = 6)] + pub unsafe fn make_min_tensor_view< + E: atomic::AtomicOrd, + const D: [i32; N], + const C: [i32; N], + >( + base: PointerTile<*mut E, { [] }>, + shape: Shape, + strides: Array, + token: Token, + ) -> MinTensor { + unreachable!() + } + /// Build a read-only partition view. Prefer `tensor.partition(tile)`. /// /// Pass `padding::None` to omit the Tile IR `padding_value` type @@ -3612,6 +4742,198 @@ pub mod core { unreachable!() } + /// Build the tiled view an `AddTensor` operation addresses. + /// + /// Internal to `AddTensor::partition`. Declared non-`pub` to mark it + /// internal, but `#[cuda_tile::variadic_op]` emits a `pub` dispatch + /// wrapper regardless, so this is intent and not enforcement. + /// + /// # Safety + /// As for `make_partition_view_mut`. The returned view must be used + /// for this one atomic operation only: it carries no disjointness + /// promise, so an ordinary store through it would race the other + /// programs' atomics. + #[cuda_tile::op(name="cuda_tile.make_partition_view", since="V13_2", + params=["tensor_view"], + output_type_params=["tensor_view", "padding_value"], + output_type_meta=["token", "tensor_view.shape()"] + )] + #[cuda_tile::variadic_op(N = 6)] + unsafe fn make_add_partition_view< + 'a, + E: atomic::AtomicAdd, + const TENSOR_SHAPE: [i32; N], + const TILE_SHAPE: [i32; N], + P: padding::Mode, + >( + tensor_view: &AddTensor, + shape: Shape, + padding_value: P, + token: Token, + ) -> AddPartition<'a, E, TILE_SHAPE> { + unreachable!() + } + + /// Build the tiled view an `AndTensor` operation addresses. + /// + /// Internal to `AndTensor::partition`. Declared non-`pub` to mark it + /// internal, but `#[cuda_tile::variadic_op]` emits a `pub` dispatch + /// wrapper regardless, so this is intent and not enforcement. + /// + /// # Safety + /// As for `make_partition_view_mut`. The returned view must be used + /// for this one atomic operation only: it carries no disjointness + /// promise, so an ordinary store through it would race the other + /// programs' atomics. + #[cuda_tile::op(name="cuda_tile.make_partition_view", since="V13_2", + params=["tensor_view"], + output_type_params=["tensor_view", "padding_value"], + output_type_meta=["token", "tensor_view.shape()"] + )] + #[cuda_tile::variadic_op(N = 6)] + unsafe fn make_and_partition_view< + 'a, + E: atomic::AtomicBitwise, + const TENSOR_SHAPE: [i32; N], + const TILE_SHAPE: [i32; N], + P: padding::Mode, + >( + tensor_view: &AndTensor, + shape: Shape, + padding_value: P, + token: Token, + ) -> AndPartition<'a, E, TILE_SHAPE> { + unreachable!() + } + + /// Build the tiled view an `OrTensor` operation addresses. + /// + /// Internal to `OrTensor::partition`. Declared non-`pub` to mark it + /// internal, but `#[cuda_tile::variadic_op]` emits a `pub` dispatch + /// wrapper regardless, so this is intent and not enforcement. + /// + /// # Safety + /// As for `make_partition_view_mut`. The returned view must be used + /// for this one atomic operation only: it carries no disjointness + /// promise, so an ordinary store through it would race the other + /// programs' atomics. + #[cuda_tile::op(name="cuda_tile.make_partition_view", since="V13_2", + params=["tensor_view"], + output_type_params=["tensor_view", "padding_value"], + output_type_meta=["token", "tensor_view.shape()"] + )] + #[cuda_tile::variadic_op(N = 6)] + unsafe fn make_or_partition_view< + 'a, + E: atomic::AtomicBitwise, + const TENSOR_SHAPE: [i32; N], + const TILE_SHAPE: [i32; N], + P: padding::Mode, + >( + tensor_view: &OrTensor, + shape: Shape, + padding_value: P, + token: Token, + ) -> OrPartition<'a, E, TILE_SHAPE> { + unreachable!() + } + + /// Build the tiled view a `XorTensor` operation addresses. + /// + /// Internal to `XorTensor::partition`. Declared non-`pub` to mark it + /// internal, but `#[cuda_tile::variadic_op]` emits a `pub` dispatch + /// wrapper regardless, so this is intent and not enforcement. + /// + /// # Safety + /// As for `make_partition_view_mut`. The returned view must be used + /// for this one atomic operation only: it carries no disjointness + /// promise, so an ordinary store through it would race the other + /// programs' atomics. + #[cuda_tile::op(name="cuda_tile.make_partition_view", since="V13_2", + params=["tensor_view"], + output_type_params=["tensor_view", "padding_value"], + output_type_meta=["token", "tensor_view.shape()"] + )] + #[cuda_tile::variadic_op(N = 6)] + unsafe fn make_xor_partition_view< + 'a, + E: atomic::AtomicBitwise, + const TENSOR_SHAPE: [i32; N], + const TILE_SHAPE: [i32; N], + P: padding::Mode, + >( + tensor_view: &XorTensor, + shape: Shape, + padding_value: P, + token: Token, + ) -> XorPartition<'a, E, TILE_SHAPE> { + unreachable!() + } + + /// Build the tiled view a `MaxTensor` operation addresses. + /// + /// Internal to `MaxTensor::partition`. Declared non-`pub` to mark it + /// internal, but `#[cuda_tile::variadic_op]` emits a `pub` dispatch + /// wrapper regardless, so this is intent and not enforcement. + /// + /// # Safety + /// As for `make_partition_view_mut`. The returned view must be used + /// for this one atomic operation only: it carries no disjointness + /// promise, so an ordinary store through it would race the other + /// programs' atomics. + #[cuda_tile::op(name="cuda_tile.make_partition_view", since="V13_2", + params=["tensor_view"], + output_type_params=["tensor_view", "padding_value"], + output_type_meta=["token", "tensor_view.shape()"] + )] + #[cuda_tile::variadic_op(N = 6)] + unsafe fn make_max_partition_view< + 'a, + E: atomic::AtomicOrd, + const TENSOR_SHAPE: [i32; N], + const TILE_SHAPE: [i32; N], + P: padding::Mode, + >( + tensor_view: &MaxTensor, + shape: Shape, + padding_value: P, + token: Token, + ) -> MaxPartition<'a, E, TILE_SHAPE> { + unreachable!() + } + + /// Build the tiled view a `MinTensor` operation addresses. + /// + /// Internal to `MinTensor::partition`. Declared non-`pub` to mark it + /// internal, but `#[cuda_tile::variadic_op]` emits a `pub` dispatch + /// wrapper regardless, so this is intent and not enforcement. + /// + /// # Safety + /// As for `make_partition_view_mut`. The returned view must be used + /// for this one atomic operation only: it carries no disjointness + /// promise, so an ordinary store through it would race the other + /// programs' atomics. + #[cuda_tile::op(name="cuda_tile.make_partition_view", since="V13_2", + params=["tensor_view"], + output_type_params=["tensor_view", "padding_value"], + output_type_meta=["token", "tensor_view.shape()"] + )] + #[cuda_tile::variadic_op(N = 6)] + unsafe fn make_min_partition_view< + 'a, + E: atomic::AtomicOrd, + const TENSOR_SHAPE: [i32; N], + const TILE_SHAPE: [i32; N], + P: padding::Mode, + >( + tensor_view: &MinTensor, + shape: Shape, + padding_value: P, + token: Token, + ) -> MinPartition<'a, E, TILE_SHAPE> { + unreachable!() + } + /// Build a mapped mutable partition view. /// /// This is used by generated entry wrappers for diff --git a/cutile-rs/cutile/src/_tileir.rs b/cutile-rs/cutile/src/_tileir.rs index ae5cd957d6..a93dfe2d63 100644 --- a/cutile-rs/cutile/src/_tileir.rs +++ b/cutile-rs/cutile/src/_tileir.rs @@ -30,8 +30,8 @@ pub mod tileir { E: ElementType, const D: [i32; N], M: atomic::Mode, - O: ordering::AtomicMode, - Sc: scope::Mode, + O: ordering::RedViewMode, + Sc: scope::RedViewScope, >( view: &mut PartitionMut, value: Tile, diff --git a/cutile-rs/cutile/src/prelude.rs b/cutile-rs/cutile/src/prelude.rs index bb5c1f3700..3dfbb17aa9 100644 --- a/cutile-rs/cutile/src/prelude.rs +++ b/cutile-rs/cutile/src/prelude.rs @@ -20,9 +20,9 @@ pub use crate::error::Error; // Tensor types and traits pub use crate::tensor::{ - IntoPartition, KernelInput, KernelInputStored, KernelOutput, KernelOutputStored, - MappedLaunchPartition, Partition, PartitionMut, Reshape, SpecializationBits, Tensor, - TensorView, ToHostVec, TryPartition, Unpartition, OWNED, + Atomic, IntoPartition, KernelAtomic, KernelInput, KernelInputStored, KernelOutput, + KernelOutputStored, MappedLaunchPartition, Partition, PartitionMut, Reshape, + SpecializationBits, Tensor, TensorView, ToHostVec, TryPartition, Unpartition, OWNED, }; // Tile kernel traits diff --git a/cutile-rs/cutile/src/tensor.rs b/cutile-rs/cutile/src/tensor.rs index 68cf53b8ce..d2a4bc3f58 100644 --- a/cutile-rs/cutile/src/tensor.rs +++ b/cutile-rs/cutile/src/tensor.rs @@ -1292,6 +1292,11 @@ impl Tensor { self.strides == contiguous_strides(&self.shape) } + /// Identity of the backing allocation, used for launch alias checks. + pub fn storage_id(&self) -> StorageId { + StorageId(Arc::as_ptr(&self.storage) as *const () as usize) + } + /// Create an `Arc>` that shares this tensor's device memory. /// /// # Safety @@ -1884,6 +1889,8 @@ pub trait KernelOutputStored: Send { fn strides_hint(&self) -> Vec; fn spec(&self) -> &SpecializationBits; fn shape_as_i32(&self) -> Vec; + /// Identity of the backing allocation. + fn storage_id(&self) -> StorageId; } /// How a `&mut Tensor` kernel param is stored during execution and recovered. @@ -1956,6 +1963,9 @@ impl KernelOutputStored for Partition> { fn shape_as_i32(&self) -> Vec { self.object.shape.clone() } + fn storage_id(&self) -> StorageId { + self.object.storage_id() + } } impl KernelOutputStored for Partition<&mut Tensor> { @@ -2014,6 +2024,9 @@ impl KernelOutputStored for Partition<&mut Tensor> { fn shape_as_i32(&self) -> Vec { self.object.shape.clone() } + fn storage_id(&self) -> StorageId { + self.object.storage_id() + } } impl KernelOutputStored for MappedLaunchPartition>> { @@ -2054,6 +2067,9 @@ impl KernelOutputStored for MappedLaunchPartition Vec { self.partition.shape_as_i32() } + fn storage_id(&self) -> StorageId { + KernelOutputStored::storage_id(&self.partition) + } } impl KernelOutputStored for MappedLaunchPartition>> { @@ -2094,6 +2110,9 @@ impl KernelOutputStored for MappedLaunchPartition Vec { self.partition.shape_as_i32() } + fn storage_id(&self) -> StorageId { + KernelOutputStored::storage_id(&self.partition) + } } impl KernelOutput for Partition> { @@ -2161,6 +2180,8 @@ pub trait KernelInputStored: Send { fn strides(&self) -> &[i32]; fn spec(&self) -> &SpecializationBits; fn dtype_str(&self) -> &'static str; + /// Identity of the backing allocation. + fn storage_id(&self) -> StorageId; } /// Converts a user-provided kernel input into a stored form for execution, @@ -2207,6 +2228,9 @@ impl KernelInputStored for Arc> { fn dtype_str(&self) -> &'static str { T::DTYPE.as_str() } + fn storage_id(&self) -> StorageId { + Tensor::storage_id(self) + } } impl KernelInputStored for &Tensor { @@ -2236,6 +2260,9 @@ impl KernelInputStored for &Tensor { fn dtype_str(&self) -> &'static str { T::DTYPE.as_str() } + fn storage_id(&self) -> StorageId { + Tensor::storage_id(self) + } } // ── KernelInput impls ─────────────────────────────────────────────────────── @@ -2304,6 +2331,9 @@ impl<'a, T: DType + Sync> KernelInputStored for &'a TensorView<'a, T> { fn dtype_str(&self) -> &'static str { T::DTYPE.as_str() } + fn storage_id(&self) -> StorageId { + self.base.storage_id() + } } impl<'a, T: DType + Sync> KernelInput for &'a TensorView<'a, T> { @@ -2324,6 +2354,161 @@ impl<'a, T: DType + Sync> IntoDeviceOp<&'a TensorView<'a, T>> for &'a TensorView } } +// ── Atomic binding ────────────────────────────────────────────────────────── + +/// Opaque identity of a backing allocation used for launch alias checks. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct StorageId(usize); + +impl StorageId { + /// Synthetic identity for launch-validation tests. + #[cfg(test)] + pub(crate) fn fake(n: usize) -> Self { + StorageId(n) + } +} + +/// Exclusive host binding for an aggregation parameter. +/// +/// ```rust,ignore +/// // `out` must already hold the identity (0.0) for `atomic_add`. +/// let mut out = api::zeros::(&[256, 256]).sync_on(&stream)?; +/// split_k(out.atomic(), lhs, rhs).grid((2, 2, 4)).sync_on(&stream)?; +/// ``` +/// +/// The backing must be initialized by the caller; floating-point atomics are +/// not reproducible. +#[derive(Debug)] +pub struct Atomic(B); + +impl Atomic> { + /// Bind an owned tensor atomically. + pub fn owned(tensor: Tensor) -> Self { + Atomic(tensor) + } +} + +impl Tensor { + /// Binds this tensor as an atomic destination. + pub fn atomic(&mut self) -> Atomic<&mut Tensor> { + Atomic(self) + } +} + +/// Host binding accepted by aggregation kernel parameters. It is sealed to the +/// two [`Atomic`] forms so their write-access tracking cannot be bypassed. +pub trait KernelAtomic: KernelInput + atomic_sealed::Sealed {} + +/// Seals [`KernelAtomic`] to exclusive [`Atomic`] backings. +mod atomic_sealed { + pub trait Sealed {} + impl Sealed for super::Atomic> {} + impl Sealed for super::Atomic<&mut super::Tensor> {} +} + +fn atomic_binding_retain( + tensor: &Tensor, + ctx: &ExecutionContext, +) -> Result<(), DeviceError> { + // Write access: the kernel mutates this storage, atomically. + tensor.storage.retain(ctx, true) +} + +fn atomic_binding_args(tensor: &Tensor, launcher: &mut AsyncKernelLaunch) { + unsafe { + launcher.push_device_ptr(tensor.cu_deviceptr()); + } + for dim in tensor.shape.iter() { + launcher.push_arg(*dim); + } + for stride in tensor.strides.iter() { + launcher.push_arg(*stride); + } +} + +impl Atomic> { + fn backing(&self) -> &Tensor { + &self.0 + } +} + +impl Atomic<&mut Tensor> { + fn backing(&self) -> &Tensor { + self.0 + } +} + +macro_rules! atomic_input_stored { + ($backing:ty) => { + impl KernelInputStored for Atomic<$backing> { + fn retain(&self, ctx: &ExecutionContext) -> Result<(), DeviceError> { + atomic_binding_retain(self.backing(), ctx) + } + fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + atomic_binding_args(self.backing(), launcher) + } + fn shape(&self) -> &[i32] { + Tensor::shape(self.backing()) + } + fn strides(&self) -> &[i32] { + Tensor::strides(self.backing()) + } + fn spec(&self) -> &SpecializationBits { + Tensor::spec(self.backing()) + } + fn dtype_str(&self) -> &'static str { + T::DTYPE.as_str() + } + fn storage_id(&self) -> StorageId { + self.backing().storage_id() + } + } + }; +} + +atomic_input_stored!(Tensor); +atomic_input_stored!(&mut Tensor); + +impl KernelInput for Atomic> { + type Stored = Atomic>; + type Returned = Tensor; + fn prepare(self) -> Self::Stored { + self + } + fn recover(stored: Self::Stored) -> Tensor { + stored.0 + } +} + +impl KernelAtomic for Atomic> {} + +impl<'a, T: DType> KernelInput for Atomic<&'a mut Tensor> { + type Stored = Atomic<&'a mut Tensor>; + type Returned = &'a mut Tensor; + fn prepare(self) -> Self::Stored { + self + } + fn recover(stored: Self::Stored) -> &'a mut Tensor { + stored.0 + } +} + +impl KernelAtomic for Atomic<&mut Tensor> {} + +impl IntoDeviceOp>> for Atomic> { + type Op = Value>>; + fn into_op(self) -> Value>> { + value(self) + } +} + +impl<'a, T: DType> IntoDeviceOp>> for Atomic<&'a mut Tensor> { + type Op = Value>>; + fn into_op(self) -> Value>> { + value(self) + } +} + #[cfg(test)] mod tests { use super::*; @@ -2355,6 +2540,34 @@ mod tests { assert_eq!(addressable_bytes::(&[0], &[1]), 0); } + /// Views over one allocation share a `StorageId`; distinct allocations do + /// not. The identity is what launch-time alias validation compares, and + /// the type stays opaque: equality is its whole public surface. + #[test] + fn storage_identity_follows_the_allocation() { + let owner: Arc = Arc::new(FakeAlloc { len_bytes: 124 }); + // SAFETY: metadata-only test; the pointer is never dereferenced and + // nothing else aliases the fake allocation. + let a = unsafe { Tensor::::from_foreign(owner.clone(), vec![4], vec![10]) }; + // SAFETY: as above; `b` shares `a`'s allocation by construction. + let b = unsafe { Tensor::::from_foreign(owner, vec![2], vec![1]) }; + let other: Arc = Arc::new(FakeAlloc { len_bytes: 124 }); + // SAFETY: as above. + let c = unsafe { Tensor::::from_foreign(other, vec![4], vec![10]) }; + assert_ne!(a.storage_id(), c.storage_id(), "distinct allocations"); + // `from_foreign` wraps each call in its own storage, so equality is + // asserted on the alias path real kernels use. + let id = c.storage_id(); + // SAFETY: metadata-only; the alias is dropped before `c`. + let alias = unsafe { c.into_shared_alias() }; + assert_eq!( + alias.storage_id(), + id, + "an alias shares its source's identity" + ); + drop((a, b, alias)); + } + #[test] fn from_foreign_accepts_allocation_covering_strided_extent() { // 4 elements at stride 10 → needs 31 f32 = 124 bytes. diff --git a/cutile-rs/cutile/src/tile_kernel.rs b/cutile-rs/cutile/src/tile_kernel.rs index 503c826e95..181fdb68f2 100644 --- a/cutile-rs/cutile/src/tile_kernel.rs +++ b/cutile-rs/cutile/src/tile_kernel.rs @@ -1088,6 +1088,93 @@ pub fn validate_grid_bounds(grid: (u32, u32, u32), bounds: &[GridBound]) -> Resu Ok(()) } +/// How one launch binds one allocation, including an aggregation operation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AggregationOp { + Add, + And, + Or, + Xor, + Max, + Min, +} + +impl AggregationOp { + fn as_str(self) -> &'static str { + match self { + AggregationOp::Add => "addition", + AggregationOp::And => "bitwise AND", + AggregationOp::Or => "bitwise OR", + AggregationOp::Xor => "bitwise XOR", + AggregationOp::Max => "maximum", + AggregationOp::Min => "minimum", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BindingAccess { + /// A whole-tensor read (`&Tensor`). + Read, + /// A tiled, disjoint write (`&mut Tensor`). + Write, + /// Overlapping atomic updates for the specified operation. + Atomic(AggregationOp), +} + +/// Rejects an atomic/ordinary alias and mismatched aggregation operations. +pub fn validate_binding_aliases( + bindings: &[(&str, crate::tensor::StorageId, BindingAccess)], +) -> Result<(), Error> { + for (i, (name, id, access)) in bindings.iter().enumerate() { + for (other_name, other_id, other_access) in bindings.iter().skip(i + 1) { + if other_id != id { + continue; + } + match (access, other_access) { + ( + BindingAccess::Read | BindingAccess::Write, + BindingAccess::Read | BindingAccess::Write, + ) => {} + (BindingAccess::Atomic(op), BindingAccess::Atomic(other_op)) if op == other_op => {} + // Allocation-level check; disjoint regions are rejected too. + (BindingAccess::Atomic(op), BindingAccess::Atomic(other_op)) => { + return kernel_launch_error_result(&format!( + "`{name}` and `{other_name}` are the same allocation, but \ + `{name}` aggregates by {} while `{other_name}` aggregates by \ + {}. The aggregation contract requires one operation per \ + allocation within a launch; bind a separate allocation, or \ + make both parameters the same operation.", + op.as_str(), + other_op.as_str() + )); + } + (BindingAccess::Atomic(_), ordinary) | (ordinary, BindingAccess::Atomic(_)) => { + let (atomic_name, ordinary_name) = if matches!(access, BindingAccess::Atomic(_)) + { + (name, other_name) + } else { + (other_name, name) + }; + let kind = match ordinary { + BindingAccess::Read => "read", + BindingAccess::Write => "written", + BindingAccess::Atomic(_) => unreachable!("the ordinary side of the pair"), + }; + return kernel_launch_error_result(&format!( + "`{atomic_name}` and `{ordinary_name}` are the same allocation, but \ + `{atomic_name}` is bound atomically while `{ordinary_name}` is {kind} \ + ordinarily. The aggregation contract admits no ordinary access to \ + the same storage within the launch; bind a separate allocation, or \ + make both parameters the same aggregation operation." + )); + } + } + } + } + Ok(()) +} + /// Runs the full set of launch-time checks before `cuLaunchKernel`: the /// built-in grid family (all partition grids match the launch grid) plus any /// compiler-emitted checks hoisted out of the device kernel. @@ -2101,6 +2188,131 @@ mod launch_check_tests { assert!(validate_grid_bounds((3, 1, 1), &[Exact((3, 1, 1)), AtMost((4, 1, 1))]).is_ok()); } + #[test] + fn an_atomic_binding_imposes_no_grid_bound() { + // The atomic capability contributes no `GridBound` at all, so a + // split-K grid with more programs than output tiles is admitted. + // That is the whole point: the exact-coverage bound a `&mut Tensor` + // output imposes would reject it. + assert_eq!( + infer_launch_grid((2, 2, 4), &[]).unwrap(), + (2, 2, 4), + "an explicit grid with no bindings to bound it must be accepted" + ); + assert!( + infer_launch_grid((2, 2, 4), &[GridBound::Exact((2, 2, 1))]).is_err(), + "the same grid under a disjoint-output bound is a launch error" + ); + // With no binding to define it, a grid is still required. + assert!(infer_launch_grid((0, 0, 0), &[]).is_err()); + } + + #[test] + fn atomic_and_ordinary_bindings_of_one_allocation_are_rejected() { + use crate::tensor::StorageId; + use AggregationOp::Add; + use BindingAccess::{Atomic, Read, Write}; + // Distinct allocations: every combination is fine. + assert!(validate_binding_aliases(&[ + ("out", StorageId::fake(1), Atomic(Add)), + ("src", StorageId::fake(2), Read) + ]) + .is_ok()); + assert!(validate_binding_aliases(&[ + ("out", StorageId::fake(1), Atomic(Add)), + ("other", StorageId::fake(2), Write) + ]) + .is_ok()); + // Two same-operation bindings over one allocation are the + // capability's point. + assert!(validate_binding_aliases(&[ + ("a", StorageId::fake(1), Atomic(Add)), + ("b", StorageId::fake(1), Atomic(Add)) + ]) + .is_ok()); + // Ordinary combinations keep whatever rules already governed them. + assert!(validate_binding_aliases(&[ + ("a", StorageId::fake(1), Read), + ("b", StorageId::fake(1), Write) + ]) + .is_ok()); + + // One allocation bound both atomically and ordinarily is the race. + let err = validate_binding_aliases(&[ + ("out", StorageId::fake(7), Atomic(Add)), + ("src", StorageId::fake(7), Read), + ]) + .unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("out") && message.contains("src") && message.contains("read"), + "the error should name both parameters and the ordinary access: {message}" + ); + // Order-independent: the atomic binding may come second. + assert!(validate_binding_aliases(&[ + ("src", StorageId::fake(7), Read), + ("out", StorageId::fake(7), Atomic(Add)) + ]) + .is_err()); + assert!(validate_binding_aliases(&[ + ("dst", StorageId::fake(7), Write), + ("out", StorageId::fake(7), Atomic(Add)) + ]) + .is_err()); + assert!(validate_binding_aliases(&[ + ("out", StorageId::fake(7), Atomic(Add)), + ("dst", StorageId::fake(7), Write) + ]) + .is_err()); + } + + #[test] + fn mixed_aggregation_operations_on_one_allocation_are_rejected() { + use crate::tensor::StorageId; + use AggregationOp::{Add, And, Max, Min, Or, Xor}; + use BindingAccess::Atomic; + let ops = [Add, And, Or, Xor, Max, Min]; + for (i, &op) in ops.iter().enumerate() { + // Same operation on one allocation: accepted, in either order. + assert!(validate_binding_aliases(&[ + ("a", StorageId::fake(1), Atomic(op)), + ("b", StorageId::fake(1), Atomic(op)) + ]) + .is_ok()); + // Different allocations: any pair of operations is accepted. + let next = ops[(i + 1) % ops.len()]; + assert!(validate_binding_aliases(&[ + ("a", StorageId::fake(1), Atomic(op)), + ("b", StorageId::fake(2), Atomic(next)) + ]) + .is_ok()); + // Different operations on one allocation: rejected, in either + // order, with both parameters and both operations named. + for &other in ops.iter().filter(|&&other| other != op) { + let err = validate_binding_aliases(&[ + ("a", StorageId::fake(3), Atomic(op)), + ("b", StorageId::fake(3), Atomic(other)), + ]) + .unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("`a`") + && message.contains("`b`") + && message.contains("one operation per"), + "the error should name both parameters and the contract: {message}" + ); + assert!( + validate_binding_aliases(&[ + ("b", StorageId::fake(3), Atomic(other)), + ("a", StorageId::fake(3), Atomic(op)) + ]) + .is_err(), + "rejection must not depend on binding order" + ); + } + } + } + #[test] fn prefix_bound_cannot_define_the_launch_grid() { use GridBound::{AtMost, Exact}; diff --git a/cutile-rs/cutile/tests/aggregation_tensor.rs b/cutile-rs/cutile/tests/aggregation_tensor.rs new file mode 100644 index 0000000000..7b5b7f4544 --- /dev/null +++ b/cutile-rs/cutile/tests/aggregation_tensor.rs @@ -0,0 +1,1671 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-only coverage for the aggregation capabilities and their launch +//! obligations. Assertions inspect the structural Tile IR module. + +use std::collections::BTreeMap; + +use cutile::compile_api::KernelCompiler; +use cutile::tile_kernel::validate_launch_checks; +use cutile_ir::bytecode::Opcode; +use cutile_ir::ir::{Attribute, Module, OpId, Operation, Type}; + +mod common; + +#[cutile::module] +mod aggregation_tensor_kernels { + use cutile::core::*; + + /// Split-K matmul through the atomic capability. + #[cutile::entry()] + fn split_k_atomic( + out: &AddTensor, + lhs: &Tensor, + rhs: &Tensor, + ) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let m: i32 = bid.0; + let n: i32 = bid.1; + let k: i32 = bid.2; + + let lhs_part: Partition = lhs.partition(shape![16, 32]); + let rhs_part: Partition = rhs.partition(shape![32, 16]); + let a: Tile = lhs_part.load([m, k]); + let b: Tile = rhs_part.load([k, n]); + let zero: Tile = constant(0.0f32, shape![16, 16]); + let partial: Tile = mma(a, b, zero); + + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([m, n], partial); + } + + #[cutile::entry()] + fn add_f16(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(f16::ONE, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_bf16(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(bf16::ONE, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_i32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_u32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn accumulate_rank1(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f32, shape![64]); + let acc: AddPartition = out.partition(shape![64]); + let _t: Token = acc.atomic_add([bid.0], v); + } + + #[cutile::entry()] + fn accumulate_rank3(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f32, shape![4, 8, 16]); + let acc: AddPartition = out.partition(shape![4, 8, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1, bid.2], v); + } + + #[cutile::entry()] + fn accumulate_const_generic(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f32, shape![BM, BN]); + let acc: AddPartition = out.partition(shape![BM, BN]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn static_ragged(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn static_divisible(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_f64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } +} + +#[cutile::module] +mod raw_intrinsic_kernels { + use cutile::core::*; + use cutile::tileir::*; + + #[cutile::entry()] + unsafe fn split_k_raw( + out: &Tensor, + lhs: &Tensor, + rhs: &Tensor, + ) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let m: i32 = bid.0; + let n: i32 = bid.1; + let k: i32 = bid.2; + + let lhs_part: Partition = lhs.partition(shape![16, 32]); + let rhs_part: Partition = rhs.partition(shape![32, 16]); + let a: Tile = lhs_part.load([m, k]); + let b: Tile = rhs_part.load([k, n]); + let zero: Tile = constant(0.0f32, shape![16, 16]); + let partial: Tile = mma(a, b, zero); + + let mut view: PartitionMut = + unsafe { out.partition_full_mut(shape![16, 16]) }; + let _done: Token = unsafe { + atomic_red_view_tko( + &mut view, + partial, + [m, n], + atomic::AddF, + ordering::Relaxed, + scope::Device, + ) + }; + } +} + +// ── Host binding ──────────────────────────────────────────────────────────── + +/// Pins a launcher's result type, so the recovery assertions below are real +/// type equalities rather than inference. +fn assert_launch_output>(_: &D) {} + +/// Compile-only: the host side of the capability typechecks end to end. +/// +/// Never called, since there is no GPU here, but building the launcher makes +/// rustc check the `KernelAtomic` bound on the exclusive borrow, the ordinary +/// `KernelInput` bound on the two shared inputs beside it, the explicit +/// split-K grid, and that the atomic parameter recovers as the backing borrow +/// rather than as the `Atomic<..>` wrapper. The negative direction (a shared +/// `Arc>` bound atomically) is a compile-fail case in `tests/ui/`. +#[allow(dead_code)] +fn borrowed_host_binding_typechecks<'a>( + out: &'a mut cutile::tensor::Tensor, + lhs: std::sync::Arc>, + rhs: std::sync::Arc>, +) { + use cutile::tile_kernel::TileKernel; + let launch = aggregation_tensor_kernels::split_k_atomic(out.atomic(), lhs, rhs).grid((2, 2, 4)); + assert_launch_output::< + ( + &'a mut cutile::tensor::Tensor, + std::sync::Arc>, + std::sync::Arc>, + ), + _, + >(&launch); +} + +/// Compile-only: an owned tensor is the other exclusive authority, and it +/// recovers as the plain `Tensor` it was. +#[allow(dead_code)] +fn owned_host_binding_typechecks( + out: cutile::tensor::Tensor, + lhs: std::sync::Arc>, + rhs: std::sync::Arc>, +) { + use cutile::tile_kernel::TileKernel; + let launch = + aggregation_tensor_kernels::split_k_atomic(cutile::tensor::Atomic::owned(out), lhs, rhs) + .grid((2, 2, 4)); + assert_launch_output::< + ( + cutile::tensor::Tensor, + std::sync::Arc>, + std::sync::Arc>, + ), + _, + >(&launch); +} + +// ── Tile IR enum values asserted below ────────────────────────────────────── +// +const MODE_ADD_F: i64 = 4; +const MODE_ADD: i64 = 3; +const MODE_AND: i64 = 0; +const MODE_OR: i64 = 1; +const MODE_XOR: i64 = 2; +const MODE_MAX: i64 = 5; +const MODE_MIN: i64 = 6; +const MODE_UMAX: i64 = 7; +const MODE_UMIN: i64 = 8; +const ORDERING_RELAXED: i64 = 1; +const SCOPE_DEVICE: i64 = 1; + +const ROW_MAJOR_1D: &[i32] = &[1]; +const ROW_MAJOR_2D: &[i32] = &[-1, 1]; +const ROW_MAJOR_3D: &[i32] = &[-1, -1, 1]; + +const MATMUL_STRIDES: &[(&str, &[i32])] = &[ + ("out", ROW_MAJOR_2D), + ("lhs", ROW_MAJOR_2D), + ("rhs", ROW_MAJOR_2D), +]; + +fn compiler( + module_ast: fn() -> cutile_compiler::ast::Module, + module: &str, + func: &str, + strides: &[(&str, &[i32])], +) -> KernelCompiler cutile_compiler::ast::Module> { + KernelCompiler::new(module_ast, module, func) + .target("sm_120") + .strides(strides) +} + +fn compile( + module_ast: fn() -> cutile_compiler::ast::Module, + module: &str, + func: &str, + strides: &[(&str, &[i32])], +) -> Module { + compile_with_counts(module_ast, module, func, strides).0 +} + +fn compile_with_counts( + module_ast: fn() -> cutile_compiler::ast::Module, + module: &str, + func: &str, + strides: &[(&str, &[i32])], +) -> (Module, cutile_compiler::compile_api::CheckPlacementCounts) { + let artifacts = compiler(module_ast, module, func, strides) + .compile() + .expect("compilation failed"); + let counts = artifacts.check_counts(); + (artifacts.into_module(), counts) +} + +fn atomic_module() -> Module { + compile( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + "split_k_atomic", + MATMUL_STRIDES, + ) +} + +fn raw_module() -> Module { + compile( + raw_intrinsic_kernels::__module_ast_self, + "raw_intrinsic_kernels", + "split_k_raw", + MATMUL_STRIDES, + ) +} + +fn ops(module: &Module) -> Vec { + fn walk(module: &Module, op: OpId, out: &mut Vec) { + out.push(op); + for ®ion in &module.op(op).regions { + for &block in &module.region(region).blocks { + for &inner in &module.block(block).ops { + walk(module, inner, out); + } + } + } + } + let mut out = Vec::new(); + for &function in &module.functions { + walk(module, function, &mut out); + } + out +} + +fn ops_with(module: &Module, opcode: Opcode) -> Vec<&Operation> { + ops(module) + .into_iter() + .map(|id| module.op(id)) + .filter(|op| op.opcode == opcode) + .collect() +} + +/// The module's single atomic reduction. +fn atomic_reduction(module: &Module) -> &Operation { + let found = ops_with(module, Opcode::AtomicRedViewTko); + assert_eq!( + found.len(), + 1, + "expected exactly one atomic reduction, got {}", + found.len() + ); + found[0] +} + +fn int_attr(op: &Operation, name: &str) -> i64 { + match op + .attributes + .iter() + .find(|(key, _)| key == name) + .map(|(_, value)| value) + { + Some(Attribute::Integer(value, _)) => *value, + other => panic!("expected an integer `{name}` attribute, got {other:?}"), + } +} + +#[test] +fn atomic_add_emits_the_expected_atomic_reduction() { + common::with_test_stack(|| { + let module = atomic_module(); + let op = atomic_reduction(&module); + + assert_eq!(int_attr(op, "mode"), MODE_ADD_F, "floating-add mode"); + assert_eq!( + int_attr(op, "memory_ordering_semantics"), + ORDERING_RELAXED, + "relaxed ordering" + ); + assert_eq!(int_attr(op, "memory_scope"), SCOPE_DEVICE, "device scope"); + + // No old-value result: the completion token is the only result. + assert_eq!( + op.result_types, + vec![Type::Token], + "the reduction must discard old values and yield only its token" + ); + + // The view operand is an unpadded partition view whose tile shape is + // the contributed tile's. + let view = module + .value_type(op.operands[0]) + .without_pointer_attribute(); + let Type::PartitionView(view) = view else { + panic!("expected a partition view operand, got {view:?}"); + }; + assert_eq!(view.tile_shape, vec![16, 16], "tile geometry from the tile"); + assert_eq!( + view.padding_value, None, + "the atomic view must be unpadded, got {:?}", + view.padding_value + ); + + // Nothing stored the contribution ordinarily. + assert!( + ops_with(&module, Opcode::StoreViewTko).is_empty(), + "an ordinary store must not stand in for the atomic reduction" + ); + + // The capability itself lowers to a plain tensor view. + assert_eq!( + ops_with(&module, Opcode::MakeTensorView).len(), + 3, + "one tensor view per tensor parameter" + ); + }); +} + +/// The attributes and operand/result arity of one op, which is the +/// "operation profile" the two kernels must agree on. Whole-module equality is +/// not used: SSA numbering and the entry symbol differ by construction. +fn op_shape(module: &Module, op: &Operation) -> String { + let mut attrs: BTreeMap<&str, String> = BTreeMap::new(); + for (key, value) in &op.attributes { + attrs.insert(key.as_str(), format!("{value:?}")); + } + let operands: Vec = op + .operands + .iter() + .map(|&v| format!("{:?}", module.value_type(v))) + .collect(); + format!( + "{:?} operands={:?} results={:?} attrs={:?}", + op.opcode, operands, op.result_types, attrs + ) +} + +fn opcode_profile(module: &Module) -> BTreeMap { + let mut counts = BTreeMap::new(); + for id in ops(module) { + *counts + .entry(format!("{:?}", module.op(id).opcode)) + .or_insert(0) += 1; + } + counts +} + +#[test] +fn atomic_add_matches_the_raw_intrinsic_kernel() { + common::with_test_stack(|| { + let atomic = atomic_module(); + let raw = raw_module(); + + // Repeatability first: comparing two kernels by a profile proves + // nothing unless the profile is stable for one kernel. + assert_eq!( + opcode_profile(&atomic), + opcode_profile(&atomic_module()), + "the opcode profile is not reproducible across compiles" + ); + assert_eq!( + op_shape(&atomic, atomic_reduction(&atomic)), + op_shape(&atomic, atomic_reduction(&atomic_module())), + "the reduction's profile is not reproducible across compiles" + ); + + // `atomic_add` bounds-checks its index where the raw intrinsic does + // not, so the profiles are compared modulo `Constant`: the discharged + // proof leaves a few folded constants behind. Every other opcode must + // match, and in particular there must be no `Assert` — a device-side + // bounds check would show up here. + let drop_constants = |m: BTreeMap| -> BTreeMap { + m.into_iter().filter(|(k, _)| k != "Constant").collect() + }; + assert_eq!( + drop_constants(opcode_profile(&atomic)), + drop_constants(opcode_profile(&raw)), + "`atomic_add` must emit the same operations as the raw intrinsic, \ + constants aside" + ); + assert_eq!( + op_shape(&atomic, atomic_reduction(&atomic)), + op_shape(&raw, atomic_reduction(&raw)), + "the reduction itself must be identical in operands, results and \ + attributes to the raw intrinsic's" + ); + }); +} + +#[test] +fn shape_is_type_metadata_not_a_device_read() { + common::with_test_stack(|| { + let module = compile( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + "split_k_atomic", + MATMUL_STRIDES, + ); + assert!( + ops_with(&module, Opcode::GetTensorShape).is_empty(), + "the view's shape metadata must not emit a device shape query" + ); + }); +} + +#[test] +fn every_float_element_reaches_the_same_reduction() { + common::with_test_stack(|| { + for func in ["add_f16", "add_bf16", "add_f64"] { + let module = compile( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + func, + &[("out", ROW_MAJOR_2D)], + ); + let op = atomic_reduction(&module); + assert_eq!( + int_attr(op, "mode"), + MODE_ADD_F, + "{func}: floating-add mode" + ); + assert_eq!( + int_attr(op, "memory_ordering_semantics"), + ORDERING_RELAXED, + "{func}: relaxed" + ); + assert_eq!(int_attr(op, "memory_scope"), SCOPE_DEVICE, "{func}: device"); + assert_eq!(op.result_types, vec![Type::Token], "{func}: token only"); + } + }); +} + +#[test] +fn every_rank_and_a_const_generic_tile_reach_the_same_reduction() { + common::with_test_stack(|| { + // (entry, stride hints, generic arguments, expected tile geometry) + let cases: Vec<(&str, &[(&str, &[i32])], Vec, Vec)> = vec![ + ( + "accumulate_rank1", + &[("out", ROW_MAJOR_1D)], + vec![], + vec![64], + ), + ( + "accumulate_rank3", + &[("out", ROW_MAJOR_3D)], + vec![], + vec![4, 8, 16], + ), + ( + "accumulate_const_generic", + &[("out", ROW_MAJOR_2D)], + vec!["16".into(), "32".into()], + vec![16, 32], + ), + ]; + + for (func, strides, generics, tile) in cases { + let module = compiler( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + func, + strides, + ) + .generics(generics) + .compile() + .unwrap_or_else(|error| panic!("{func} failed to compile: {error}")) + .into_module(); + + let op = atomic_reduction(&module); + assert_eq!(int_attr(op, "mode"), MODE_ADD_F, "{func}: floating-add"); + assert_eq!( + int_attr(op, "memory_ordering_semantics"), + ORDERING_RELAXED, + "{func}: relaxed" + ); + assert_eq!(int_attr(op, "memory_scope"), SCOPE_DEVICE, "{func}: device"); + assert_eq!(op.result_types, vec![Type::Token], "{func}: token only"); + + let view = module + .value_type(op.operands[0]) + .without_pointer_attribute(); + let Type::PartitionView(view) = view else { + panic!("{func}: expected a partition view operand, got {view:?}"); + }; + assert_eq!(view.tile_shape, tile, "{func}: tile geometry"); + assert_eq!(view.padding_value, None, "{func}: unpadded"); + + assert_eq!( + op.operands.len(), + tile.len() + 3, + "{func}: one index component per axis" + ); + } + }); +} + +/// `atomic_add` is bounds-checked, and for a `tile_block_id` coordinate the +/// check costs nothing on the device: the block-id axiom rung reduces it to a +/// claim over launch-known operands. This asserts the placement rather than +/// assuming it — the counters are the compiler's own report, and an `Assert` +/// op in the module would mean a per-thread branch instead. +#[test] +fn a_block_id_index_discharges_its_bounds_check() { + common::with_test_stack(|| { + let (module, counts) = compile_with_counts( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + "split_k_atomic", + MATMUL_STRIDES, + ); + assert!( + counts.in_place == 0, + "expected no device-side check, got {counts:?}" + ); + assert!( + counts.discharged > 0, + "expected the index proof to be discharged, got {counts:?}" + ); + assert!( + ops_with(&module, Opcode::Assert).is_empty(), + "a discharged check must leave no `Assert` in the module" + ); + }); +} + +/// The divisibility obligations the compiler evacuated for parameter +/// `param`, as `axis -> divisor`. +/// +/// Reads the [`Predicate`] structurally, not the cause string. The same +/// kernel also evacuates grid-bound checks stated over the *same* extents +/// (`num_tile_blocks(k) <= ceil(extent(out, k)/16)`), so "a launch check +/// exists" is not evidence that the tiling obligation was emitted; only a +/// `DivisibleBy` whose term is exactly `1 * Dim { param, axis }` is. +fn divisibility_obligations( + checks: &[cuda_async::predicate::LaunchCheck], + param: usize, +) -> BTreeMap { + use cuda_async::predicate::{Atom, Predicate}; + let mut found = BTreeMap::new(); + for check in checks { + let Predicate::DivisibleBy { term, divisor } = &check.predicate else { + continue; + }; + let coeffs = term.coeffs(); + if coeffs.len() != 1 || term.constant_part() != 0 { + continue; + } + let (&atom, &coeff) = coeffs.iter().next().expect("one coefficient"); + let Atom::Dim { param: p, axis } = atom else { + continue; + }; + if p != param || coeff != 1 { + continue; + } + let previous = found.insert(axis, *divisor); + assert!( + previous.is_none(), + "two divisibility obligations for param {param} axis {axis}" + ); + } + found +} + +/// Compiles `split_k_atomic` and hands back its hoisted launch checks. +fn split_k_launch_checks() -> Vec { + compiler( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + "split_k_atomic", + MATMUL_STRIDES, + ) + .compile() + .expect("compilation failed") + .launch_checks() + .to_vec() +} + +/// A ragged destination must be rejected before any GPU work. +/// +/// A Tile IR partition index space is a ceiling division, so a 20-row tensor +/// tiled by 16 has a legitimate index 1 selecting a *partial* tile. Tile IR +/// documents masking for ordinary stores to partial tiles but does not grant +/// the same to `atomic_red_view_tko`, whose view must be unpadded — so a safe +/// reduction requires an exact tiling. `partition` emits that as a launch +/// obligation per axis; this asserts the obligations are there on *both* +/// output axes, and then evaluates the real emitted checks against a shape +/// that is ragged on each axis in turn. +#[test] +fn a_ragged_destination_is_rejected_before_launch() { + common::with_test_stack(|| { + let checks = split_k_launch_checks(); + + // `out` is parameter 0 and the kernel tiles it [16, 16]. Both axes + // must carry the obligation: one would leave the other edge ragged. + assert_eq!( + divisibility_obligations(&checks, 0), + BTreeMap::from([(0usize, 16i64), (1usize, 16i64)]), + "both output axes must carry a divisibility obligation naming the \ + tile extent; got {checks:#?}" + ); + // And the causes must explain themselves, since a failing launch + // shows the user nothing else. + for check in &checks { + if matches!( + check.predicate, + cuda_async::predicate::Predicate::DivisibleBy { .. } + ) { + assert!( + check + .cause + .contains("multiple of the aggregation tile extent 16"), + "unhelpful cause: {}", + check.cause + ); + } + } + + // The inputs are read through ordinary partitions, which tolerate a + // partial edge tile; no divisibility obligation belongs to them. + assert!( + divisibility_obligations(&checks, 1).is_empty() + && divisibility_obligations(&checks, 2).is_empty(), + "only the aggregation destination needs an exact tiling: {checks:#?}" + ); + + // Now evaluate the checks the compiler actually emitted. The grid is + // (2, 3, 4) throughout, and every `out` shape below satisfies the + // grid-bound checks, so the only thing that can reject a launch here + // is the tiling. + let grid = (2u32, 3u32, 4u32); + let shapes = |out: Vec| vec![out, vec![32, 128], vec![128, 48]]; + + let exact = shapes(vec![32, 48]); + validate_launch_checks(&checks, &exact, &exact, grid) + .expect("a tiling that divides exactly must be accepted"); + + for (ragged, axis) in [(vec![20, 48], 0), (vec![32, 50], 1)] { + let ragged_shapes = shapes(ragged.clone()); + let Err(err) = validate_launch_checks(&checks, &ragged_shapes, &ragged_shapes, grid) + else { + panic!("a destination ragged on axis {axis} ({ragged:?}) must be rejected"); + }; + let message = err.to_string(); + assert!( + message.contains(&format!("axis {axis} must be a multiple of")), + "the rejection must name the ragged axis ({axis}): {message}" + ); + } + }); +} + +/// The converse of the static case: an extent the signature pins to a value +/// the tile does not divide is an impossibility, so the JIT rejects the kernel +/// outright rather than emitting a launch check that could never pass. +#[test] +fn a_statically_ragged_destination_fails_to_compile() { + common::with_test_stack(|| { + let result = compiler( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + "static_ragged", + &[("out", ROW_MAJOR_2D)], + ) + .compile(); + let Err(error) = result else { + panic!("a [20, 48] destination tiled [16, 16] must not compile"); + }; + let message = error.to_string(); + assert!( + message.contains("declared extent 20") && message.contains("multiple of"), + "the error must name the declared extent and the requirement: {message}" + ); + }); +} + +/// Exactly-divisible compile-time extents cost nothing. +/// +/// The obligation is the same one, but the signature pins the extents, so the +/// compiler decides it and leaves no divisibility check for the host to run +/// per launch. Asserted against the dynamic kernel, which does leave two: the +/// difference is the claim. +#[test] +fn an_exactly_divisible_static_shape_leaves_no_launch_check() { + common::with_test_stack(|| { + let artifacts = compiler( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + "static_divisible", + &[("out", ROW_MAJOR_2D)], + ) + .compile() + .expect("compilation failed"); + let checks = artifacts.launch_checks(); + + assert!( + divisibility_obligations(checks, 0).is_empty(), + "a [32, 48] destination tiled [16, 16] divides exactly on both axes, \ + so nothing should be left for the host to re-check: {checks:#?}" + ); + // The reduction is still there — the check was discharged, not the + // kernel silently reshaped. + assert_eq!( + atomic_reduction(artifacts.module()).result_types, + vec![Type::Token] + ); + // Control: the same tiling over `{ [-1, -1] }` does leave two. + assert_eq!( + divisibility_obligations(&split_k_launch_checks(), 0).len(), + 2 + ); + }); +} + +#[test] +fn the_capability_is_not_partitioned_by_the_entry_wrapper() { + common::with_test_stack(|| { + let module = atomic_module(); + // Three partition views: `lhs`, `rhs`, and the one the reduction + // builds internally. A fourth would mean the entry wrapper + // partitioned the capability the way it partitions a `&mut Tensor`. + assert_eq!( + ops_with(&module, Opcode::MakePartitionView).len(), + 3, + "expected one view per input partition plus the reduction's own" + ); + }); +} + +/// One method spelling, two Tile IR modes, selected by the element type. +/// +/// `atomic_add` is eight concrete-element impls rather than one generic impl, +/// because the JIT reads the mode from the literal spelled at the call to +/// `atomic_red_view_raw`. A generic body cannot supply it, and two impls with +/// disjoint bounds collide as duplicate definitions. This asserts the split +/// actually lands: integers reach `Add`, floats reach `AddF`. +#[test] +fn an_integer_element_reaches_the_integer_add_mode() { + common::with_test_stack(|| { + for func in ["add_i32", "add_u32"] { + let module = compile( + aggregation_tensor_kernels::__module_ast_self, + "aggregation_tensor_kernels", + func, + &[("out", ROW_MAJOR_2D)], + ); + let op = atomic_reduction(&module); + assert_eq!(int_attr(op, "mode"), MODE_ADD, "{func}: integer add mode"); + assert_ne!( + int_attr(op, "mode"), + MODE_ADD_F, + "{func}: must not reach the floating-add mode" + ); + } + }); +} + +/// Kernels spelling the semantic comparison markers at the raw op, one per +/// signedness. The marker is resolved by the compiler from the Rust element +/// type, which is the only place signedness survives: the Tile IR scalar for +/// `u32` is `i32`. +#[cutile::module] +mod semantic_marker_kernels { + use cutile::core::*; + + #[cutile::entry()] + fn max_i32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Maximum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn max_u32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Maximum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn max_i64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Maximum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn max_u64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Maximum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn min_i32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Minimum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn min_u32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Minimum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn min_i64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Minimum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn min_u64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Minimum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + + #[cutile::entry()] + fn max_f32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; the index is a block id inside the grid and + // `partition` established exact tiling. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Maximum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } + /// The destination view and contribution tile deliberately disagree in + /// signedness: the raw signature ties nothing together, and both erase to + /// the same IR scalar, so only the compiler's cross-check can see it. + #[cutile::entry()] + fn max_mismatched(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + // SAFETY: probe kernel; never reaches the device. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Maximum, + ordering::Relaxed, + scope::Device, + None, + ) + }; + } +} + +fn semantic_marker_mode(func: &str) -> i64 { + let module = compile( + semantic_marker_kernels::__module_ast_self, + "semantic_marker_kernels", + func, + &[("out", ROW_MAJOR_2D)], + ); + int_attr(atomic_reduction(&module), "mode") +} + +/// `Maximum` and `Minimum` resolve the comparison from the element's +/// signedness: signed integers reach `Max`/`Min`, unsigned reach +/// `Umax`/`Umin`. Both directions are asserted, because the failure mode is +/// a wrong comparison, not a compile error: a signed mode applied to `u32` +/// values straddling the high bit silently picks the smaller value. +#[test] +fn comparison_markers_resolve_by_signedness() { + common::with_test_stack(|| { + for (func, expect, wrong) in [ + ("max_i32", MODE_MAX, MODE_UMAX), + ("max_u32", MODE_UMAX, MODE_MAX), + ("max_i64", MODE_MAX, MODE_UMAX), + ("max_u64", MODE_UMAX, MODE_MAX), + ("min_i32", MODE_MIN, MODE_UMIN), + ("min_u32", MODE_UMIN, MODE_MIN), + ("min_i64", MODE_MIN, MODE_UMIN), + ("min_u64", MODE_UMIN, MODE_MIN), + ] { + let mode = semantic_marker_mode(func); + assert_eq!(mode, expect, "{func}"); + assert_ne!(mode, wrong, "{func}: wrong signedness selected"); + } + }); +} + +/// Tile IR has no atomic floating-point maximum, so `Maximum` over a float +/// element is rejected in the JIT with a message naming both. +#[test] +fn a_float_comparison_marker_is_rejected() { + common::with_test_stack(|| { + let msg = match compiler( + semantic_marker_kernels::__module_ast_self, + "semantic_marker_kernels", + "max_f32", + &[("out", ROW_MAJOR_2D)], + ) + .compile() + { + Ok(_) => panic!("a float atomic maximum must not compile"), + Err(err) => err.to_string(), + }; + assert!( + msg.contains("`Maximum` has no atomic reduction mode for `f32`"), + "unexpected message: {msg}" + ); + }); +} + +/// The partition view an operation addresses must be unpadded: Tile IR +/// rejects a padded view for an atomic reduction, and the safe constructors +/// never create one. +fn assert_unpadded_view(module: &Module, func: &str) { + let views = ops_with(module, Opcode::MakePartitionView); + assert!(!views.is_empty(), "{func}: expected a partition view"); + for op in views { + for ty in &op.result_types { + if let Type::PartitionView(view) = ty { + assert!( + view.padding_value.is_none(), + "{func}: the view must be unpadded" + ); + } + } + } +} + +/// One kernel per legal operation/element pair: the full 28-cell matrix the +/// three element traits admit, every kernel identical except the capability +/// type, the method, and the element. +#[cutile::module] +mod op_matrix_kernels { + use cutile::core::*; + + #[cutile::entry()] + fn add_f16(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(f16::ONE, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_bf16(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(bf16::ONE, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_f32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_f64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_i32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_u32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_i64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn add_u64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn and_i32(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn and_u32(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn and_i64(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn and_u64(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn or_i32(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn or_u32(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn or_i64(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn or_u64(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn xor_i32(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn xor_u32(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn xor_i64(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn xor_u64(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn max_i32(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn max_u32(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn max_i64(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn max_u64(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn min_i32(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn min_u32(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn min_i64(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } + + #[cutile::entry()] + fn min_u64(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } +} + +/// Every supported operation/element pair compiles and emits exactly the +/// expected Tile IR: the right mode for the pair (signedness included), +/// relaxed ordering, device scope, an unpadded view, and a completion token +/// as the only result. +#[test] +fn every_operation_element_pair_reaches_its_mode() { + common::with_test_stack(|| { + for (func, expected_mode) in [ + ("add_f16", MODE_ADD_F), + ("add_bf16", MODE_ADD_F), + ("add_f32", MODE_ADD_F), + ("add_f64", MODE_ADD_F), + ("add_i32", MODE_ADD), + ("add_u32", MODE_ADD), + ("add_i64", MODE_ADD), + ("add_u64", MODE_ADD), + ("and_i32", MODE_AND), + ("and_u32", MODE_AND), + ("and_i64", MODE_AND), + ("and_u64", MODE_AND), + ("or_i32", MODE_OR), + ("or_u32", MODE_OR), + ("or_i64", MODE_OR), + ("or_u64", MODE_OR), + ("xor_i32", MODE_XOR), + ("xor_u32", MODE_XOR), + ("xor_i64", MODE_XOR), + ("xor_u64", MODE_XOR), + ("max_i32", MODE_MAX), + ("max_u32", MODE_UMAX), + ("max_i64", MODE_MAX), + ("max_u64", MODE_UMAX), + ("min_i32", MODE_MIN), + ("min_u32", MODE_UMIN), + ("min_i64", MODE_MIN), + ("min_u64", MODE_UMIN), + ] { + let module = compile( + op_matrix_kernels::__module_ast_self, + "op_matrix_kernels", + func, + &[("out", ROW_MAJOR_2D)], + ); + let op = atomic_reduction(&module); + assert_eq!(int_attr(op, "mode"), expected_mode, "{func}: mode"); + assert_eq!( + int_attr(op, "memory_ordering_semantics"), + ORDERING_RELAXED, + "{func}: ordering" + ); + assert_eq!(int_attr(op, "memory_scope"), SCOPE_DEVICE, "{func}: scope"); + assert_eq!(op.result_types, vec![Type::Token], "{func}: result"); + assert_unpadded_view(&module, func); + } + }); +} + +/// A module-local type alias names a capability the same way it names a +/// `Tensor`, and element inference reads through it. +#[cutile::module] +mod alias_kernels { + use cutile::core::*; + + type Bins = AddTensor; + + #[cutile::entry()] + fn histogram(out: &Bins) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![64]); + let acc: AddPartition = out.partition(shape![64]); + let _t: Token = acc.atomic_add([bid.0], v); + } +} + +/// The capability reached through a module-local alias lowers identically: +/// integer add mode through the alias, with the same view and token shape. +#[test] +fn a_type_alias_reaches_the_same_reduction() { + common::with_test_stack(|| { + let module = compile( + alias_kernels::__module_ast_self, + "alias_kernels", + "histogram", + &[("out", &[1])], + ); + let op = atomic_reduction(&module); + assert_eq!( + int_attr(op, "mode"), + MODE_ADD, + "integer add through the alias" + ); + assert_eq!(op.result_types, vec![Type::Token]); + }); +} + +/// A statically ragged destination fails for every family, not only addition: +/// the divisibility obligation is emitted by each capability's `partition`. +#[test] +fn a_statically_ragged_max_destination_fails_to_compile() { + common::with_test_stack(|| { + let msg = match compiler( + ragged_max_kernels::__module_ast_self, + "ragged_max_kernels", + "static_ragged_max", + &[("out", ROW_MAJOR_2D)], + ) + .compile() + { + Ok(_) => panic!("a statically ragged MaxTensor must not compile"), + Err(err) => err.to_string(), + }; + assert!( + msg.contains("not a multiple of the aggregation tile extent"), + "unexpected message: {msg}" + ); + }); +} + +#[cutile::module] +mod ragged_max_kernels { + use cutile::core::*; + + /// 20 is not divisible by 16: index 1 would select a partial tile. + #[cutile::entry()] + fn static_ragged_max(out: &MaxTensor) { + let v: Tile = constant(1i32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([0i32, 0i32], v); + } +} + +/// A semantic marker resolves from the *destination* view, and a +/// contribution tile of a different element is rejected by name: both erase +/// to the same IR scalar, so resolving from the contribution would silently +/// select the wrong comparison. +#[test] +fn a_mismatched_contribution_element_is_rejected() { + common::with_test_stack(|| { + let msg = match compiler( + semantic_marker_kernels::__module_ast_self, + "semantic_marker_kernels", + "max_mismatched", + &[("out", ROW_MAJOR_2D)], + ) + .compile() + { + Ok(_) => panic!("a u32 destination with an i32 contribution must not compile"), + Err(err) => err.to_string(), + }; + assert!( + msg.contains("`Maximum` destination is `u32`, but contribution tile is `i32`"), + "unexpected message: {msg}" + ); + }); +} + +/// The token a checked operation returns is the reduction's own completion +/// token: consumed as a raw store's ordering operand, that operand's producer +/// is the `AtomicRedViewTko` itself and its type is `Token`. The weaker +/// `result_types == [Token]` cannot tell the op's token from any other. +#[test] +fn the_returned_token_is_the_reductions_own() { + common::with_test_stack(|| { + for func in [ + "consume_add", + "consume_and", + "consume_or", + "consume_xor", + "consume_max", + "consume_min", + ] { + let module = compile( + token_consumer_kernels::__module_ast_self, + "token_consumer_kernels", + func, + &[("out", ROW_MAJOR_2D)], + ); + let reductions = ops_with(&module, Opcode::AtomicRedViewTko); + assert_eq!(reductions.len(), 2, "{func}: expected both reductions"); + let token = *module + .op(ops(&module) + .into_iter() + .filter(|&id| module.op(id).opcode == Opcode::AtomicRedViewTko) + .nth(1) + .expect("second reduction")) + .operands + .last() + .expect("a reduction has operands"); + let producer = match module.value_data(token).producer { + cutile_ir::ir::ValueProducer::OpResult { op, .. } => Some(module.op(op).opcode), + cutile_ir::ir::ValueProducer::BlockArg { .. } => None, + }; + assert_eq!( + producer, + Some(Opcode::AtomicRedViewTko), + "{func}: the consumed token must be the first reduction's result" + ); + assert_eq!(module.value_type(token), &Type::Token, "{func}"); + } + }); +} + +#[cutile::module] +mod token_consumer_kernels { + use cutile::core::*; + + #[cutile::entry()] + fn consume_add(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let done: Token = acc.atomic_add([bid.0, bid.1], v); + // SAFETY: same view, in-grid index, exact tiling from `partition`. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Sum, + ordering::Relaxed, + scope::Device, + Some(done), + ) + }; + } + #[cutile::entry()] + fn consume_and(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let done: Token = acc.atomic_and([bid.0, bid.1], v); + // SAFETY: same view, in-grid index, exact tiling from `partition`. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::And, + ordering::Relaxed, + scope::Device, + Some(done), + ) + }; + } + #[cutile::entry()] + fn consume_or(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let done: Token = acc.atomic_or([bid.0, bid.1], v); + // SAFETY: same view, in-grid index, exact tiling from `partition`. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Or, + ordering::Relaxed, + scope::Device, + Some(done), + ) + }; + } + #[cutile::entry()] + fn consume_xor(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let done: Token = acc.atomic_xor([bid.0, bid.1], v); + // SAFETY: same view, in-grid index, exact tiling from `partition`. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Xor, + ordering::Relaxed, + scope::Device, + Some(done), + ) + }; + } + #[cutile::entry()] + fn consume_max(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let done: Token = acc.atomic_max([bid.0, bid.1], v); + // SAFETY: same view, in-grid index, exact tiling from `partition`. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Maximum, + ordering::Relaxed, + scope::Device, + Some(done), + ) + }; + } + #[cutile::entry()] + fn consume_min(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let done: Token = acc.atomic_min([bid.0, bid.1], v); + // SAFETY: same view, in-grid index, exact tiling from `partition`. + let _t: Token = unsafe { + atomic_red_view_raw( + &acc, + v, + [bid.0, bid.1], + atomic::Minimum, + ordering::Relaxed, + scope::Device, + Some(done), + ) + }; + } +} diff --git a/cutile-rs/cutile/tests/gpu.rs b/cutile-rs/cutile/tests/gpu.rs index 10e48e0459..e4f58c01a7 100644 --- a/cutile-rs/cutile/tests/gpu.rs +++ b/cutile-rs/cutile/tests/gpu.rs @@ -77,6 +77,9 @@ mod mapped_partition_schedule_matrix; #[path = "gpu/atomic_red_view.rs"] mod atomic_red_view; +#[path = "gpu/aggregation_tensor.rs"] +mod aggregation_tensor; + #[path = "gpu/borrow_raw_parts.rs"] mod borrow_raw_parts; diff --git a/cutile-rs/cutile/tests/gpu/aggregation_tensor.rs b/cutile-rs/cutile/tests/gpu/aggregation_tensor.rs new file mode 100644 index 0000000000..5b0f67dc36 --- /dev/null +++ b/cutile-rs/cutile/tests/gpu/aggregation_tensor.rs @@ -0,0 +1,1003 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Device tests for the aggregation capabilities, `AddTensor` through +//! `MinTensor`. +//! +//! The compile-only suite (`tests/aggregation_tensor.rs`) proves each checked +//! operation emits the intended `cuda_tile.atomic_red_view_tko` with the +//! right mode. These tests prove the emitted op does the intended arithmetic — every program's contribution +//! lands, none is lost to a racing read-modify-write, and the sum matches a +//! host oracle — through both host bindings, the exclusive borrow +//! (`tensor.atomic()`) and the owned one (`Atomic::owned`). +//! +//! The value kernels are split-K: the grid's last axis carries several +//! programs per output tile, which is exactly the overlap an ordinary +//! `&mut Tensor` output would forbid. +//! +//! The last two tests are about the host's within-launch alias rule rather +//! than arithmetic: one allocation bound both atomically and ordinarily is +//! rejected before any GPU work, while two atomic bindings over one +//! allocation stay accepted. + +use crate::common; +use cutile::half::{bf16, f16}; +use cutile::prelude::*; +use std::sync::Arc; + +const BM: usize = 16; +const BN: usize = 16; +const BK: usize = 32; + +#[cutile::module] +mod atomic_value_kernels { + use cutile::core::*; + + /// Every program adds a ones-tile into output tile `[m, n]`. With a grid + /// of `(tiles_m, tiles_n, splits)` each tile must end at exactly + /// `splits`; a non-atomic store would lose updates and land lower. + #[cutile::entry()] + fn accumulate_ones(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let ones: Tile = constant(1.0f32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([bid.0, bid.1], ones); + } + + #[cutile::entry()] + fn accumulate_f16_ones(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let ones: Tile = constant(f16::ONE, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([bid.0, bid.1], ones); + } + + #[cutile::entry()] + fn accumulate_bf16_ones(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let ones: Tile = constant(bf16::ONE, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([bid.0, bid.1], ones); + } + + /// Split-K matmul: program `(m, n, k)` contributes `lhs[m, k] @ rhs[k, n]` + /// to output tile `[m, n]`. The K axis is divided across the grid, so the + /// output tile is only correct once every program's partial product has + /// been accumulated. + #[cutile::entry()] + fn split_k_matmul( + out: &AddTensor, + lhs: &Tensor, + rhs: &Tensor, + ) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let m: i32 = bid.0; + let n: i32 = bid.1; + let k: i32 = bid.2; + + let lhs_part: Partition = lhs.partition(shape![16, 32]); + let rhs_part: Partition = rhs.partition(shape![32, 16]); + let a: Tile = lhs_part.load([m, k]); + let b: Tile = rhs_part.load([k, n]); + let zero: Tile = constant(0.0f32, shape![16, 16]); + let partial: Tile = mma(a, b, zero); + + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([m, n], partial); + } +} + +use atomic_value_kernels::{ + accumulate_bf16_ones, accumulate_f16_ones, accumulate_ones, split_k_matmul, +}; + +/// Lost updates are the failure this capability exists to prevent, so the +/// split count is high enough that a racing read-modify-write loses some. +const SPLITS: u32 = 64; + +fn assert_atomic_add_ones(mut out: Tensor, grid: (u32, u32, u32), expected: T, launch: L) +where + T: DType + PartialEq + std::fmt::Debug + Copy, + L: FnOnce(Atomic<&mut Tensor>, (u32, u32, u32)), +{ + launch(out.atomic(), grid); + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_uniform(&host, expected, "atomic add"); +} + +#[test] +fn f16_atomic_add_accumulates_every_contribution() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let tiles = (2u32, 3u32); + let out = api::zeros::(&[tiles.0 as usize * BM, tiles.1 as usize * BN]) + .sync() + .expect("zeros"); + assert_atomic_add_ones( + out, + (tiles.0, tiles.1, SPLITS), + f16::from_f32(SPLITS as f32), + |out, grid| { + let _ = accumulate_f16_ones(out) + .grid(grid) + .sync() + .expect("f16 atomic add"); + }, + ); + }); +} + +#[test] +fn bf16_atomic_add_accumulates_every_contribution() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[ + cutile_ir::requirements::Feature::AtomicRedView, + cutile_ir::requirements::Feature::Bf16AtomicAdd, + ]) { + return; + } + let tiles = (2u32, 3u32); + let out = api::zeros::(&[tiles.0 as usize * BM, tiles.1 as usize * BN]) + .sync() + .expect("zeros"); + assert_atomic_add_ones( + out, + (tiles.0, tiles.1, SPLITS), + bf16::from_f32(SPLITS as f32), + |out, grid| { + let _ = accumulate_bf16_ones(out) + .grid(grid) + .sync() + .expect("bf16 atomic add"); + }, + ); + }); +} + +#[test] +fn every_contribution_lands() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let tiles_m = 2u32; + let tiles_n = 3u32; + let rows = tiles_m as usize * BM; + let cols = tiles_n as usize * BN; + + // `atomic_add` accumulates; the identity has to be there first. + let mut out = api::zeros::(&[rows, cols]).sync().expect("zeros"); + // The borrowed binding writes in place; dropping the recovered + // `&mut` here ends the borrow so `out` can be read below. + let _ = accumulate_ones(out.atomic()) + .grid((tiles_m, tiles_n, SPLITS)) + .sync() + .expect("launch"); + + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host.len(), rows * cols); + let expected = SPLITS as f32; + // Integers up to 2^24 are exact in f32, so this is an equality, not a + // tolerance: any deviation is a lost or duplicated update. + let wrong = host.iter().filter(|&&v| v != expected).count(); + assert_eq!( + wrong, + 0, + "{wrong} of {} elements are not {expected}; first few: {:?}", + host.len(), + &host[..8.min(host.len())] + ); + }); +} + +#[test] +fn split_k_matmul_matches_a_host_oracle() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let tiles_m = 2usize; + let tiles_n = 2usize; + let splits = 4usize; + let (m, n, k) = (tiles_m * BM, tiles_n * BN, splits * BK); + + // Small signed values keep every partial product exact in f32, so the + // only inexactness left is the order the atomics add them in. + let lhs_host: Vec = (0..m * k).map(|i| ((i % 7) as f32) - 3.0).collect(); + let rhs_host: Vec = (0..k * n).map(|i| ((i % 5) as f32) - 2.0).collect(); + + let mut expected = vec![0.0f32; m * n]; + for row in 0..m { + for col in 0..n { + let mut acc = 0.0f32; + for step in 0..k { + acc += lhs_host[row * k + step] * rhs_host[step * n + col]; + } + expected[row * n + col] = acc; + } + } + + let lhs: Arc> = api::copy_host_vec_to_device(&Arc::new(lhs_host)) + .reshape(&[m, k]) + .sync() + .expect("lhs") + .into(); + let rhs: Arc> = api::copy_host_vec_to_device(&Arc::new(rhs_host)) + .reshape(&[k, n]) + .sync() + .expect("rhs") + .into(); + + let mut out = api::zeros::(&[m, n]).sync().expect("zeros"); + let _ = split_k_matmul(out.atomic(), lhs, rhs) + .grid((tiles_m as u32, tiles_n as u32, splits as u32)) + .sync() + .expect("launch"); + + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host.len(), m * n); + + // Atomic float addition has no fixed summation order, so this is a + // tolerance and not an equality — and deliberately so: the capability + // does not promise bit-reproducibility. A lost contribution is a whole + // tile's worth of error and cannot hide inside this bound. + let mut worst = 0.0f32; + for (got, want) in host.iter().zip(expected.iter()) { + let scale = want.abs().max(1.0); + worst = worst.max((got - want).abs() / scale); + } + assert!( + worst < 1e-5, + "worst relative error {worst} exceeds tolerance; \ + got {:?} want {:?}", + &host[..4.min(host.len())], + &expected[..4.min(expected.len())] + ); + }); +} + +// ── The owned binding ─────────────────────────────────────────────────────── + +/// The owned binding, end to end on a device. +/// +/// `Atomic::owned` is the other exclusive host authority: ownership rather +/// than an exclusive borrow. Until now it was only typechecked +/// (`tests/aggregation_tensor.rs::owned_host_binding_typechecks`), so nothing +/// exercised preparation, launch, and recovery through it, and nothing checked +/// that the tensor the launch hands back holds the accumulated values. The +/// launch *takes* the tensor, so recovery is the only way back to the data — +/// which makes the recovered `Tensor` the thing under test, not a +/// convenience. +#[test] +fn an_owned_binding_accumulates_and_recovers_the_tensor() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let tiles_m = 2u32; + let tiles_n = 3u32; + let rows = tiles_m as usize * BM; + let cols = tiles_n as usize * BN; + + let out = api::zeros::(&[rows, cols]).sync().expect("zeros"); + // A one-parameter kernel recovers a one-tuple; the annotation pins + // that the owned binding comes back as the plain `Tensor` it was, + // not as the `Atomic<..>` wrapper. + let (recovered,): (Tensor,) = accumulate_ones(Atomic::owned(out)) + .grid((tiles_m, tiles_n, SPLITS)) + .sync() + .expect("launch"); + + let host: Vec = recovered.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host.len(), rows * cols); + let expected = SPLITS as f32; + // Integers up to 2^24 are exact in f32, so this is an equality, not a + // tolerance: any deviation is a lost or duplicated update. + let wrong = host.iter().filter(|&&v| v != expected).count(); + assert_eq!( + wrong, + 0, + "{wrong} of {} elements are not {expected}; first few: {:?}", + host.len(), + &host[..8.min(host.len())] + ); + }); +} + +// ── One allocation, two bindings ──────────────────────────────────────────── + +#[cutile::module] +mod alias_kernels { + use cutile::core::*; + + /// One atomic destination and one ordinary input. Binding both to the + /// same allocation is the race the host must reject: the two parameters + /// are independent views, so no token ordering inside the kernel fixes + /// it. + /// + /// The contribution is `other + 1`, which is nonzero whatever `other` + /// holds, so a launch that reached the device leaves a mark in `out`. + #[cutile::entry()] + fn atomic_and_ordinary(out: &AddTensor, other: &Tensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let part: Partition = other.partition(shape![16, 16]); + let seen: Tile = part.load([bid.0, bid.1]); + let ones: Tile = constant(1.0f32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([bid.0, bid.1], seen + ones); + } + + /// Two atomic destinations. Binding both to one allocation is exactly + /// what the capability is for, because overlapping atomic updates + /// compose, so the alias check must leave it alone. + #[cutile::entry()] + fn two_atomic_destinations( + first: &AddTensor, + second: &AddTensor, + ) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let one: Tile = constant(1.0f32, shape![16, 16]); + let two: Tile = constant(2.0f32, shape![16, 16]); + let a: AddPartition = first.partition(shape![16, 16]); + let b: AddPartition = second.partition(shape![16, 16]); + let _x: Token = a.atomic_add([bid.0, bid.1], one); + let _y: Token = b.atomic_add([bid.0, bid.1], two); + } +} + +use alias_kernels::{atomic_and_ordinary, two_atomic_destinations}; + +/// One allocation bound both atomically and ordinarily is rejected, and +/// rejected before anything runs. +/// +/// "Before any GPU work" is asserted, not inferred. Two independent facts +/// carry it: +/// +/// 1. The error is *that* rejection — it names both parameters and says one +/// is bound atomically — so a launch that failed for some other reason +/// cannot pass this test. +/// 2. The destination is unchanged. It was zeroed, and every program of this +/// kernel adds a strictly positive tile to its output tile, so any program +/// that reached the device would have left `out` non-zero. +/// +/// Neither fact depends on a valid kernel executing afterwards: the test +/// still fails if nothing on this machine could ever launch. +#[test] +fn one_allocation_bound_atomically_and_ordinarily_is_rejected_before_any_gpu_work() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let tiles = 2u32; + let rows = tiles as usize * BM; + let cols = tiles as usize * BN; + let mut out = api::zeros::(&[rows, cols]).sync().expect("zeros"); + + // Deliberate aliasing: one allocation, two bindings. `into_shared_alias` + // is `unsafe` for exactly the reason under test, and the launch is + // supposed to be what catches it. + let alias: Arc> = unsafe { out.into_shared_alias() }; + assert_eq!( + alias.storage_id(), + out.storage_id(), + "the alias must name the same allocation, or this test proves nothing" + ); + + let result = atomic_and_ordinary(out.atomic(), alias) + .grid((tiles, tiles, SPLITS)) + .sync(); + let Err(error) = result else { + panic!("binding one allocation atomically and ordinarily must fail the launch"); + }; + let message = error.to_string(); + for fragment in ["`out`", "`other`", "same allocation", "is bound atomically"] { + assert!( + message.contains(fragment), + "the rejection must be the alias check and must explain itself; \ + missing {fragment:?} in: {message}" + ); + } + + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host.len(), rows * cols); + let touched = host.iter().filter(|&&v| v != 0.0).count(); + assert_eq!( + touched, + 0, + "{touched} elements changed, so the launch reached the device before \ + being rejected; first few: {:?}", + &host[..8.min(host.len())] + ); + }); +} + +/// Two atomic bindings over one allocation stay accepted — the alias check +/// rejects a mixed pair, not aliasing itself. +/// +/// Both programs' contributions land in the same storage, so each element +/// ends at `1.0 + 2.0`. +#[test] +fn two_atomic_bindings_over_one_allocation_are_accepted() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let tiles = 2u32; + let rows = tiles as usize * BM; + let cols = tiles as usize * BN; + let mut out = api::zeros::(&[rows, cols]).sync().expect("zeros"); + + // A second *owned* tensor over the same storage, which is what the + // owned atomic binding needs. A freshly minted alias is the only + // reference to its own `Arc`, so unwrapping it always succeeds. + let Ok(second) = Arc::try_unwrap(unsafe { out.into_shared_alias() }) else { + panic!("a freshly minted alias is the only reference to itself"); + }; + assert_eq!( + second.storage_id(), + out.storage_id(), + "both bindings must name the same allocation" + ); + + // The recovered bindings are dropped here, which ends the borrow of + // `out` so it can be read below. + let _ = two_atomic_destinations(out.atomic(), Atomic::owned(second)) + .grid((tiles, tiles, 1)) + .sync() + .expect("two atomic bindings over one allocation must be accepted"); + + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_eq!(host.len(), rows * cols); + let wrong = host.iter().filter(|&&v| v != 3.0).count(); + assert_eq!( + wrong, + 0, + "{wrong} of {} elements are not 3.0; first few: {:?}", + host.len(), + &host[..8.min(host.len())] + ); + }); +} + +#[cutile::module] +mod per_operation_kernels { + use cutile::core::*; + + /// OR is idempotent, so any number of overlapping programs must leave + /// exactly `seed | 0x00A5_0F3C`. + #[cutile::entry()] + fn or_pattern(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x00A5_0F3Cu32, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + + /// AND is idempotent, so starting from all-ones the result must be + /// exactly the mask. + #[cutile::entry()] + fn and_mask(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_F0F0u32, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + + /// Every program XORs the same value; an even number of programs must + /// restore the seed exactly, so a lost or duplicated update shows. + #[cutile::entry()] + fn xor_toggle(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0xA5A5_A5A5u32, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + + /// Unsigned maximum where contribution and seed straddle the high bit: + /// the seed is `0x7FFF_FFF0` and the contribution `0x8000_0010`. Under + /// the correct unsigned comparison the high-bit value wins; a signed + /// comparison would read it as negative and keep the seed, a visibly + /// wrong result. + #[cutile::entry()] + fn max_straddles_the_high_bit(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x8000_0010u32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + + /// Unsigned minimum across the same straddle, reversed: the seed is + /// `0x8000_0010` and the contribution `0x7FFF_FFF0`, which must win; a + /// signed comparison would keep the "negative" seed. + #[cutile::entry()] + fn min_straddles_the_high_bit(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x7FFF_FFF0u32, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } + + /// Signed maximum of seed `2` and contribution `-5`: the signed answer + /// keeps `2`; an unsigned comparison would read `-5` as `0xFFFF_FFFB` + /// and take it. + #[cutile::entry()] + fn max_signed_negatives(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(-5i32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + + /// Integer addition is exact at any count, unlike the float split-K + /// above, so the per-element count is an equality. + #[cutile::entry()] + fn add_u32_ones(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } +} + +/// Seeds a `(2 tiles, 3 tiles)` u32 destination with `seed`, launches over a +/// `(2, 3, splits)` grid, and returns the device result. +fn run_u32(seed: u32, splits: u32, launch: L) -> Vec +where + L: FnOnce(Atomic<&mut Tensor>, (u32, u32, u32)), +{ + let (tiles_m, tiles_n) = (2u32, 3u32); + let rows = tiles_m as usize * BM; + let cols = tiles_n as usize * BN; + let mut out = api::full::(seed, &[rows, cols]).sync().expect("seed"); + launch(out.atomic(), (tiles_m, tiles_n, splits)); + out.dup().to_host_vec().sync().expect("to_host") +} + +/// Every element of `host` must equal `expected`; reports how many differ +/// and shows the first few, so a wrong comparison or lost update is legible. +fn assert_uniform(host: &[T], expected: T, what: &str) { + let wrong = host.iter().filter(|&&v| v != expected).count(); + assert_eq!( + wrong, + 0, + "{what}: {wrong} of {} elements are not {expected:?}; first few: {:?}", + host.len(), + &host[..8.min(host.len())] + ); +} + +/// OR, AND and XOR against bit-exact references: idempotent OR and AND land +/// exactly on `seed | v` and `seed & v`, and an even XOR toggle count +/// restores the seed, so any lost or duplicated update is visible. +#[test] +fn bitwise_operations_match_bit_exact_references() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let or = run_u32(0x5000_0001, 4, |out, grid| { + let _ = per_operation_kernels::or_pattern(out) + .grid(grid) + .sync() + .expect("or"); + }); + assert_uniform(&or, 0x5000_0001 | 0x00A5_0F3C, "atomic_or"); + + let and = run_u32(!0u32, 4, |out, grid| { + let _ = per_operation_kernels::and_mask(out) + .grid(grid) + .sync() + .expect("and"); + }); + assert_uniform(&and, 0x0F0F_F0F0, "atomic_and"); + + let xor = run_u32(0x5A5A_5A5A, 8, |out, grid| { + let _ = per_operation_kernels::xor_toggle(out) + .grid(grid) + .sync() + .expect("xor"); + }); + assert_uniform(&xor, 0x5A5A_5A5A, "atomic_xor: 8 toggles must cancel"); + }); +} + +/// The comparisons respect the element's signedness on real hardware, with +/// seed and contribution straddling the high bit so the wrong comparison +/// gives a visibly wrong value rather than a coincidence. +#[test] +fn comparisons_respect_signedness_on_device() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let max = run_u32(0x7FFF_FFF0, 1, |out, grid| { + let _ = per_operation_kernels::max_straddles_the_high_bit(out) + .grid(grid) + .sync() + .expect("max u32"); + }); + assert_uniform(&max, 0x8000_0010u32, "atomic_max u32: unsigned comparison"); + + let min = run_u32(0x8000_0010, 1, |out, grid| { + let _ = per_operation_kernels::min_straddles_the_high_bit(out) + .grid(grid) + .sync() + .expect("min u32"); + }); + assert_uniform(&min, 0x7FFF_FFF0u32, "atomic_min u32: unsigned comparison"); + + let (tiles_m, tiles_n) = (2u32, 3u32); + let rows = tiles_m as usize * BM; + let cols = tiles_n as usize * BN; + let mut out = api::full::(2, &[rows, cols]).sync().expect("seed"); + let _ = per_operation_kernels::max_signed_negatives(out.atomic()) + .grid((tiles_m, tiles_n, 1)) + .sync() + .expect("max i32"); + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + assert_uniform(&host, 2i32, "atomic_max i32: signed comparison"); + }); +} + +/// Integer addition is exact, so the count is an equality at any size. +#[test] +fn integer_addition_is_exact_on_device() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let splits = 64u32; + let host = run_u32(0, splits, |out, grid| { + let _ = per_operation_kernels::add_u32_ones(out) + .grid(grid) + .sync() + .expect("add"); + }); + assert_uniform(&host, splits, "atomic_add u32"); + }); +} + +#[cutile::module] +mod mixed_operation_kernels { + use cutile::core::*; + + /// Two different aggregation operations over what the host binds as one + /// allocation: the launcher must reject the launch, so this body is + /// never reached on any device. + #[cutile::entry()] + fn add_and_max(a: &AddTensor, b: &MaxTensor) { + let v: Tile = constant(1i32, shape![16, 16]); + let add: AddPartition = a.partition(shape![16, 16]); + let max: MaxPartition = b.partition(shape![16, 16]); + let _t: Token = add.atomic_add([0i32, 0i32], v); + let _t: Token = max.atomic_max([0i32, 0i32], v); + } +} + +/// One allocation bound to two *different* aggregation operations is +/// rejected before any GPU work, naming both parameters and both operations. +#[test] +fn mixed_operations_on_one_allocation_are_rejected_before_any_gpu_work() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let mut out = api::zeros::(&[32, 48]).sync().expect("zeros"); + // A second owned tensor over the same storage, as in the + // two-atomic-bindings test above. + let Ok(second) = Arc::try_unwrap(unsafe { out.into_shared_alias() }) else { + panic!("a freshly minted alias is the only reference to itself"); + }; + let err = mixed_operation_kernels::add_and_max(out.atomic(), Atomic::owned(second)) + .grid((2, 3, 1)) + .sync() + .expect_err("two operations over one allocation must not launch"); + let message = err.to_string(); + assert!( + message.contains("addition") && message.contains("maximum"), + "the rejection must name both operations: {message}" + ); + }); +} + +#[cutile::module] +mod matrix_kernels { + use cutile::core::*; + + #[cutile::entry()] + fn matrix_add_i32(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_add_i64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1i64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_add_u64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1u64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_add_f64(out: &AddTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(1.0f64, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_add([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_or_i32(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0Fi32, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_and_i32(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0Fi32, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_xor_i32(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0Fi32, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_or_i64(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0F_0F0F_0F0Fi64, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_and_i64(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0F_0F0F_0F0Fi64, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_xor_i64(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0F_0F0F_0F0Fi64, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_or_u64(out: &OrTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0F_0F0F_0F0Fu64, shape![16, 16]); + let acc: OrPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_or([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_and_u64(out: &AndTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0F_0F0F_0F0Fu64, shape![16, 16]); + let acc: AndPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_and([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_xor_u64(out: &XorTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x0F0F_0F0F_0F0F_0F0Fu64, shape![16, 16]); + let acc: XorPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_xor([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_max_i64(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(-5i64, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_max_u64(out: &MaxTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x8000_0000_0000_0010u64, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_max([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_min_i32(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(-5i32, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_min_i64(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(-5i64, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } + #[cutile::entry()] + fn matrix_min_u64(out: &MinTensor) { + let bid: (i32, i32, i32) = get_tile_block_id(); + let v: Tile = constant(0x7FFF_FFFF_FFFF_FFF0u64, shape![16, 16]); + let acc: MinPartition = out.partition(shape![16, 16]); + let _t: Token = acc.atomic_min([bid.0, bid.1], v); + } +} + +/// One device value test per remaining operation/element cell, completing +/// hardware coverage of the legal matrix (f16/bf16/f32/u32/i32 cells are +/// covered above). Exact references throughout: counting for addition, +/// idempotent masks for OR and AND, an even toggle count for XOR, and +/// seed/contribution pairs that distinguish the comparison's signedness, +/// including 64-bit values straddling the high bit. +#[test] +fn the_remaining_matrix_cells_hold_on_device() { + common::with_test_stack(|| { + if !common::supports_tile_ir(&[cutile_ir::requirements::Feature::AtomicRedView]) { + return; + } + let (tiles_m, tiles_n) = (2u32, 2u32); + let shape = [tiles_m as usize * BM, tiles_n as usize * BN]; + let grid = (tiles_m, tiles_n, 1u32); + let grid8 = (tiles_m, tiles_n, 8u32); + + fn check( + what: &str, + out: &mut Tensor, + expected: T, + ) { + let host: Vec = out.dup().to_host_vec().sync().expect("to_host"); + let wrong = host.iter().filter(|&&v| v != expected).count(); + assert_eq!(wrong, 0, "{what}: {wrong} of {} wrong", host.len()); + } + + // Addition counts exactly at 8 overlapping programs per tile. + let mut o = api::zeros::(&shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_add_i32(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("add i32", &mut o, 8i32); + let mut o = api::zeros::(&shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_add_i64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("add i64", &mut o, 8i64); + let mut o = api::zeros::(&shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_add_u64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("add u64", &mut o, 8u64); + let mut o = api::zeros::(&shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_add_f64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("add f64", &mut o, 8.0f64); + + // OR and AND are idempotent; XOR's even toggle count restores the seed. + let mut o = api::full::(0x5000_0001, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_or_i32(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("or i32", &mut o, 0x5000_0001 | 0x0F0F_0F0F); + let mut o = api::full::(-1, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_and_i32(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("and i32", &mut o, 0x0F0F_0F0F); + let mut o = api::full::(0x5A5A_5A5A, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_xor_i32(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("xor i32 (8 toggles cancel)", &mut o, 0x5A5A_5A5A); + + let mut o = api::full::(0x5000_0000_0000_0001, &shape) + .sync() + .expect("seed"); + let _ = matrix_kernels::matrix_or_i64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check( + "or i64", + &mut o, + 0x5000_0000_0000_0001 | 0x0F0F_0F0F_0F0F_0F0F, + ); + let mut o = api::full::(-1, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_and_i64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("and i64", &mut o, 0x0F0F_0F0F_0F0F_0F0F); + let mut o = api::full::(0x5A5A_5A5A_5A5A_5A5A, &shape) + .sync() + .expect("seed"); + let _ = matrix_kernels::matrix_xor_i64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("xor i64", &mut o, 0x5A5A_5A5A_5A5A_5A5A); + + let mut o = api::full::(0x9000_0000_0000_0001, &shape) + .sync() + .expect("seed"); + let _ = matrix_kernels::matrix_or_u64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check( + "or u64", + &mut o, + 0x9000_0000_0000_0001 | 0x0F0F_0F0F_0F0F_0F0F, + ); + let mut o = api::full::(!0u64, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_and_u64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("and u64", &mut o, 0x0F0F_0F0F_0F0F_0F0F); + let mut o = api::full::(0xA5A5_A5A5_A5A5_A5A5, &shape) + .sync() + .expect("seed"); + let _ = matrix_kernels::matrix_xor_u64(o.atomic()) + .grid(grid8) + .sync() + .expect("launch"); + check("xor u64", &mut o, 0xA5A5_A5A5_A5A5_A5A5); + + // Comparisons: signed keeps 2 over -5; unsigned 64-bit straddles the + // high bit in both directions. + let mut o = api::full::(2, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_max_i64(o.atomic()) + .grid(grid) + .sync() + .expect("launch"); + check("max i64 (signed)", &mut o, 2i64); + let mut o = api::full::(0x7FFF_FFFF_FFFF_FFF0, &shape) + .sync() + .expect("seed"); + let _ = matrix_kernels::matrix_max_u64(o.atomic()) + .grid(grid) + .sync() + .expect("launch"); + check("max u64 (unsigned)", &mut o, 0x8000_0000_0000_0010u64); + let mut o = api::full::(2, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_min_i32(o.atomic()) + .grid(grid) + .sync() + .expect("launch"); + check("min i32 (signed)", &mut o, -5i32); + let mut o = api::full::(2, &shape).sync().expect("seed"); + let _ = matrix_kernels::matrix_min_i64(o.atomic()) + .grid(grid) + .sync() + .expect("launch"); + check("min i64 (signed)", &mut o, -5i64); + let mut o = api::full::(0x8000_0000_0000_0010, &shape) + .sync() + .expect("seed"); + let _ = matrix_kernels::matrix_min_u64(o.atomic()) + .grid(grid) + .sync() + .expect("launch"); + check("min u64 (unsigned)", &mut o, 0x7FFF_FFFF_FFFF_FFF0u64); + }); +} diff --git a/cutile-rs/cutile/tests/ui.rs b/cutile-rs/cutile/tests/ui.rs index 6876a80a9a..394b9f0bfc 100644 --- a/cutile-rs/cutile/tests/ui.rs +++ b/cutile-rs/cutile/tests/ui.rs @@ -27,4 +27,9 @@ fn ui() { t.compile_fail("tests/ui/entry_non_literal_value.rs"); t.compile_fail("tests/ui/global_*.rs"); t.compile_fail("tests/ui/unordered_partition_requires_unsafe.rs"); + // The four `tests/ui/aggregation_tensor_*.rs` cases are registered in + // `tests/ui_aggregation.rs`, not here: that target is in CPU CI and this + // one cannot be (see that file's header). Registering them in both would + // build each case twice and give their expectations two places to be + // re-blessed from. } diff --git a/cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.rs b/cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.rs new file mode 100644 index 0000000000..786f87f48e --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.rs @@ -0,0 +1,12 @@ +// `KernelAtomic` is sealed. A downstream wrapper that delegated to +// `Arc>`'s launch plumbing would retain the storage as a *read*, +// bypassing the exclusive host authority an aggregation binding promises, so +// no type outside the crate may implement the trait. +use cutile::prelude::*; +use std::sync::Arc; + +struct SharedWrapper(Arc>); + +impl cutile::tensor::KernelAtomic for SharedWrapper {} + +fn main() {} diff --git a/cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.stderr b/cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.stderr new file mode 100644 index 0000000000..ab3d9d9cc7 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_binding_cannot_be_reimplemented.stderr @@ -0,0 +1,51 @@ +error[E0277]: the trait bound `SharedWrapper: tensor::atomic_sealed::Sealed` is not satisfied + --> tests/ui/aggregation_binding_cannot_be_reimplemented.rs:10:44 + | +10 | impl cutile::tensor::KernelAtomic for SharedWrapper {} + | ^^^^^^^^^^^^^ unsatisfied trait bound + | +help: the trait `tensor::atomic_sealed::Sealed` is not implemented for `SharedWrapper` + --> tests/ui/aggregation_binding_cannot_be_reimplemented.rs:8:1 + | + 8 | struct SharedWrapper(Arc>); + | ^^^^^^^^^^^^^^^^^^^^ +help: the following other types implement trait `tensor::atomic_sealed::Sealed` + --> src/tensor.rs + | + | impl Sealed for super::Atomic> {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `cutile::prelude::Atomic>` + | impl Sealed for super::Atomic<&mut super::Tensor> {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `cutile::prelude::Atomic<&mut cutile::prelude::Tensor>` +note: required by a bound in `KernelAtomic` + --> src/tensor.rs + | + | pub trait KernelAtomic: KernelInput + atomic_sealed::Sealed {} + | ^^^^^^^^^^^^^^^^^^^^^ required by this bound in `KernelAtomic` + = note: `KernelAtomic` is a "sealed trait", because to implement it you also need to implement `cutile::tensor::atomic_sealed::Sealed`, which is not accessible; this is usually done to force you to use one of the provided types that already implement it + = help: the following types implement the trait: + cutile::prelude::Atomic> + cutile::prelude::Atomic<&mut cutile::prelude::Tensor> + +error[E0277]: the trait bound `SharedWrapper: KernelInput` is not satisfied + --> tests/ui/aggregation_binding_cannot_be_reimplemented.rs:10:44 + | +10 | impl cutile::tensor::KernelAtomic for SharedWrapper {} + | ^^^^^^^^^^^^^ unsatisfied trait bound + | +help: the trait `KernelInput` is not implemented for `SharedWrapper` + --> tests/ui/aggregation_binding_cannot_be_reimplemented.rs:8:1 + | + 8 | struct SharedWrapper(Arc>); + | ^^^^^^^^^^^^^^^^^^^^ + = help: the following other types implement trait `KernelInput`: + &'a cutile::prelude::Tensor + &'a cutile::prelude::TensorView<'a, T> + Arc> + cutile::prelude::Atomic<&'a mut cutile::prelude::Tensor> + cutile::prelude::Atomic> + cutile::prelude::Tensor +note: required by a bound in `KernelAtomic` + --> src/tensor.rs + | + | pub trait KernelAtomic: KernelInput + atomic_sealed::Sealed {} + | ^^^^^^^^^^^^^^ required by this bound in `KernelAtomic` diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.rs b/cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.rs new file mode 100644 index 0000000000..a74cb5445c --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.rs @@ -0,0 +1,24 @@ +// Binding an atomic destination requires exclusive mutable authority over the +// backing tensor for the launch. A shared `Arc>` is not that: another +// holder could read the storage while the atomics are in flight. +use cutile::prelude::*; + +#[cutile::module] +mod kernels { + use cutile::core::*; + + #[cutile::entry()] + fn accumulate(out: &AddTensor, src: &Tensor) { + let part: Partition = src.partition(shape![16, 16]); + let bid: (i32, i32, i32) = get_tile_block_id(); + let tile: Tile = part.load([bid.0, bid.1]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([bid.0, bid.1], tile); + } +} + +fn shared(out: Arc>, src: Arc>) { + let _ = kernels::accumulate(out, src); +} + +fn main() {} diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.stderr b/cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.stderr new file mode 100644 index 0000000000..d513c16389 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_needs_exclusive_host_authority.stderr @@ -0,0 +1,20 @@ +error[E0277]: the trait bound `std::sync::Arc>: KernelAtomic` is not satisfied + --> tests/ui/aggregation_tensor_needs_exclusive_host_authority.rs:21:13 + | +21 | let _ = kernels::accumulate(out, src); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `KernelAtomic` is not implemented for `std::sync::Arc>` + | +help: the following other types implement trait `KernelAtomic` + --> src/tensor.rs + | + | impl KernelAtomic for Atomic> {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `cutile::prelude::Atomic>` +... + | impl KernelAtomic for Atomic<&mut Tensor> {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `cutile::prelude::Atomic<&mut cutile::prelude::Tensor>` +note: required by a bound in `accumulate` + --> tests/ui/aggregation_tensor_needs_exclusive_host_authority.rs:6:1 + | + 6 | #[cutile::module] + | ^^^^^^^^^^^^^^^^^ required by this bound in `accumulate` + = note: this error originates in the attribute macro `cutile::module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.rs b/cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.rs new file mode 100644 index 0000000000..f089950a2b --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.rs @@ -0,0 +1,16 @@ +// The capability is shared by construction: every program in the launch holds +// it and the atomics resolve the collisions. `&mut` would promise the +// disjointness the operation deliberately does not have. +#[cutile::module] +mod kernels { + use cutile::core::*; + + #[cutile::entry()] + fn accumulate(out: &mut AddTensor) { + let ones: Tile = constant(1.0f32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([0i32, 0i32], ones); + } +} + +fn main() {} diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.stderr b/cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.stderr new file mode 100644 index 0000000000..ce2d20f284 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_rejects_mutable_reference.stderr @@ -0,0 +1,5 @@ +error: `&mut AddTensor<..>` is not a kernel parameter type. Atomic updates go through a shared `&AddTensor<..>`, like Rust's own atomics. + --> tests/ui/aggregation_tensor_rejects_mutable_reference.rs:9:24 + | +9 | fn accumulate(out: &mut AddTensor) { + | ^ diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.rs b/cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.rs new file mode 100644 index 0000000000..188463f4c5 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.rs @@ -0,0 +1,18 @@ +// `atomic::AtomicAdd` admits the element types Tile IR can add atomically: +// `f16`, `bf16`, `f32`, `f64` through the `AddF` mode, and `i32`, `u32`, +// `i64`, `u64` through `Add`. `i8` is an `ElementType` the IR cannot reduce +// atomically at all, so the pairing must fail at the type level rather than +// reach the JIT. +#[cutile::module] +mod kernels { + use cutile::core::*; + + #[cutile::entry()] + fn accumulate(out: &AddTensor) { + let ones: Tile = constant(1i8, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_add([0i32, 0i32], ones); + } +} + +fn main() {} diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.stderr b/cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.stderr new file mode 100644 index 0000000000..34bd9439be --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_requires_a_supported_element.stderr @@ -0,0 +1,77 @@ +error[E0277]: the trait bound `i8: AtomicAdd` is not satisfied + --> tests/ui/aggregation_tensor_requires_a_supported_element.rs:6:1 + | +6 | #[cutile::module] + | ^^^^^^^^^^^^^^^^^ the trait `AtomicAdd` is not implemented for `i8` + | + = help: the following other types implement trait `AtomicAdd`: + f32 + f64 + i32 + i64 + u32 + u64 +note: required by a bound in `cutile::core::AddTensor_2` + --> src/_core.rs + | + | pub struct AddTensor { + | ^^^^^^^^^^^^^^^^^ required by this bound in `AddTensor_2` + = note: this error originates in the attribute macro `cutile::module` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: the method `partition` exists for reference `&cutile::core::AddTensor_2`, but its trait bounds were not satisfied + --> tests/ui/aggregation_tensor_requires_a_supported_element.rs:13:55 + | +13 | let acc: AddPartition = out.partition(shape![16, 16]); + | ^^^^^^^^^ method cannot be called on `&cutile::core::AddTensor_2` due to unsatisfied trait bounds + | + ::: src/_core.rs + | + | pub struct AddTensor { + | ------------------------------------------------------------- doesn't satisfy `<_ as DeviceOp>::Output = _`, `_: DeviceOp`, `_: PartitionOp<_, AddTensor_2>` or `cutile::core::AddTensor_2: Iterator` + | + = note: the following trait bounds were not satisfied: + `i8: AtomicAdd` + `<&cutile::core::AddTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &cutile::core::AddTensor_2>` + `&cutile::core::AddTensor_2: cutile::prelude::DeviceOp` + which is required by `&cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &cutile::core::AddTensor_2>` + `<&&cutile::core::AddTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&&cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &&cutile::core::AddTensor_2>` + `&&cutile::core::AddTensor_2: cutile::prelude::DeviceOp` + which is required by `&&cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &&cutile::core::AddTensor_2>` + `<&mut &cutile::core::AddTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&mut &cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &mut &cutile::core::AddTensor_2>` + `&mut &cutile::core::AddTensor_2: cutile::prelude::DeviceOp` + which is required by `&mut &cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &mut &cutile::core::AddTensor_2>` + `&cutile::core::AddTensor_2: Iterator` + which is required by `&mut &cutile::core::AddTensor_2: Iterator` + ` as cutile::prelude::DeviceOp>::Output = _` + which is required by `cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, cutile::core::AddTensor_2>` + `cutile::core::AddTensor_2: cutile::prelude::DeviceOp` + which is required by `cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, cutile::core::AddTensor_2>` + `<&mut cutile::core::AddTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&mut cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &mut cutile::core::AddTensor_2>` + `&mut cutile::core::AddTensor_2: cutile::prelude::DeviceOp` + which is required by `&mut cutile::core::AddTensor_2: cutile::prelude::PartitionOp<_, &mut cutile::core::AddTensor_2>` + `cutile::core::AddTensor_2: Iterator` + which is required by `&mut cutile::core::AddTensor_2: Iterator` + +error[E0277]: the trait bound `i8: AtomicAdd` is not satisfied + --> tests/ui/aggregation_tensor_requires_a_supported_element.rs:6:1 + | +6 | #[cutile::module] + | ^^^^^^^^^^^^^^^^^ the trait `AtomicAdd` is not implemented for `i8` + | + = help: the following other types implement trait `AtomicAdd`: + f32 + f64 + i32 + i64 + u32 + u64 +note: required by a bound in `cutile::core::AddPartition_2` + --> src/_core.rs + | + | pub struct AddPartition<'a, E: atomic::AtomicAdd, const D: [i32; N]> { + | ^^^^^^^^^^^^^^^^^ required by this bound in `AddPartition_2` + = note: this error originates in the attribute macro `cutile::module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.rs b/cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.rs new file mode 100644 index 0000000000..0c2c0307b3 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.rs @@ -0,0 +1,33 @@ +// The aggregation capability exposes its reductions and nothing else. An +// ordinary load or store through the same storage would race the other +// programs' reductions, so neither the tensor nor the view it hands out +// offers one — and `partition` yields an `AddPartition`, never an +// ordinary `Partition` or `PartitionMut`. +#[cutile::module] +mod kernels { + use cutile::core::*; + + #[cutile::entry()] + fn accumulate(out: &AddTensor) { + let zero: Tile = constant(0.0f32, shape![16, 16]); + + // Nothing ordinary on the capability itself. `shape()` is internal + // too: it exists only because the view constructor's + // `output_type_meta` evaluates it, and tile counts come from + // `num_tiles(&partition, axis)` instead. + let _shape: Shape<{ [-1, -1] }> = out.shape(); + let _loaded: Tile = out.load([0i32, 0i32]); + out.store(zero); + let _mut_view: PartitionMut = out.partition_mut(shape![16, 16]); + + // `partition` exists, but it cannot produce an ordinary view. + let _ordinary: Partition = out.partition(shape![16, 16]); + + // Nor does the view it does produce permit an ordinary access. + let acc: AddPartition = out.partition(shape![16, 16]); + let _tile: Tile = acc.load([0i32, 0i32]); + let _tok: Token = acc.store(zero, [0i32, 0i32]); + } +} + +fn main() {} diff --git a/cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.stderr b/cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.stderr new file mode 100644 index 0000000000..f59761f6cc --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_tensor_withholds_ordinary_access.stderr @@ -0,0 +1,58 @@ +error[E0624]: method `shape` is private + --> tests/ui/aggregation_tensor_withholds_ordinary_access.rs:18:47 + | +18 | let _shape: Shape<{ [-1, -1] }> = out.shape(); + | ^^^^^ private method + | + ::: src/_core.rs + | + | #[cutile_macro::module(tile_rust_crate = true)] + | ----------------------------------------------- private method defined here + +error[E0599]: no method named `load` found for reference `&cutile::core::AddTensor_2` in the current scope + --> tests/ui/aggregation_tensor_withholds_ordinary_access.rs:19:52 + | +19 | let _loaded: Tile = out.load([0i32, 0i32]); + | ^^^^ method not found in `&cutile::core::AddTensor_2` + +error[E0599]: no method named `store` found for reference `&cutile::core::AddTensor_2` in the current scope + --> tests/ui/aggregation_tensor_withholds_ordinary_access.rs:20:13 + | +20 | out.store(zero); + | ^^^^^ method not found in `&cutile::core::AddTensor_2` + +error[E0599]: no method named `partition_mut` found for reference `&cutile::core::AddTensor_2` in the current scope + --> tests/ui/aggregation_tensor_withholds_ordinary_access.rs:21:62 + | +21 | let _mut_view: PartitionMut = out.partition_mut(shape![16, 16]); + | ^^^^^^^^^^^^^ + | +help: there is a method `partition` with a similar name + | +21 - let _mut_view: PartitionMut = out.partition_mut(shape![16, 16]); +21 + let _mut_view: PartitionMut = out.partition(shape![16, 16]); + | + +error[E0308]: mismatched types + --> tests/ui/aggregation_tensor_withholds_ordinary_access.rs:24:55 + | + 6 | #[cutile::module] + | ----------------- expected due to this +... +24 | let _ordinary: Partition = out.partition(shape![16, 16]); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `Partition_2<'_, f32, 16, 16>`, found `AddPartition_2<'_, f32, 16, 16>` + | + = note: expected struct `cutile::core::Partition_2<'_, f32, 16, 16>` + found struct `cutile::core::AddPartition_2<'_, f32, 16, 16>` + +error[E0599]: no method named `load` found for struct `cutile::core::AddPartition_2<'a, E, D0, D1>` in the current scope + --> tests/ui/aggregation_tensor_withholds_ordinary_access.rs:28:50 + | +28 | let _tile: Tile = acc.load([0i32, 0i32]); + | ^^^^ method not found in `cutile::core::AddPartition_2<'_, f32, 16, 16>` + +error[E0599]: no method named `store` found for struct `cutile::core::AddPartition_2<'a, E, D0, D1>` in the current scope + --> tests/ui/aggregation_tensor_withholds_ordinary_access.rs:29:31 + | +29 | let _tok: Token = acc.store(zero, [0i32, 0i32]); + | ^^^^^ method not found in `cutile::core::AddPartition_2<'_, f32, 16, 16>` diff --git a/cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.rs b/cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.rs new file mode 100644 index 0000000000..aa34ce6881 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.rs @@ -0,0 +1,17 @@ +// One capability is one operation. An `AddPartition` offers `atomic_add` and +// nothing else; combining by maximum into the same storage takes a +// `MaxTensor` binding, and the launcher rejects binding both to one +// allocation. The wrong-operation method must not exist. +#[cutile::module] +mod kernels { + use cutile::core::*; + + #[cutile::entry()] + fn accumulate(out: &AddTensor) { + let ones: Tile = constant(1i32, shape![16, 16]); + let acc: AddPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_max([0i32, 0i32], ones); + } +} + +fn main() {} diff --git a/cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.stderr b/cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.stderr new file mode 100644 index 0000000000..73b213adc2 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/aggregation_withholds_other_operations.stderr @@ -0,0 +1,11 @@ +error[E0599]: no method named `atomic_max` found for struct `cutile::core::AddPartition_2<'a, E, D0, D1>` in the current scope + --> tests/ui/aggregation_withholds_other_operations.rs:13:32 + | +13 | let _done: Token = acc.atomic_max([0i32, 0i32], ones); + | ^^^^^^^^^^ + | +help: there is a method `atomic_add` with a similar name + | +13 - let _done: Token = acc.atomic_max([0i32, 0i32], ones); +13 + let _done: Token = acc.atomic_add([0i32, 0i32], ones); + | diff --git a/cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.rs b/cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.rs new file mode 100644 index 0000000000..5943418a2d --- /dev/null +++ b/cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.rs @@ -0,0 +1,17 @@ +// `atomic::AtomicOrd` admits the element types Tile IR can compare +// atomically: `i32`, `u32`, `i64`, `u64`. There is no atomic floating-point +// maximum, so `MaxTensor` must fail at the type level rather than +// reach the JIT. +#[cutile::module] +mod kernels { + use cutile::core::*; + + #[cutile::entry()] + fn running_max(out: &MaxTensor) { + let ones: Tile = constant(1.0f32, shape![16, 16]); + let acc: MaxPartition = out.partition(shape![16, 16]); + let _done: Token = acc.atomic_max([0i32, 0i32], ones); + } +} + +fn main() {} diff --git a/cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.stderr b/cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.stderr new file mode 100644 index 0000000000..5c97365391 --- /dev/null +++ b/cutile-rs/cutile/tests/ui/max_tensor_requires_an_ordered_element.stderr @@ -0,0 +1,85 @@ +error[E0277]: the trait bound `f32: AtomicOrd` is not satisfied + --> tests/ui/max_tensor_requires_an_ordered_element.rs:5:1 + | +5 | #[cutile::module] + | ^^^^^^^^^^^^^^^^^ the trait `AtomicOrd` is not implemented for `f32` + | +help: the following other types implement trait `AtomicOrd` + --> src/_core.rs + | + | impl AtomicOrd for i32 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `i32` + | impl AtomicOrd for u32 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `u32` + | impl AtomicOrd for i64 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `i64` + | impl AtomicOrd for u64 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `u64` +note: required by a bound in `cutile::core::MaxTensor_2` + --> src/_core.rs + | + | pub struct MaxTensor { + | ^^^^^^^^^^^^^^^^^ required by this bound in `MaxTensor_2` + = note: this error originates in the attribute macro `cutile::module` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: the method `partition` exists for reference `&cutile::core::MaxTensor_2`, but its trait bounds were not satisfied + --> tests/ui/max_tensor_requires_an_ordered_element.rs:12:56 + | +12 | let acc: MaxPartition = out.partition(shape![16, 16]); + | ^^^^^^^^^ method cannot be called on `&cutile::core::MaxTensor_2` due to unsatisfied trait bounds + | + ::: src/_core.rs + | + | pub struct MaxTensor { + | ------------------------------------------------------------- doesn't satisfy `<_ as DeviceOp>::Output = _`, `_: DeviceOp`, `_: PartitionOp<_, MaxTensor_2>` or `cutile::core::MaxTensor_2: Iterator` + | + = note: the following trait bounds were not satisfied: + `f32: AtomicOrd` + `<&cutile::core::MaxTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &cutile::core::MaxTensor_2>` + `&cutile::core::MaxTensor_2: cutile::prelude::DeviceOp` + which is required by `&cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &cutile::core::MaxTensor_2>` + `<&&cutile::core::MaxTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&&cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &&cutile::core::MaxTensor_2>` + `&&cutile::core::MaxTensor_2: cutile::prelude::DeviceOp` + which is required by `&&cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &&cutile::core::MaxTensor_2>` + `<&mut &cutile::core::MaxTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&mut &cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &mut &cutile::core::MaxTensor_2>` + `&mut &cutile::core::MaxTensor_2: cutile::prelude::DeviceOp` + which is required by `&mut &cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &mut &cutile::core::MaxTensor_2>` + `&cutile::core::MaxTensor_2: Iterator` + which is required by `&mut &cutile::core::MaxTensor_2: Iterator` + ` as cutile::prelude::DeviceOp>::Output = _` + which is required by `cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, cutile::core::MaxTensor_2>` + `cutile::core::MaxTensor_2: cutile::prelude::DeviceOp` + which is required by `cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, cutile::core::MaxTensor_2>` + `<&mut cutile::core::MaxTensor_2 as cutile::prelude::DeviceOp>::Output = _` + which is required by `&mut cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &mut cutile::core::MaxTensor_2>` + `&mut cutile::core::MaxTensor_2: cutile::prelude::DeviceOp` + which is required by `&mut cutile::core::MaxTensor_2: cutile::prelude::PartitionOp<_, &mut cutile::core::MaxTensor_2>` + `cutile::core::MaxTensor_2: Iterator` + which is required by `&mut cutile::core::MaxTensor_2: Iterator` + +error[E0277]: the trait bound `f32: AtomicOrd` is not satisfied + --> tests/ui/max_tensor_requires_an_ordered_element.rs:5:1 + | +5 | #[cutile::module] + | ^^^^^^^^^^^^^^^^^ the trait `AtomicOrd` is not implemented for `f32` + | +help: the following other types implement trait `AtomicOrd` + --> src/_core.rs + | + | impl AtomicOrd for i32 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `i32` + | impl AtomicOrd for u32 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `u32` + | impl AtomicOrd for i64 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `i64` + | impl AtomicOrd for u64 {} + | ^^^^^^^^^^^^^^^^^^^^^^ `u64` +note: required by a bound in `cutile::core::MaxPartition_2` + --> src/_core.rs + | + | pub struct MaxPartition<'a, E: atomic::AtomicOrd, const D: [i32; N]> { + | ^^^^^^^^^^^^^^^^^ required by this bound in `MaxPartition_2` + = note: this error originates in the attribute macro `cutile::module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cutile-rs/cutile/tests/ui_aggregation.rs b/cutile-rs/cutile/tests/ui_aggregation.rs new file mode 100644 index 0000000000..53367c3ead --- /dev/null +++ b/cutile-rs/cutile/tests/ui_aggregation.rs @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail tests for the aggregation capabilities. +//! +//! A lane of its own so CPU CI can run it. The whole `tests/ui` suite cannot: +//! `graph_scope_rejects_allocating_input.stderr` is a stale snapshot holding a +//! `$WORKSPACE/cuda-async/...` path the merged-repository layout cannot +//! reproduce, so `--test ui` fails on unmodified `main`. Rewriting that +//! snapshot would bless whatever the diagnostic says today, which is a +//! separate decision from this capability's; these four cases are moved here +//! instead, where their expectations are checked on every CPU CI run. +//! +//! Each case in `tests/ui/*.rs` is a program that must NOT compile, paired +//! with the expected diagnostics in the `.stderr` file next to it. These need +//! no GPU; they only build against `cutile`. Regenerate the expectations with +//! `TRYBUILD=overwrite cargo test -p cutile --test ui_aggregation` after an +//! intentional diagnostic change — and read the regenerated file, because the +//! same command will also bless a real regression. + +/// An aggregation capability is a narrow capability, not a tensor. `partition` is +/// its only view constructor and neither it nor `AggregationPartition` +/// exposes ordinary load or store access; `atomic_add` is restricted to the +/// float elements the operation matrix admits; the reference is shared by +/// construction; and binding one needs exclusive host authority over the +/// backing tensor. +#[test] +fn ui_aggregation() { + let t = trybuild::TestCases::new(); + t.compile_fail("tests/ui/aggregation_tensor_withholds_ordinary_access.rs"); + t.compile_fail("tests/ui/aggregation_tensor_requires_a_supported_element.rs"); + t.compile_fail("tests/ui/aggregation_withholds_other_operations.rs"); + t.compile_fail("tests/ui/max_tensor_requires_an_ordered_element.rs"); + t.compile_fail("tests/ui/aggregation_binding_cannot_be_reimplemented.rs"); + t.compile_fail("tests/ui/aggregation_tensor_rejects_mutable_reference.rs"); + t.compile_fail("tests/ui/aggregation_tensor_needs_exclusive_host_authority.rs"); +} diff --git a/cutile-rs/scripts/run_cpu_tests.sh b/cutile-rs/scripts/run_cpu_tests.sh index 52f457cd3c..0cdda12458 100755 --- a/cutile-rs/scripts/run_cpu_tests.sh +++ b/cutile-rs/scripts/run_cpu_tests.sh @@ -46,6 +46,7 @@ run_step \ cargo test -p cutile --doc for test_target in \ + aggregation_tensor \ basics_and_inlining \ binary_math_ops \ bitwise_and_bitcast_ops \ @@ -73,6 +74,7 @@ for test_target in \ token_threading \ two_cga_trait_impl \ type_inference_sanity \ + ui_aggregation \ unary_math_ops \ use_classifier do