From c576839b4d05284d49f17185c7a34f06fece32ec Mon Sep 17 00:00:00 2001 From: Robert Ream <154010+robertream@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:58:16 -0700 Subject: [PATCH] fix(cutile-compiler): a view op declared `-> Token` evaluates to its token Signed-off-by: Robert Ream <154010+robertream@users.noreply.github.com> --- .../src/compiler/compile_cuda_tile_op.rs | 32 +- cutile-rs/cutile/tests/ir_inspection/mod.rs | 96 +++++ .../cutile/tests/memory_and_atomic_ops.rs | 347 ++++++++++++++++++ cutile-rs/cutile/tests/token_threading.rs | 48 ++- 4 files changed, 500 insertions(+), 23 deletions(-) create mode 100644 cutile-rs/cutile/tests/ir_inspection/mod.rs 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..5ddb427155 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 @@ -1679,16 +1679,12 @@ impl<'m> CUDATileFunctionCompiler<'m> { .build(module); append_op(module, block_id, op_id); let _old = super::shared_utils::update_token(view_arg, results[0], ctx); - let Some(var_arg_ident) = get_ident_from_expr(view_arg) else { - return self.jit_error_result(&view_arg.span(), "Unexpected expression"); - }; - let Some(result) = ctx.vars.get(var_arg_ident.to_string().as_str()) else { - return self.jit_error_result( - &view_arg.span(), - &format!("Unexpected state: Expected {var_arg_ident} in ctx"), - ); - }; - Ok(Some(result.clone())) + let token_type = self + .compile_type(&syn::parse_quote!(Token), generic_args, &HashMap::new())? + .ok_or_else(|| self.jit_error(&call_expr.span(), "unable to resolve Token"))?; + Ok(Some(TileRustValue::new_primitive( + results[0], token_type, None, + ))) } /// Reads `memory_ordering`, `memory_scope` ZST args from a view op call. @@ -2105,16 +2101,12 @@ impl<'m> CUDATileFunctionCompiler<'m> { let (op_id, results) = op_builder.build(module); append_op(module, block_id, op_id); let _old = super::shared_utils::update_token(view_arg, results[0], ctx); - let Some(var_arg_ident) = get_ident_from_expr(view_arg) else { - return self.jit_error_result(&view_arg.span(), "Unexpected expression"); - }; - let Some(result) = ctx.vars.get(var_arg_ident.to_string().as_str()) else { - return self.jit_error_result( - &view_arg.span(), - &format!("Unexpected state: Expected {var_arg_ident} in ctx"), - ); - }; - Ok(Some(result.clone())) + let token_type = self + .compile_type(&syn::parse_quote!(Token), generic_args, &HashMap::new())? + .ok_or_else(|| self.jit_error(&call_expr.span(), "unable to resolve Token"))?; + Ok(Some(TileRustValue::new_primitive( + results[0], token_type, None, + ))) } fn compile_print_tko( diff --git a/cutile-rs/cutile/tests/ir_inspection/mod.rs b/cutile-rs/cutile/tests/ir_inspection/mod.rs new file mode 100644 index 0000000000..6fa1ef11c8 --- /dev/null +++ b/cutile-rs/cutile/tests/ir_inspection/mod.rs @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural inspection of `cutile_ir::Module`. +//! +//! `ir_text()` cannot answer questions about which SSA value an operand uses: +//! the printer renders several operations as `` and does not +//! identify an operand's producer. These helpers inspect `cutile_ir::Module` +//! directly, allowing tests to assert the producer and type of a specific +//! operand. + +use cutile_ir::bytecode::Opcode; +use cutile_ir::ir::{OpId, Type, Value, ValueProducer}; +use cutile_ir::Module; + +/// A borrowed module, queried structurally. `Copy`, so the iterators below +/// borrow the module rather than the inspection. +#[derive(Clone, Copy)] +pub struct ModuleInspection<'a> { + module: &'a Module, +} + +/// Grants [`Module`] the [`inspect`](ModuleInspectionExt::inspect) method. +/// Import it at the call site for `.inspect()` to resolve. +pub trait ModuleInspectionExt { + /// Borrows this module for structural queries. + fn inspect(&self) -> ModuleInspection<'_>; +} + +impl ModuleInspectionExt for Module { + fn inspect(&self) -> ModuleInspection<'_> { + ModuleInspection { module: self } + } +} + +impl<'a> ModuleInspection<'a> { + /// Immediate child operations of `op`, in region/block order. + fn child_ops(self, op: OpId) -> impl Iterator + 'a { + let module = self.module; + module + .op(op) + .regions + .iter() + .copied() + .flat_map(move |region| module.region(region).blocks.iter().copied()) + .flat_map(move |block| module.block(block).ops.iter().copied()) + } + + /// `op` and all its descendants, in preorder. + fn nested_ops(self, op: OpId) -> Box + 'a> { + Box::new( + std::iter::once(op).chain( + self.child_ops(op) + .flat_map(move |child| self.nested_ops(child)), + ), + ) + } + + /// Every operation in the module, nested regions included. + /// + /// A bounds check or a loop puts operations inside a nested region, so a + /// flat walk over a function's own blocks misses them. + fn ops_in_order(self) -> impl Iterator + 'a { + self.module + .functions + .iter() + .copied() + .flat_map(move |function| self.nested_ops(function)) + } + + /// Operations with opcode `opcode`, in order. + pub fn ops_with_opcode(self, opcode: Opcode) -> Vec { + let module = self.module; + self.ops_in_order() + .filter(move |&op| module.op(op).opcode == opcode) + .collect() + } + + /// `(producing opcode, value type)` of `value`; `None` for a block + /// argument. + /// + /// Producer *and* type. SSA identity on its own proves nothing about a + /// token slot: a handler that returned its view instead of its token + /// handed back a value that was also reachable from the view, so two + /// consumers of "the same value" can agree while both hold the wrong kind + /// of thing. + pub fn producer_and_type(self, value: Value) -> (Option, Type) { + let producer = match self.module.value_data(value).producer { + ValueProducer::OpResult { op, .. } => Some(self.module.op(op).opcode), + ValueProducer::BlockArg { .. } => None, + }; + (producer, self.module.value_type(value).clone()) + } +} diff --git a/cutile-rs/cutile/tests/memory_and_atomic_ops.rs b/cutile-rs/cutile/tests/memory_and_atomic_ops.rs index a1ed670c3a..b226586faf 100644 --- a/cutile-rs/cutile/tests/memory_and_atomic_ops.rs +++ b/cutile-rs/cutile/tests/memory_and_atomic_ops.rs @@ -2,9 +2,15 @@ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +use cutile_compiler::compile_api::KernelCompiler; use cutile_compiler::compiler::utils::CompileOptions; +use cutile_ir::bytecode::{BytecodeVersion, Opcode}; +use cutile_ir::ir::Type; mod common; +mod ir_inspection; + +use crate::ir_inspection::ModuleInspectionExt; #[cutile::module] mod memory_and_atomic_ops_module { @@ -552,6 +558,347 @@ mod atomic_red_view_module { } } +/// Both spellings of the atomic reduction, each threading the token it +/// returns into a following store's ordering slot. +/// +/// `atomic_red_view_tko` and `atomic_red_view_raw` lower to the same +/// `cuda_tile.atomic_red_view_tko` op through different compiler handlers, and +/// both declare `-> Token`. The only way to see which value a call actually +/// evaluates to is to consume it where a token is required and look at the +/// operand. +#[cutile::module] +mod atomic_red_view_token_module { + use cutile::core::*; + use cutile::tileir::*; + + #[cutile::entry()] + unsafe fn convenience_form(output: &mut Tensor) { + let mut part: PartitionMut = output.partition_mut(shape![16, 16]); + let pid: (i32, i32, i32) = get_tile_block_id(); + let tile: Tile = constant(1.0f32, shape![16, 16]); + let reduced: Token = unsafe { + atomic_red_view_tko( + &mut part, + tile, + [pid.0, 0i32], + atomic::AddF, + ordering::Relaxed, + scope::Device, + ) + }; + let _stored: Token = unsafe { + store_view_raw( + &part, + tile, + [pid.0, 0i32], + Some(reduced), + ordering::Weak, + scope::Device, + tma::Disabled, + None, + [false, false], + ) + }; + } + + #[cutile::entry()] + unsafe fn raw_form(output: &mut Tensor) { + let mut part: PartitionMut = output.partition_mut(shape![16, 16]); + let pid: (i32, i32, i32) = get_tile_block_id(); + let tile: Tile = constant(1.0f32, shape![16, 16]); + let reduced: Token = unsafe { + atomic_red_view_raw( + &mut part, + tile, + [pid.0, 0i32], + atomic::AddF, + ordering::Relaxed, + scope::Device, + None, + ) + }; + let _stored: Token = unsafe { + store_view_raw( + &part, + tile, + [pid.0, 0i32], + Some(reduced), + ordering::Weak, + scope::Device, + tma::Disabled, + None, + [false, false], + ) + }; + } +} + +/// The ordering token of the module's single store, as `(producer opcode, +/// operand type)`. +/// +/// Read off `cutile_ir::Module`, not the IR text: the printer renders +/// `atomic_red_view_tko` as ``, so text matching cannot see +/// this. The token is the store's last operand: `operandSegmentSizes` is +/// `[view, index.., value, token]`. +fn store_ordering_token(function_name: &str) -> (Opcode, Type) { + let module = KernelCompiler::new( + atomic_red_view_token_module::__module_ast_self, + "atomic_red_view_token_module", + function_name, + ) + .target("sm_120") + .strides(&[("output", &[-1, 1])]) + .compile() + .expect("compilation failed") + .into_module(); + + let stores = module.inspect().ops_with_opcode(Opcode::StoreViewTko); + assert_eq!(stores.len(), 1, "{function_name}: expected one store"); + let token = *module + .op(stores[0]) + .operands + .last() + .expect("a store has operands"); + let (producer, ty) = module.inspect().producer_and_type(token); + let producer = producer + .unwrap_or_else(|| panic!("{function_name}: the ordering token is a block argument")); + (producer, ty) +} + +/// `atomic_red_view_tko` evaluates to its completion token, like the raw op. +#[test] +fn atomic_red_view_tko_evaluates_to_its_completion_token() { + common::with_test_stack(|| { + let raw = store_ordering_token("raw_form"); + assert_eq!( + raw, + (Opcode::AtomicRedViewTko, Type::Token), + "`atomic_red_view_raw`'s token must order the following store" + ); + assert_eq!( + store_ordering_token("convenience_form"), + raw, + "`atomic_red_view_tko` must evaluate to the same thing the raw op does" + ); + }); +} + +/// Every DSL store that routes to the `store_view_tko` compiler handler, +/// each threading the token it returns into a `gdc_launch_dependents_tko` +/// and then making a second store through the same view. +/// +/// Three declarations share that handler, `store_view_tko_mut`, +/// `store_view_tko_bounded` and `store_view_tko_mapped_mut`, and all three +/// are declared `-> Token`. The signal's operand is the value the call +/// evaluated to; the second store's token operand is the token the handler +/// threaded onto the view. Both must be the first store's completion token. +#[cutile::module] +mod store_view_token_module { + use cutile::core::*; + + /// `PartitionMut::store`, the safe, bounds-checked store. + #[cutile::entry()] + fn partition_mut_store(output: &mut Tensor) { + let mut part: PartitionMut = output.partition_mut(shape![16, 16]); + let tile: Tile = constant(1.0f32, shape![16, 16]); + let stored: Token = part.store(tile, [0i32, 0i32]); + let _signal: Token = unsafe { gdc_launch_dependents_tko(Some(stored)) }; + let _second: Token = part.store(tile, [0i32, 0i32]); + } + + /// `PartitionMut::store_unchecked`, the same op without the bounds proof. + #[cutile::entry()] + fn partition_mut_store_unchecked(output: &mut Tensor) { + let mut part: PartitionMut = output.partition_mut(shape![16, 16]); + let tile: Tile = constant(1.0f32, shape![16, 16]); + let stored: Token = unsafe { part.store_unchecked(tile, [0i32, 0i32]) }; + let _signal: Token = unsafe { gdc_launch_dependents_tko(Some(stored)) }; + let _second: Token = unsafe { part.store_unchecked(tile, [0i32, 0i32]) }; + } + + /// `BoundedPartitionMut::store`, the branded-coordinate store, reached + /// through the non-deprecated spelling of `PartitionMut::with_bounds`. + #[cutile::entry()] + fn bounded_partition_mut_store(output: &mut Tensor) { + let part: PartitionMut = output.partition_mut(shape![16, 16]); + let rows: Dim = num_tiles(&part, 0).into_dim(); + let cols: Dim = num_tiles(&part, 1).into_dim(); + let mut bounded: BoundedPartitionMut = + partition_with_bounds_mut(part, (rows, cols)); + let tile: Tile = constant(1.0f32, shape![16, 16]); + for i in rows { + for j in cols { + let stored: Token = bounded.store(tile, coord((i, j))); + let _signal: Token = unsafe { gdc_launch_dependents_tko(Some(stored)) }; + let _second: Token = bounded.store(tile, coord((i, j))); + } + } + } + + /// `MappedPartitionMut::store`, the schedule-minted store. + #[cutile::entry()] + fn mapped_partition_mut_store( + mut output: MappedPartitionMut, + ) { + let tile: Tile = constant(1.0f32, shape![16, 16]); + for index in output.iter_indices() { + let stored: Token = output.store(tile, index); + let _signal: Token = unsafe { gdc_launch_dependents_tko(Some(stored)) }; + let _second: Token = output.store(tile, index); + } + } +} + +/// For one kernel of `store_view_token_module`: the facts about the token the +/// store *evaluated to* (the `gdc_launch_dependents_tko` operand) and about +/// the token the handler threaded onto the view (the second store's ordering +/// operand). +/// +/// Read off `cutile_ir::Module`, not `ir_text()`: the text printer renders +/// some of these ops as ``, and it never names the producing +/// operation at all. +fn store_token_facts( + function_name: &str, + generics: &[&str], + strides: &[(&str, &[i32])], +) -> ((Option, Type), (Option, Type)) { + let module = KernelCompiler::new( + store_view_token_module::__module_ast_self, + "store_view_token_module", + function_name, + ) + .target("sm_120") + .bytecode_version(BytecodeVersion::V13_4) + .generics(generics.iter().map(|s| s.to_string()).collect::>()) + .strides(strides) + .options(CompileOptions::default()) + .compile() + .unwrap_or_else(|e| panic!("compile {function_name}: {e}")) + .into_module(); + + let stores = module.inspect().ops_with_opcode(Opcode::StoreViewTko); + assert_eq!( + stores.len(), + 2, + "{function_name}: expected exactly two stores through the view" + ); + let signals = module + .inspect() + .ops_with_opcode(Opcode::GdcLaunchDependentsTko); + assert_eq!( + signals.len(), + 1, + "{function_name}: expected exactly one launch-dependents signal" + ); + + // `gdc_launch_dependents_tko`'s only operand is the token it signals + // after: the value the first store evaluated to. + let returned = *module + .op(signals[0]) + .operands + .first() + .expect("the signal has a token operand"); + // The second store's last operand is its ordering token, read from the + // view's `token` metadata field, what `update_token` put there. + let threaded = *module + .op(stores[1]) + .operands + .last() + .expect("a store has operands"); + ( + module.inspect().producer_and_type(returned), + module.inspect().producer_and_type(threaded), + ) +} + +/// `(Opcode::StoreViewTko, Type::Token)`, what both tokens must be. +fn store_completion_token() -> (Option, Type) { + (Some(Opcode::StoreViewTko), Type::Token) +} + +/// `PartitionMut::store`, the safe store path, evaluates to its completion +/// token and threads the same token onto the view. +#[test] +fn partition_mut_store_evaluates_to_its_completion_token() { + common::with_test_stack(|| { + let (returned, threaded) = + store_token_facts("partition_mut_store", &[], &[("output", &[-1, 1])]); + assert_eq!( + returned, + store_completion_token(), + "`PartitionMut::store` must evaluate to its completion token" + ); + assert_eq!( + threaded, + store_completion_token(), + "`PartitionMut::store` must also thread that token onto the view" + ); + }); +} + +/// The same for `PartitionMut::store_unchecked`, which shares the op. +#[test] +fn partition_mut_store_unchecked_evaluates_to_its_completion_token() { + common::with_test_stack(|| { + let (returned, threaded) = store_token_facts( + "partition_mut_store_unchecked", + &[], + &[("output", &[-1, 1])], + ); + assert_eq!( + returned, + store_completion_token(), + "`PartitionMut::store_unchecked` must evaluate to its completion token" + ); + assert_eq!( + threaded, + store_completion_token(), + "`PartitionMut::store_unchecked` must also thread that token onto the view" + ); + }); +} + +/// The same for `BoundedPartitionMut::store`, via `store_view_tko_bounded`. +#[test] +fn bounded_partition_mut_store_evaluates_to_its_completion_token() { + common::with_test_stack(|| { + let (returned, threaded) = + store_token_facts("bounded_partition_mut_store", &[], &[("output", &[-1, 1])]); + assert_eq!( + returned, + store_completion_token(), + "`BoundedPartitionMut::store` must evaluate to its completion token" + ); + assert_eq!( + threaded, + store_completion_token(), + "`BoundedPartitionMut::store` must also thread that token onto the view" + ); + }); +} + +/// The same for `MappedPartitionMut::store`, via `store_view_tko_mapped_mut`. +#[test] +fn mapped_partition_mut_store_evaluates_to_its_completion_token() { + common::with_test_stack(|| { + let (returned, threaded) = store_token_facts( + "mapped_partition_mut_store", + &["4", "4"], + &[("output", &[-1, 1])], + ); + assert_eq!( + returned, + store_completion_token(), + "`MappedPartitionMut::store` must evaluate to its completion token" + ); + assert_eq!( + threaded, + store_completion_token(), + "`MappedPartitionMut::store` must also thread that token onto the view" + ); + }); +} + fn compile_ir(function_name: &str, generics: &[String], strides: &[(&str, &[i32])]) -> String { common::compile_to_ir( __module_ast_self, diff --git a/cutile-rs/cutile/tests/token_threading.rs b/cutile-rs/cutile/tests/token_threading.rs index 9b8bf95fc2..be890f8f0a 100644 --- a/cutile-rs/cutile/tests/token_threading.rs +++ b/cutile-rs/cutile/tests/token_threading.rs @@ -13,9 +13,13 @@ use cutile_compiler::compile_api::KernelCompiler; use cutile_compiler::compiler::utils::CompileOptions; -use cutile_ir::bytecode::BytecodeVersion; +use cutile_ir::bytecode::{BytecodeVersion, Opcode}; +use cutile_ir::ir::Type; mod common; +mod ir_inspection; + +use crate::ir_inspection::ModuleInspectionExt; #[cutile::module] mod token_module { @@ -266,6 +270,13 @@ mod token_module { use token_module::__module_ast_self as token_module_ast; fn compile(entry: &str, args: &[(&str, &[i32])]) -> String { + compile_artifacts(entry, args).ir_text() +} + +fn compile_artifacts( + entry: &str, + args: &[(&str, &[i32])], +) -> cutile_compiler::compile_api::CompileArtifacts { KernelCompiler::new(token_module_ast, "token_module", entry) .target("sm_120") .bytecode_version(BytecodeVersion::V13_4) @@ -273,7 +284,6 @@ fn compile(entry: &str, args: &[(&str, &[i32])]) -> String { .options(CompileOptions::default()) .compile() .unwrap_or_else(|e| panic!("compile {entry}: {e:?}")) - .ir_text() } /// Lines of `ir` containing `op`, in order. @@ -318,10 +328,22 @@ fn second_store_consumes_the_first_stores_token() { }); } +/// `Tensor::store`'s result and `Tensor::token` are the same value, *and* +/// that value is the store's completion token. +/// +/// SSA identity alone would be a weak claim: it holds for any value both +/// spellings happen to read, the view included. `Tensor::store` reaches +/// `store_view_tko_mut` through `store_tile_*`, which discards the call's +/// result and installs `get_partition_token_mut(&view)` on the tensor +/// instead, so both spellings read the view's threaded token and the +/// identity holds no matter what the op's handler evaluates to. The producer +/// and type assertions below are what pins the shared value to the store. #[test] fn token_method_reads_the_stores_token() { common::with_test_stack(|| { - let ir = compile("store_token_is_tensor_token", &[("z", &[1]), ("x", &[1])]); + let artifacts = + compile_artifacts("store_token_is_tensor_token", &[("z", &[1]), ("x", &[1])]); + let ir = artifacts.ir_text(); let store = result_of(lines_with(&ir, "store_view_tko")[0]); let signals = lines_with(&ir, "gdc_launch_dependents_tko"); assert_eq!(signals.len(), 2, "IR:\n{ir}"); @@ -332,6 +354,26 @@ fn token_method_reads_the_stores_token() { "both spellings must yield the store's token.\nIR:\n{ir}" ); } + + // Producer and type, off `cutile_ir::Module`: the value both + // spellings agree on must be a `Token` produced by the store. + let module = artifacts.module(); + let signals = module + .inspect() + .ops_with_opcode(Opcode::GdcLaunchDependentsTko); + assert_eq!(signals.len(), 2, "IR:\n{ir}"); + for signal in signals { + let token = *module + .op(signal) + .operands + .first() + .expect("the signal has a token operand"); + assert_eq!( + module.inspect().producer_and_type(token), + (Some(Opcode::StoreViewTko), Type::Token), + "both spellings must yield the store's completion token.\nIR:\n{ir}" + ); + } }); }