From b5428c4ef3bce328f0484b5b6e6d63871a20d984 Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:17:25 +0900 Subject: [PATCH 1/4] perf(cute): look up ordinary provenance facts by carrier `OrdinaryFacts` kept one flat map keyed by (value, field path), so `entries(value)` scanned every key in the map. The ordinary plumbing and the poison pass call it for every cast, aggregate, store, load and block argument of a scope, which made each sweep quadratic in the scope's size. Key the facts by SSA value first and by field path second. `entries`, `at_path` and `copy_entries` now touch only the named carrier, and `copy_entries` still forwards a snapshot of the source taken before the copy. The shared-MMA use closure walks the nested map and tests scope membership through a set instead of a linear search. Signed-off-by: juliankang4 <200967947+juliankang4@users.noreply.github.com> --- .../dialect-cute/src/verify/mma.rs | 79 ++++++++++--------- .../dialect-cute/src/verify/provenance.rs | 34 ++++---- .../dialect-cute/src/verify/tests.rs | 49 +++++++++++- 3 files changed, 104 insertions(+), 58 deletions(-) diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs index ef3389c4cf..a8bfe5b72d 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs @@ -23,9 +23,9 @@ use super::{ ghosts::smem_mma_semantic_ids, invalid, provenance::{ - CarrierKey, OrdinaryFacts, branch_successor_operands, is_semantic_pointer_carrier_cast, - local_cell, local_cell_is_closed, poison_missing_forwarding_sources, - propagate_ordinary_plumbing, seed_unknown_integer_carriers, + OrdinaryFacts, branch_successor_operands, is_semantic_pointer_carrier_cast, local_cell, + local_cell_is_closed, poison_missing_forwarding_sources, propagate_ordinary_plumbing, + seed_unknown_integer_carriers, }, }; use crate::{ @@ -360,7 +360,8 @@ fn mma_branch_forwards_fact( fn mma_ordinary_use_forwards_fact( ctx: &Context, facts: &OrdinaryFacts, - key: &CarrierKey, + value: Value, + path: &[u32], fact: MmaFact, user: Ptr, operand_index: usize, @@ -372,14 +373,14 @@ fn mma_ordinary_use_forwards_fact( && is_semantic_pointer_carrier_cast(ctx, user) && operation.get_num_results() == 1 { - return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + return mma_fact_at_path(facts, operation.get_result(0), path, fact); } if (opid == MirConstructStructOp::get_opid_static() || opid == MirConstructTupleOp::get_opid_static()) && operation.get_num_results() == 1 { let mut destination_path = vec![operand_index as u32]; - destination_path.extend_from_slice(&key.path); + destination_path.extend_from_slice(path); return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); } if opid == MirInsertFieldOp::get_opid_static() @@ -390,14 +391,14 @@ fn mma_ordinary_use_forwards_fact( return false; }; if operand_index == 0 { - if key.path.first().copied() == Some(index.0) { + if path.first().copied() == Some(index.0) { return true; } - return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + return mma_fact_at_path(facts, operation.get_result(0), path, fact); } if operand_index == 1 { let mut destination_path = vec![index.0]; - destination_path.extend_from_slice(&key.path); + destination_path.extend_from_slice(path); return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); } return false; @@ -409,10 +410,10 @@ fn mma_ordinary_use_forwards_fact( let Some(index) = MirExtractFieldOp::new(user).get_attr_index(ctx) else { return false; }; - if key.path.first().copied() != Some(index.0) { + if path.first().copied() != Some(index.0) { return true; } - return mma_fact_at_path(facts, operation.get_result(0), &key.path[1..], fact); + return mma_fact_at_path(facts, operation.get_result(0), &path[1..], fact); } if opid == MirStoreOp::get_opid_static() && operation.get_num_operands() == 2 { let store = MirStoreOp::new(user); @@ -424,7 +425,7 @@ fn mma_ordinary_use_forwards_fact( && let Some((root, mut destination_path)) = local_cell(ctx, store.address_opd(ctx)) && local_cell_is_closed(ctx, root) { - destination_path.extend_from_slice(&key.path); + destination_path.extend_from_slice(path); return mma_fact_at_path(facts, root, &destination_path, fact); } return false; @@ -440,7 +441,7 @@ fn mma_ordinary_use_forwards_fact( if !local_cell_is_closed(ctx, root) { return false; } - let Some(result_path) = key.path.strip_prefix(cell_path.as_slice()) else { + let Some(result_path) = path.strip_prefix(cell_path.as_slice()) else { return true; }; return mma_fact_at_path(facts, operation.get_result(0), result_path, fact); @@ -458,7 +459,7 @@ fn mma_ordinary_use_forwards_fact( if opid == MirCondBranchOp::get_opid_static() && operand_index == 0 { return false; } - mma_branch_forwards_fact(ctx, facts, user, key.value, &key.path, fact) + mma_branch_forwards_fact(ctx, facts, user, value, path, fact) } /// Reject ordinary uses that discard or expose semantic MMA carriers outside recognized forwarding. @@ -467,31 +468,35 @@ pub(super) fn verify_mma_use_closure( operations: &[Ptr], facts: &OrdinaryFacts, ) -> Result { - for (key, found) in &facts.facts { - for fact in found.iter().copied() { - if matches!(fact, MmaFact::Scalar(_) | MmaFact::Unknown(_)) { - continue; - } - for r#use in key.value.uses(ctx) { - let user = r#use.user_op(); - let operand_index = r#use.find_index(ctx); - if operations.contains(&user) - && (mma_semantic_consumer(ctx, user, operand_index, fact) - || mma_ordinary_use_forwards_fact( - ctx, - facts, - key, - fact, - user, - operand_index, - )) - { + let scope: HashSet> = operations.iter().copied().collect(); + for (value, paths) in &facts.facts { + for (path, found) in paths { + for fact in found.iter().copied() { + if matches!(fact, MmaFact::Scalar(_) | MmaFact::Unknown(_)) { continue; } - return Err(invalid(format!( - "shared-MMA carrier has an unsupported outgoing use by `{}` operand {operand_index}", - Operation::get_opid(user, ctx) - ))); + for r#use in value.uses(ctx) { + let user = r#use.user_op(); + let operand_index = r#use.find_index(ctx); + if scope.contains(&user) + && (mma_semantic_consumer(ctx, user, operand_index, fact) + || mma_ordinary_use_forwards_fact( + ctx, + facts, + *value, + path, + fact, + user, + operand_index, + )) + { + continue; + } + return Err(invalid(format!( + "shared-MMA carrier has an unsupported outgoing use by `{}` operand {operand_index}", + Operation::get_opid(user, ctx) + ))); + } } } } diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs index 7bd4ae9a83..e682e76803 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs @@ -35,21 +35,15 @@ use pliron::{ value::Value, }; -/// Identity of a whole SSA carrier or one nested aggregate field. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub(super) struct CarrierKey { - /// The whole SSA carrier before following any aggregate field path. - pub(super) value: Value, - /// Empty means the value itself. Non-empty paths describe a field inside - /// an ordinary aggregate before a matching `mir.extract_field` exposes it. - pub(super) path: Vec, -} - /// Monotone sets of possible semantic origins, indexed by carrier and field path. #[derive(Clone)] pub(super) struct OrdinaryFacts { - /// Accumulated origins; multiple entries at one key represent a merge that consumers must resolve. - pub(super) facts: HashMap>, + /// Accumulated origins by SSA carrier, then by field path inside it. An + /// empty path means the value itself; a non-empty path describes a field + /// inside an ordinary aggregate before a matching `mir.extract_field` + /// exposes it. Multiple origins at one path represent a merge that + /// consumers must resolve. + pub(super) facts: HashMap, HashSet>>, } impl Default for OrdinaryFacts { @@ -67,7 +61,9 @@ where /// Add one origin to a field path and report whether the fixed-point state changed. pub(super) fn insert_path(&mut self, value: Value, path: Vec, fact: F) -> bool { self.facts - .entry(CarrierKey { value, path }) + .entry(value) + .or_default() + .entry(path) .or_default() .insert(fact) } @@ -79,10 +75,7 @@ where /// Borrow the possible origins at an exact carrier field path. pub(super) fn at_path(&self, value: Value, path: &[u32]) -> Option<&HashSet> { - self.facts.get(&CarrierKey { - value, - path: path.to_vec(), - }) + self.facts.get(&value)?.get(path) } /// Borrow origins attached directly to a whole SSA carrier. @@ -93,9 +86,10 @@ where /// Collect all nested field origins carried by one SSA value. pub(super) fn entries(&self, value: Value) -> Vec<(Vec, F)> { self.facts - .iter() - .filter(|(key, _)| key.value == value) - .flat_map(|(key, facts)| facts.iter().cloned().map(|fact| (key.path.clone(), fact))) + .get(&value) + .into_iter() + .flatten() + .flat_map(|(path, facts)| facts.iter().cloned().map(|fact| (path.clone(), fact))) .collect() } diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs index 30d90d2765..18acd059c9 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs @@ -34,7 +34,8 @@ use super::{ mma::{build_mma_facts, require_smem_overlay_pair, verify_mma_use_closure}, pipeline::require_producer_pipeline_order, provenance::{ - is_local_cell_pointer_alias_cast, is_semantic_pointer_carrier_cast, local_cell_is_closed, + OrdinaryFacts, is_local_cell_pointer_alias_cast, is_semantic_pointer_carrier_cast, + local_cell_is_closed, }, scheduler::{build_scheduler_facts, scheduler_root}, *, @@ -388,6 +389,52 @@ fn factless_block_argument_store_poisons_a_valid_scheduler_cell() { assert!(error.contains("merges values"), "{error}"); } +/// Fact lookups see only the named carrier, and a copy forwards the source as it was when the copy began. +#[test] +fn ordinary_facts_forward_a_snapshot_of_one_carrier() { + let mut ctx = Context::new(); + let (_, block) = module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let source = undef(&mut ctx, block, u64_ty); + let other = undef(&mut ctx, block, u64_ty); + let destination = undef(&mut ctx, block, u64_ty); + let sorted = |facts: &OrdinaryFacts, value| { + let mut entries = facts.entries(value); + entries.sort(); + entries + }; + + let mut facts = OrdinaryFacts::default(); + assert!(facts.insert_path(source, vec![0], 1)); + assert!(facts.insert_path(source, vec![1, 2], 2)); + assert!(!facts.insert_path(source, vec![1, 2], 2)); + assert!(facts.insert(other, 3)); + assert_eq!(sorted(&facts, source), vec![(vec![0], 1), (vec![1, 2], 2)]); + assert_eq!(sorted(&facts, other), vec![(vec![], 3)]); + assert!(sorted(&facts, destination).is_empty()); + assert_eq!( + facts.at_path(source, &[1, 2]).map(|found| found.len()), + Some(1) + ); + assert!(facts.at_path(source, &[1]).is_none()); + assert!(facts.at(other).is_some_and(|found| found.contains(&3))); + + assert!(facts.copy_entries(source, source, &[], &[5])); + assert_eq!( + sorted(&facts, source), + vec![ + (vec![0], 1), + (vec![1, 2], 2), + (vec![5, 0], 1), + (vec![5, 1, 2], 2) + ] + ); + assert!(facts.copy_entries(source, destination, &[1], &[])); + assert!(!facts.copy_entries(source, destination, &[1], &[])); + assert_eq!(sorted(&facts, destination), vec![(vec![2], 2)]); + assert_eq!(sorted(&facts, other), vec![(vec![], 3)]); +} + /// Casts that change pointee, address space, or write capability cannot establish closed-cell aliases. #[test] fn non_identity_pointer_casts_cannot_alias_a_local_semantic_cell() { From 3ddd00775a5fda6260fd7d032b6b27d82395908b Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:17:25 +0900 Subject: [PATCH 2/4] perf(cute): skip a family's fact analysis in scopes without its operations The scheduler, pipeline and shared-MMA stories built their provenance facts for every scope, including every function with no CuTe operation at all. Each analysis seeds every integer value of the scope, so this cost was paid for every ordinary kernel. Return before building the facts when the scope holds no operation of the story's family. Every check in those stories starts from such an operation, so a scope without one was already accepted after the analysis; consumer-only (malformed) stories still reach their checks because their operations belong to the family. Signed-off-by: juliankang4 <200967947+juliankang4@users.noreply.github.com> --- .../cute-rs-preview/dialect-cute/src/verify/cfg.rs | 9 ++++++++- .../cute-rs-preview/dialect-cute/src/verify/mma.rs | 5 +++++ .../cute-rs-preview/dialect-cute/src/verify/pipeline.rs | 6 +++++- .../cute-rs-preview/dialect-cute/src/verify/scheduler.rs | 6 +++++- 4 files changed, 23 insertions(+), 3 deletions(-) diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs index 6a4da1be59..5094602dd0 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs @@ -12,7 +12,7 @@ use pliron::{ basic_block::BasicBlock, context::{Context, Ptr}, linked_list::ContainsLinkedList, - op::Op, + op::{Op, OpId}, operation::Operation, }; @@ -53,6 +53,13 @@ pub(super) fn operations_by_scope( scopes.into_values().collect() } +/// Test whether a scope holds any operation of one semantic family. +pub(super) fn scope_has_any(ctx: &Context, scope: &[Ptr], ids: &[OpId]) -> bool { + scope + .iter() + .any(|operation| ids.contains(&Operation::get_opid(*operation, ctx))) +} + /// Test whether any successor path in the same region reaches the destination block. pub(super) fn block_reaches( ctx: &Context, diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs index a8bfe5b72d..d70afd2e59 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs @@ -20,6 +20,7 @@ use pliron::{ use super::{ Result, + cfg::scope_has_any, ghosts::smem_mma_semantic_ids, invalid, provenance::{ @@ -523,6 +524,10 @@ fn require_same_mma_plan( /// Validate shared overlays, lazy B partitions, scale/A fragments, and accumulator provenance together. pub(super) fn verify_smem_mma_story(ctx: &Context, all_ops: &[Ptr]) -> Result { let ids = smem_mma_semantic_ids(); + // Facts start only at this family's operations; a scope without one has nothing to check. + if !scope_has_any(ctx, all_ops, &ids) { + return Ok(()); + } let facts = build_mma_facts(ctx, all_ops); let mut overlays = HashSet::new(); let mut used_overlays = HashSet::new(); diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs index 355f94342e..4b6cb8318f 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs @@ -17,7 +17,7 @@ use pliron::{ use super::{ Result, - cfg::{operation_dominates, operation_reaches}, + cfg::{operation_dominates, operation_reaches, scope_has_any}, ghosts::pipeline_semantic_ids, invalid, provenance::{ @@ -302,6 +302,10 @@ pub(super) fn require_producer_pipeline_order( /// Validate handle ownership, complete cursor lifecycles, CFG order, and expected TMA byte totals. pub(super) fn verify_pipeline_story(ctx: &Context, all_ops: &[Ptr]) -> Result { let ids = pipeline_semantic_ids(); + // Facts start only at this family's operations; a scope without one has nothing to check. + if !scope_has_any(ctx, all_ops, &ids) { + return Ok(()); + } let facts = build_pipeline_facts(ctx, all_ops); let makes: Vec<_> = all_ops .iter() diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs index 0d4e052952..e34b57e919 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs @@ -19,7 +19,7 @@ use pliron::{ use super::{ Result, - cfg::{block_reaches, operation_reaches}, + cfg::{block_reaches, operation_reaches, scope_has_any}, ghosts::scheduler_semantic_ids, invalid, provenance::{ @@ -198,6 +198,10 @@ fn scheduler_guard( /// Check one scheduler's state origins, loop advance, work guard, and coordinate consumers. pub(super) fn verify_scheduler_story(ctx: &Context, all_ops: &[Ptr]) -> Result { let ids = scheduler_semantic_ids(); + // Facts start only at this family's operations; a scope without one has nothing to check. + if !scope_has_any(ctx, all_ops, &ids) { + return Ok(()); + } let facts = build_scheduler_facts(ctx, all_ops); let mut roots = HashSet::new(); let mut has_work_roots = HashSet::new(); From 02702d971eb3e93690e719062b51f7a993d3a28d Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:17:26 +0900 Subject: [PATCH 3/4] perf(cute): keep at most two opaque origins per provenance fact set Seeding gives every factless integer value an origin that names the value itself (`Unknown`, and `Scalar` for shared MMA). Propagation merges those origins through block arguments and closed local cells, so a value's set grew with the length of every chain that reaches it, and each sweep of the fixed point copied the grown sets: verification time grew with the cube of the scope size. Consumers read a set only through its size, its sole origin, or whether it holds a given semantic origin, and they skip opaque origins when walking a set. Two opaque origins therefore report a merge exactly as any larger number does. `insert_path` now drops a third one, through a `CarrierFact` trait that each family's fact type implements. Path presence and semantic origins are unchanged, so every accept or reject decision and the diagnostic each consumer reports stay the same. Tests cover the cap, the snapshot copy, a scheduler state carried through 48 blocks and cells, and the rejection wording for cursors that merge one to eight unrelated values. Signed-off-by: juliankang4 <200967947+juliankang4@users.noreply.github.com> --- .../dialect-cute/src/verify/mma.rs | 14 +- .../dialect-cute/src/verify/pipeline.rs | 8 +- .../dialect-cute/src/verify/provenance.rs | 36 +++-- .../dialect-cute/src/verify/scheduler.rs | 8 +- .../dialect-cute/src/verify/tests.rs | 135 ++++++++++++++++-- 5 files changed, 177 insertions(+), 24 deletions(-) diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs index d70afd2e59..c8deda70a6 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs @@ -24,9 +24,9 @@ use super::{ ghosts::smem_mma_semantic_ids, invalid, provenance::{ - OrdinaryFacts, branch_successor_operands, is_semantic_pointer_carrier_cast, local_cell, - local_cell_is_closed, poison_missing_forwarding_sources, propagate_ordinary_plumbing, - seed_unknown_integer_carriers, + CarrierFact, OrdinaryFacts, branch_successor_operands, is_semantic_pointer_carrier_cast, + local_cell, local_cell_is_closed, poison_missing_forwarding_sources, + propagate_ordinary_plumbing, seed_unknown_integer_carriers, }, }; use crate::{ @@ -102,6 +102,12 @@ pub(super) enum MmaFact { Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), } +impl CarrierFact for MmaFact { + fn is_opaque(&self) -> bool { + matches!(self, Self::Scalar(_) | Self::Unknown(_)) + } +} + /// Return the sole carrier fact, rejecting absent or conflicting facts; callers validate its kind. pub(super) fn one_mma_fact( facts: &OrdinaryFacts, @@ -473,7 +479,7 @@ pub(super) fn verify_mma_use_closure( for (value, paths) in &facts.facts { for (path, found) in paths { for fact in found.iter().copied() { - if matches!(fact, MmaFact::Scalar(_) | MmaFact::Unknown(_)) { + if fact.is_opaque() { continue; } for r#use in value.uses(ctx) { diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs index 4b6cb8318f..d4951a504a 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs @@ -21,7 +21,7 @@ use super::{ ghosts::pipeline_semantic_ids, invalid, provenance::{ - OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, + CarrierFact, OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, seed_unknown_integer_carriers, }, }; @@ -175,6 +175,12 @@ enum PipelineFact { Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), } +impl CarrierFact for PipelineFact { + fn is_opaque(&self) -> bool { + matches!(self, Self::Unknown(_)) + } +} + /// Require an unambiguous slot or phase fact rooted at cursor creation. fn pipeline_state_root( facts: &OrdinaryFacts, diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs index e682e76803..12fcdde90c 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs @@ -35,6 +35,19 @@ use pliron::{ value::Value, }; +/// One possible origin of an ordinary carrier. +pub(super) trait CarrierFact: Clone + Eq + std::hash::Hash { + /// True for a fact that names an ordinary SSA value instead of a semantic + /// origin: an input without provenance, or an ordinary scalar. + fn is_opaque(&self) -> bool; +} + +/// Consumers read an origin set only through its size, its sole origin, and +/// whether it holds a given semantic origin. Two opaque origins therefore show +/// a merge exactly as any larger number does, and keeping more would let the +/// sets grow with every merge they flow through. +const OPAQUE_FACTS_KEPT: usize = 2; + /// Monotone sets of possible semantic origins, indexed by carrier and field path. #[derive(Clone)] pub(super) struct OrdinaryFacts { @@ -54,18 +67,21 @@ impl Default for OrdinaryFacts { } } -impl OrdinaryFacts -where - F: Clone + Eq + std::hash::Hash, -{ +impl OrdinaryFacts { /// Add one origin to a field path and report whether the fixed-point state changed. pub(super) fn insert_path(&mut self, value: Value, path: Vec, fact: F) -> bool { - self.facts + let found = self + .facts .entry(value) .or_default() .entry(path) - .or_default() - .insert(fact) + .or_default(); + if fact.is_opaque() + && found.iter().filter(|kept| kept.is_opaque()).count() >= OPAQUE_FACTS_KEPT + { + return false; + } + found.insert(fact) } /// Add an origin for the whole SSA carrier, with an empty field path. @@ -273,7 +289,7 @@ pub(super) fn propagate_ordinary_plumbing( facts: &mut OrdinaryFacts, ) -> bool where - F: Clone + Eq + std::hash::Hash, + F: CarrierFact, { let mut changed = false; let mut blocks = HashSet::new(); @@ -427,7 +443,7 @@ pub(super) fn seed_unknown_integer_carriers( width: u32, unknown: impl Fn(Value) -> F, ) where - F: Clone + Eq + std::hash::Hash, + F: CarrierFact, { let mut values = Vec::new(); let mut blocks = HashSet::new(); @@ -465,7 +481,7 @@ pub(super) fn poison_missing_forwarding_sources( unknown: impl Fn(Value) -> F + Copy, ) -> bool where - F: Clone + Eq + std::hash::Hash, + F: CarrierFact, { let mut changed = false; let mut blocks = HashSet::new(); diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs index e34b57e919..d65437e580 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs @@ -23,7 +23,7 @@ use super::{ ghosts::scheduler_semantic_ids, invalid, provenance::{ - OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, + CarrierFact, OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, seed_unknown_integer_carriers, }, }; @@ -74,6 +74,12 @@ pub(super) enum SchedulerFact { Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), } +impl CarrierFact for SchedulerFact { + fn is_opaque(&self) -> bool { + matches!(self, Self::Unknown(_)) + } +} + /// Require exactly one current-index or stride origin of the requested kind. pub(super) fn scheduler_root( facts: &OrdinaryFacts, diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs index 18acd059c9..6f1298d802 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs @@ -34,8 +34,8 @@ use super::{ mma::{build_mma_facts, require_smem_overlay_pair, verify_mma_use_closure}, pipeline::require_producer_pipeline_order, provenance::{ - OrdinaryFacts, is_local_cell_pointer_alias_cast, is_semantic_pointer_carrier_cast, - local_cell_is_closed, + CarrierFact, OrdinaryFacts, is_local_cell_pointer_alias_cast, + is_semantic_pointer_carrier_cast, local_cell_is_closed, }, scheduler::{build_scheduler_facts, scheduler_root}, *, @@ -190,12 +190,15 @@ fn cond_branch( } /// Construct one valid persistent-scheduler function with configurable body/exit placement. +/// The scheduler state reaches the loop header through `pad` ordinary blocks, each passing it +/// as block arguments and through a closed local cell. fn append_complete_scheduler_function( ctx: &mut Context, module_block: Ptr, name: &str, grid: CuteTileGridAttr, misroute_work_edge: bool, + pad: usize, ) { let function_type = FunctionType::get(ctx, vec![], vec![]); let function = Operation::new( @@ -223,12 +226,20 @@ fn append_complete_scheduler_function( let start = CuteSchedulerNew1dOp::new(ctx, grid).get_operation(); start.insert_at_back(entry, ctx); let start = CuteSchedulerNew1dOp::wrap(start); - goto( - ctx, - entry, - header, - vec![start.current(ctx), start.stride(ctx)], - ); + let mut carried = vec![start.current(ctx), start.stride(ctx)]; + let mut previous = entry; + for _ in 0..pad { + let block = BasicBlock::new(ctx, None, vec![u64_ty, u64_ty]); + block.insert_at_back(region, ctx); + goto(ctx, previous, block, carried); + let current = block.deref(ctx).get_argument(0); + let stride = block.deref(ctx).get_argument(1); + let cell = local_cell(ctx, block, u64_ty); + store(ctx, block, cell, current); + carried = vec![load(ctx, block, cell, u64_ty), stride]; + previous = block; + } + goto(ctx, previous, header, carried); let current = header.deref(ctx).get_argument(0); let stride = header.deref(ctx).get_argument(1); @@ -306,6 +317,7 @@ fn accepts_independent_scheduler_stories_in_two_functions_without_mutation() { "first", CuteTileGridAttr::new(4, 4, 1), false, + 0, ); append_complete_scheduler_function( &mut ctx, @@ -313,6 +325,7 @@ fn accepts_independent_scheduler_stories_in_two_functions_without_mutation() { "second", CuteTileGridAttr::new(5, 4, 1), false, + 0, ); let mut before = Vec::new(); collect_ops(&ctx, module, &mut before); @@ -335,6 +348,7 @@ fn rejects_a_scheduler_body_on_the_exit_edge_even_when_preorder_looks_ordered() "misrouted", CuteTileGridAttr::new(4, 4, 1), true, + 0, ); let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); @@ -433,6 +447,111 @@ fn ordinary_facts_forward_a_snapshot_of_one_carrier() { assert!(!facts.copy_entries(source, destination, &[1], &[])); assert_eq!(sorted(&facts, destination), vec![(vec![2], 2)]); assert_eq!(sorted(&facts, other), vec![(vec![], 3)]); + + // Opaque origins stop at two per path; semantic origins are always kept. + assert!(facts.insert(other, 100)); + assert!(facts.insert(other, 101)); + assert!(!facts.insert(other, 102)); + assert!(!facts.insert(other, 100)); + assert!(facts.insert(other, 4)); + assert_eq!( + sorted(&facts, other), + vec![(vec![], 3), (vec![], 4), (vec![], 100), (vec![], 101)] + ); + assert!(facts.insert(destination, 102)); + assert!(facts.copy_entries(destination, other, &[], &[])); + assert_eq!( + sorted(&facts, other), + vec![ + (vec![], 3), + (vec![], 4), + (vec![], 100), + (vec![], 101), + (vec![2], 2) + ] + ); +} + +/// Origins at or above 100 stand for opaque facts in the fact-set tests. +impl CarrierFact for u32 { + fn is_opaque(&self) -> bool { + *self >= 100 + } +} + +/// A scheduler cursor that passes through many blocks and cells keeps its one origin. +#[test] +fn accepts_scheduler_state_carried_through_many_blocks_and_cells() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "padded", + CuteTileGridAttr::new(4, 4, 1), + false, + 48, + ); + verify_cute_semantics(&ctx, module).unwrap(); +} + +/// Build a function whose `cute.scheduler_has_work` reads a cursor with no scheduler origin: +/// a plain value, or a block argument merging `incoming` unrelated values. +fn scheduler_has_work_without_origin(incoming: usize) -> String { + let mut ctx = Context::new(); + let (module, module_block) = module_top(&mut ctx); + let function_type = FunctionType::get(&ctx, vec![], vec![]); + let function = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function_op = MirFuncOp::new(&mut ctx, function, TypeAttr::new(function_type.into())); + function_op.set_symbol_name(&mut ctx, "unrooted".try_into().unwrap()); + function.insert_at_back(module_block, &ctx); + let region = function.deref(&ctx).get_region(0); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + let mut cursor = undef(&mut ctx, entry, u64_ty); + let mut current = entry; + if incoming > 0 { + let join = BasicBlock::new(&mut ctx, None, vec![u64_ty]); + join.insert_at_back(region, &ctx); + goto(&mut ctx, entry, join, vec![cursor]); + for _ in 1..incoming { + let side = BasicBlock::new(&mut ctx, None, vec![]); + side.insert_at_back(region, &ctx); + let value = undef(&mut ctx, side, u64_ty); + goto(&mut ctx, side, join, vec![value]); + } + cursor = join.deref(&ctx).get_argument(0); + current = join; + } + let grid = CuteTileGridAttr::new(4, 4, 1); + CuteSchedulerHasWorkOp::new(&mut ctx, cursor, grid) + .get_operation() + .insert_at_back(current, &ctx); + verify_cute_semantics(&ctx, module).unwrap_err().to_string() +} + +/// Rejections keep their wording however many unrelated values a cursor merges. +#[test] +fn scheduler_cursor_without_origin_is_rejected_with_the_same_wording() { + let prefix = "invalid CuTe semantic module: cute.scheduler_has_work current"; + assert_eq!( + scheduler_has_work_without_origin(0), + format!("{prefix} is not derived from the required scheduler_new_1d result") + ); + for incoming in [1, 3, 8] { + assert_eq!( + scheduler_has_work_without_origin(incoming), + format!("{prefix} merges values from different scheduler stories") + ); + } } /// Casts that change pointee, address space, or write capability cannot establish closed-cell aliases. From 4bb62e04db459718bb2424a2dbd5a62d9aec3c7a Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:28:05 +0900 Subject: [PATCH 4/4] perf(cute): check each local cell's aliases once per analysis Every load and store that forwards provenance through a closed local cell called `local_cell_is_closed`, which walks all uses of the cell and its aliases. A cell with N accesses therefore cost N walks of N uses on every fixed-point sweep, so verification time grew with the square of the number of accesses to one cell. The verifier does not change the IR it reads, so a cell's closedness is fixed for the whole analysis. `LocalCells` walks the alias graph of each `mir.alloca` in the scope once, and the scheduler, pipeline and shared-MMA fact builders, the shared-MMA use closure and the epilogue plumbing check look cells up there. A cell allocated outside the scope is still walked on demand, so every answer is the one the walk gives. The ordinary plumbing and the poison pass also read each incoming edge's successor operands once per block instead of once per block argument. A new test compares the table with the alias walk for closed, aliased and escaped cells, with and without their allocations in the scope. Signed-off-by: juliankang4 <200967947+juliankang4@users.noreply.github.com> --- .../dialect-cute/src/verify/epilogue.rs | 30 +++++--- .../dialect-cute/src/verify/mma.rs | 37 +++++---- .../dialect-cute/src/verify/pipeline.rs | 18 +++-- .../dialect-cute/src/verify/provenance.rs | 77 +++++++++++++------ .../dialect-cute/src/verify/scheduler.rs | 18 +++-- .../dialect-cute/src/verify/tests.rs | 40 +++++++++- 6 files changed, 157 insertions(+), 63 deletions(-) diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs index 66aae6d026..fe2794e177 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs @@ -29,7 +29,7 @@ use super::{ ghosts::epilogue_semantic_ids, invalid, mma::{MmaFact, build_mma_facts, one_mma_fact}, - provenance::{is_semantic_pointer_carrier_cast, local_cell, local_cell_is_closed}, + provenance::{LocalCells, is_semantic_pointer_carrier_cast}, tma::direct_tma_view_producer, }; use crate::{ @@ -49,21 +49,25 @@ fn is_epilogue_pointer_identity_cast(ctx: &Context, operation: Ptr) - } /// Recognize carrier-only operations allowed between writer synchronization anchors. -fn is_epilogue_writer_plumbing(ctx: &Context, operation: Ptr) -> bool { +fn is_epilogue_writer_plumbing( + ctx: &Context, + cells: &LocalCells, + operation: Ptr, +) -> bool { let opid = Operation::get_opid(operation, ctx); opid == MirGotoOp::get_opid_static() || opid == MirStorageLiveOp::get_opid_static() || opid == MirStorageDeadOp::get_opid_static() || opid == MirExtractFieldOp::get_opid_static() || is_epilogue_pointer_identity_cast(ctx, operation) - || is_closed_local_cell_load(ctx, operation) + || is_closed_local_cell_load(ctx, cells, operation) } /// A non-volatile `mir.load` from a closed compiler-created local cell only /// re-materializes carrier state (e.g. the epilogue handle struct spilled to /// an alloca slot); it cannot touch tile memory, so it is admissible plumbing /// between the writer-protocol anchor operations. -fn is_closed_local_cell_load(ctx: &Context, operation: Ptr) -> bool { +fn is_closed_local_cell_load(ctx: &Context, cells: &LocalCells, operation: Ptr) -> bool { if Operation::get_opid(operation, ctx) != MirLoadOp::get_opid_static() { return false; } @@ -77,10 +81,7 @@ fn is_closed_local_cell_load(ctx: &Context, operation: Ptr) -> bool { if load.is_volatile(ctx) { return false; } - let Some((root, _path)) = local_cell(ctx, load.address_opd(ctx)) else { - return false; - }; - local_cell_is_closed(ctx, root) + cells.closed_cell(ctx, load.address_opd(ctx)).is_some() } /// Accept only the dead immutable reference Rust leaves around a no-inline @@ -149,9 +150,13 @@ fn is_dead_epilogue_store_pipeline_ref(ctx: &Context, operation: Ptr) } /// Recognize harmless carrier assembly and indexing between TMA issuer anchors. -fn is_epilogue_issuer_plumbing(ctx: &Context, operation: Ptr) -> bool { +fn is_epilogue_issuer_plumbing( + ctx: &Context, + cells: &LocalCells, + operation: Ptr, +) -> bool { let opid = Operation::get_opid(operation, ctx); - is_epilogue_writer_plumbing(ctx, operation) + is_epilogue_writer_plumbing(ctx, cells, operation) || opid == MirUndefOp::get_opid_static() || opid == MirInsertFieldOp::get_opid_static() || opid == MirConstructStructOp::get_opid_static() @@ -528,6 +533,7 @@ pub(super) fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr]) - )?; } + let cells = LocalCells::new(ctx, all_ops); let writer_range = require_linear_protocol_order( ctx, "epilogue writer protocol", @@ -535,7 +541,7 @@ pub(super) fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr]) - )?; for operation in &writer_range { let expected = [reusable_sync, store_fragment, ready_sync].contains(operation); - if !expected && !is_epilogue_writer_plumbing(ctx, *operation) { + if !expected && !is_epilogue_writer_plumbing(ctx, &cells, *operation) { return Err(invalid(format!( "epilogue writer protocol has interposed non-carrier operation `{}`", Operation::get_opid(*operation, ctx) @@ -663,7 +669,7 @@ pub(super) fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr]) - commit, ] .contains(&operation); - if !allowed_tma_view && !expected && !is_epilogue_issuer_plumbing(ctx, operation) { + if !allowed_tma_view && !expected && !is_epilogue_issuer_plumbing(ctx, &cells, operation) { return Err(invalid(format!( "epilogue TMA issuer protocol has interposed operation `{opid}`" ))); diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs index c8deda70a6..f154312134 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs @@ -24,8 +24,8 @@ use super::{ ghosts::smem_mma_semantic_ids, invalid, provenance::{ - CarrierFact, OrdinaryFacts, branch_successor_operands, is_semantic_pointer_carrier_cast, - local_cell, local_cell_is_closed, poison_missing_forwarding_sources, + CarrierFact, LocalCells, OrdinaryFacts, branch_successor_operands, + is_semantic_pointer_carrier_cast, poison_missing_forwarding_sources, propagate_ordinary_plumbing, seed_unknown_integer_carriers, }, }; @@ -198,6 +198,7 @@ pub(super) fn build_mma_facts( ctx: &Context, operations: &[Ptr], ) -> OrdinaryFacts { + let cells = LocalCells::new(ctx, operations); let mut facts = OrdinaryFacts::default(); for operation in operations { let opid = Operation::get_opid(*operation, ctx); @@ -240,7 +241,7 @@ pub(super) fn build_mma_facts( } loop { - let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + let mut changed = propagate_ordinary_plumbing(ctx, operations, &cells, &mut facts); for operation in operations { if Operation::get_opid(*operation, ctx) != CuteTiledGemmOp::get_opid_static() { continue; @@ -266,9 +267,14 @@ pub(super) fn build_mma_facts( seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, MmaFact::Scalar); seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, MmaFact::Scalar); loop { - let poisoned = - poison_missing_forwarding_sources(ctx, operations, &mut facts, MmaFact::Unknown); - let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + let poisoned = poison_missing_forwarding_sources( + ctx, + operations, + &cells, + &mut facts, + MmaFact::Unknown, + ); + let propagated = propagate_ordinary_plumbing(ctx, operations, &cells, &mut facts); if !poisoned && !propagated { break; } @@ -364,8 +370,10 @@ fn mma_branch_forwards_fact( } /// Check that ordinary casts, aggregates, or closed local cells forward the tracked fact without escape. +#[allow(clippy::too_many_arguments)] fn mma_ordinary_use_forwards_fact( ctx: &Context, + cells: &LocalCells, facts: &OrdinaryFacts, value: Value, path: &[u32], @@ -425,12 +433,11 @@ fn mma_ordinary_use_forwards_fact( if opid == MirStoreOp::get_opid_static() && operation.get_num_operands() == 2 { let store = MirStoreOp::new(user); if operand_index == 0 { - return local_cell(ctx, store.address_opd(ctx)) - .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + return cells.closed_cell(ctx, store.address_opd(ctx)).is_some(); } if operand_index == 1 - && let Some((root, mut destination_path)) = local_cell(ctx, store.address_opd(ctx)) - && local_cell_is_closed(ctx, root) + && let Some((root, mut destination_path)) = + cells.closed_cell(ctx, store.address_opd(ctx)) { destination_path.extend_from_slice(path); return mma_fact_at_path(facts, root, &destination_path, fact); @@ -442,12 +449,9 @@ fn mma_ordinary_use_forwards_fact( && operation.get_num_results() == 1 { let load = MirLoadOp::new(user); - let Some((root, cell_path)) = local_cell(ctx, load.address_opd(ctx)) else { + let Some((_, cell_path)) = cells.closed_cell(ctx, load.address_opd(ctx)) else { return false; }; - if !local_cell_is_closed(ctx, root) { - return false; - } let Some(result_path) = path.strip_prefix(cell_path.as_slice()) else { return true; }; @@ -460,8 +464,7 @@ fn mma_ordinary_use_forwards_fact( .get_attr_field_index(ctx) .is_some() { - return local_cell(ctx, operation.get_result(0)) - .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + return cells.closed_cell(ctx, operation.get_result(0)).is_some(); } if opid == MirCondBranchOp::get_opid_static() && operand_index == 0 { return false; @@ -476,6 +479,7 @@ pub(super) fn verify_mma_use_closure( facts: &OrdinaryFacts, ) -> Result { let scope: HashSet> = operations.iter().copied().collect(); + let cells = LocalCells::new(ctx, operations); for (value, paths) in &facts.facts { for (path, found) in paths { for fact in found.iter().copied() { @@ -489,6 +493,7 @@ pub(super) fn verify_mma_use_closure( && (mma_semantic_consumer(ctx, user, operand_index, fact) || mma_ordinary_use_forwards_fact( ctx, + &cells, facts, *value, path, diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs index d4951a504a..3787024c30 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs @@ -21,8 +21,8 @@ use super::{ ghosts::pipeline_semantic_ids, invalid, provenance::{ - CarrierFact, OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, - seed_unknown_integer_carriers, + CarrierFact, LocalCells, OrdinaryFacts, poison_missing_forwarding_sources, + propagate_ordinary_plumbing, seed_unknown_integer_carriers, }, }; use crate::{ @@ -212,6 +212,7 @@ fn build_pipeline_facts( ctx: &Context, operations: &[Ptr], ) -> OrdinaryFacts { + let cells = LocalCells::new(ctx, operations); let mut facts = OrdinaryFacts::default(); for operation in operations { if Operation::get_opid(*operation, ctx) == CutePipelineStateNewOp::get_opid_static() { @@ -222,7 +223,7 @@ fn build_pipeline_facts( } loop { - let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + let mut changed = propagate_ordinary_plumbing(ctx, operations, &cells, &mut facts); for operation in operations { if Operation::get_opid(*operation, ctx) != CutePipelineStateAdvanceOp::get_opid_static() { @@ -254,9 +255,14 @@ fn build_pipeline_facts( seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, PipelineFact::Unknown); loop { - let poisoned = - poison_missing_forwarding_sources(ctx, operations, &mut facts, PipelineFact::Unknown); - let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + let poisoned = poison_missing_forwarding_sources( + ctx, + operations, + &cells, + &mut facts, + PipelineFact::Unknown, + ); + let propagated = propagate_ordinary_plumbing(ctx, operations, &cells, &mut facts); if !poisoned && !propagated { break; } diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs index 12fcdde90c..3732e626c4 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs @@ -27,6 +27,7 @@ use dialect_mir::{ types::{MirArrayType, MirPtrType, MirStructType, MirTupleType, address_space}, }; use pliron::{ + basic_block::BasicBlock, builtin::{op_interfaces::BranchOpInterface, types::IntegerType}, context::{Context, Ptr}, op::{Op, op_cast}, @@ -278,6 +279,51 @@ pub(super) fn local_cell_is_closed(ctx: &Context, root: Value) -> bool { true } +/// Closedness of every local cell allocated in one scope. The verifier never +/// mutates the IR it reads, so each cell's alias graph is walked once and +/// shared by all accesses and fixed-point iterations. +pub(super) struct LocalCells { + closed: HashMap, +} + +impl LocalCells { + /// Check the alias graph of each `mir.alloca` result in the scope. + pub(super) fn new(ctx: &Context, operations: &[Ptr]) -> Self { + let closed = operations + .iter() + .filter(|operation| { + Operation::get_opid(**operation, ctx) == MirAllocaOp::get_opid_static() + && operation.deref(ctx).get_num_results() == 1 + }) + .map(|operation| { + let root = operation.deref(ctx).get_result(0); + (root, local_cell_is_closed(ctx, root)) + }) + .collect(); + Self { closed } + } + + /// Resolve a pointer to a closed local cell and the field path inside it. + pub(super) fn closed_cell(&self, ctx: &Context, pointer: Value) -> Option<(Value, Vec)> { + let (root, path) = local_cell(ctx, pointer)?; + let closed = self + .closed + .get(&root) + .copied() + .unwrap_or_else(|| local_cell_is_closed(ctx, root)); + closed.then_some((root, path)) + } +} + +/// Read the operands each incoming edge passes to a block, once per edge. +fn incoming_edge_operands(ctx: &Context, block: Ptr) -> Vec> { + block + .uses(ctx) + .into_iter() + .filter_map(|edge| branch_successor_operands(ctx, edge.user_op(), edge.find_index(ctx))) + .collect() +} + /// Propagate semantic facts through the exact ordinary SSA adapters that can /// remain after common CuTe preparation. This is deliberately a monotone /// fixed-point transfer: loop block arguments become known once their entry @@ -286,6 +332,7 @@ pub(super) fn local_cell_is_closed(ctx: &Context, root: Value) -> bool { pub(super) fn propagate_ordinary_plumbing( ctx: &Context, operations: &[Ptr], + cells: &LocalCells, facts: &mut OrdinaryFacts, ) -> bool where @@ -350,8 +397,7 @@ where } else if opid == MirStoreOp::get_opid_static() && operation_ref.get_num_operands() == 2 { let store = MirStoreOp::new(*operation); if !store.is_volatile(ctx) - && let Some((root, path)) = local_cell(ctx, store.address_opd(ctx)) - && local_cell_is_closed(ctx, root) + && let Some((root, path)) = cells.closed_cell(ctx, store.address_opd(ctx)) { changed |= facts.copy_entries(store.value_opd(ctx), root, &[], &path); } @@ -361,8 +407,7 @@ where { let load = MirLoadOp::new(*operation); if !load.is_volatile(ctx) - && let Some((root, path)) = local_cell(ctx, load.address_opd(ctx)) - && local_cell_is_closed(ctx, root) + && let Some((root, path)) = cells.closed_cell(ctx, load.address_opd(ctx)) { changed |= facts.copy_entries(root, operation_ref.get_result(0), &path, &[]); } @@ -370,16 +415,11 @@ where } for block in blocks { + let edges = incoming_edge_operands(ctx, block); let block_ref = block.deref(ctx); for argument_index in 0..block_ref.get_num_arguments() { let argument = block_ref.get_argument(argument_index); - for edge in block.uses(ctx) { - let terminator = edge.user_op(); - let Some(operands) = - branch_successor_operands(ctx, terminator, edge.find_index(ctx)) - else { - continue; - }; + for operands in &edges { if let Some(incoming) = operands.get(argument_index).copied() { changed |= facts.copy_entries(incoming, argument, &[], &[]); } @@ -477,6 +517,7 @@ pub(super) fn seed_unknown_integer_carriers( pub(super) fn poison_missing_forwarding_sources( ctx: &Context, operations: &[Ptr], + cells: &LocalCells, facts: &mut OrdinaryFacts, unknown: impl Fn(Value) -> F + Copy, ) -> bool @@ -493,12 +534,9 @@ where continue; } let store = MirStoreOp::new(*operation); - let Some((root, cell_path)) = local_cell(ctx, store.address_opd(ctx)) else { + let Some((root, cell_path)) = cells.closed_cell(ctx, store.address_opd(ctx)) else { continue; }; - if !local_cell_is_closed(ctx, root) { - continue; - } for (path, _) in facts.entries(root) { let Some(source_path) = path.strip_prefix(cell_path.as_slice()) else { continue; @@ -514,6 +552,7 @@ where } for block in blocks { + let edges = incoming_edge_operands(ctx, block); let block_ref = block.deref(ctx); for argument_index in 0..block_ref.get_num_arguments() { let argument = block_ref.get_argument(argument_index); @@ -521,13 +560,7 @@ where if argument_entries.is_empty() { continue; } - for edge in block.uses(ctx) { - let terminator = edge.user_op(); - let Some(operands) = - branch_successor_operands(ctx, terminator, edge.find_index(ctx)) - else { - continue; - }; + for operands in &edges { let Some(incoming) = operands.get(argument_index).copied() else { continue; }; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs index d65437e580..205e60a15d 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs @@ -23,8 +23,8 @@ use super::{ ghosts::scheduler_semantic_ids, invalid, provenance::{ - CarrierFact, OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, - seed_unknown_integer_carriers, + CarrierFact, LocalCells, OrdinaryFacts, poison_missing_forwarding_sources, + propagate_ordinary_plumbing, seed_unknown_integer_carriers, }, }; use crate::{ @@ -111,6 +111,7 @@ pub(super) fn build_scheduler_facts( ctx: &Context, operations: &[Ptr], ) -> OrdinaryFacts { + let cells = LocalCells::new(ctx, operations); let mut facts = OrdinaryFacts::default(); for operation in operations { if Operation::get_opid(*operation, ctx) == CuteSchedulerNew1dOp::get_opid_static() { @@ -121,7 +122,7 @@ pub(super) fn build_scheduler_facts( } loop { - let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + let mut changed = propagate_ordinary_plumbing(ctx, operations, &cells, &mut facts); for operation in operations { if Operation::get_opid(*operation, ctx) != CuteSchedulerAdvanceOp::get_opid_static() { continue; @@ -152,9 +153,14 @@ pub(super) fn build_scheduler_facts( // itself first, while still poisoning a merge with an arbitrary edge. seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, SchedulerFact::Unknown); loop { - let poisoned = - poison_missing_forwarding_sources(ctx, operations, &mut facts, SchedulerFact::Unknown); - let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + let poisoned = poison_missing_forwarding_sources( + ctx, + operations, + &cells, + &mut facts, + SchedulerFact::Unknown, + ); + let propagated = propagate_ordinary_plumbing(ctx, operations, &cells, &mut facts); if !poisoned && !propagated { break; } diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs index 6f1298d802..4df2c3f17e 100644 --- a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs @@ -34,7 +34,7 @@ use super::{ mma::{build_mma_facts, require_smem_overlay_pair, verify_mma_use_closure}, pipeline::require_producer_pipeline_order, provenance::{ - CarrierFact, OrdinaryFacts, is_local_cell_pointer_alias_cast, + CarrierFact, LocalCells, OrdinaryFacts, is_local_cell_pointer_alias_cast, is_semantic_pointer_carrier_cast, local_cell_is_closed, }, scheduler::{build_scheduler_facts, scheduler_root}, @@ -403,6 +403,44 @@ fn factless_block_argument_store_poisons_a_valid_scheduler_cell() { assert!(error.contains("merges values"), "{error}"); } +/// A scope's cell table answers every pointer as walking that pointer's cell aliases does, +/// including cells allocated outside the scope. +#[test] +fn local_cells_agree_with_the_alias_walk() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let closed = local_cell(&mut ctx, block, u64_ty); + let read_only_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, false).into(); + let alias_cast = pointer_cast(&mut ctx, block, closed, read_only_type); + let alias = alias_cast.deref(&ctx).get_result(0); + let value = undef(&mut ctx, block, u64_ty); + store(&mut ctx, block, closed, value); + load(&mut ctx, block, alias, u64_ty); + let escaped = local_cell(&mut ctx, block, u64_ty); + let pointer_type = escaped.get_type(&ctx); + let holder = local_cell(&mut ctx, block, pointer_type); + store(&mut ctx, block, holder, escaped); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + let without_allocas: Vec<_> = operations + .iter() + .copied() + .filter(|operation| Operation::get_opid(*operation, &ctx) != MirAllocaOp::get_opid_static()) + .collect(); + + for scope in [&operations, &without_allocas] { + let cells = LocalCells::new(&ctx, scope); + for pointer in [closed, alias, escaped, holder] { + let walked = super::provenance::local_cell(&ctx, pointer) + .filter(|(root, _)| local_cell_is_closed(&ctx, *root)); + assert_eq!(cells.closed_cell(&ctx, pointer), walked); + } + assert!(cells.closed_cell(&ctx, alias).is_some()); + assert!(cells.closed_cell(&ctx, escaped).is_none()); + } +} + /// Fact lookups see only the named carrier, and a copy forwards the source as it was when the copy began. #[test] fn ordinary_facts_forward_a_snapshot_of_one_carrier() {