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