perf(cute): speed up the CuTe verifier's provenance analysis

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<Operation>], 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,

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<Operation>) -
 }
 
 /// Recognize carrier-only operations allowed between writer synchronization anchors.
-fn is_epilogue_writer_plumbing(ctx: &Context, operation: Ptr<Operation>) -> bool {
+fn is_epilogue_writer_plumbing(
+    ctx: &Context,
+    cells: &LocalCells,
+    operation: Ptr<Operation>,
+) -> 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<Operation>) -> bool {
+fn is_closed_local_cell_load(ctx: &Context, cells: &LocalCells, operation: Ptr<Operation>) -> 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<Operation>) -> 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<Operation>)
 }
 
 /// Recognize harmless carrier assembly and indexing between TMA issuer anchors.
-fn is_epilogue_issuer_plumbing(ctx: &Context, operation: Ptr<Operation>) -> bool {
+fn is_epilogue_issuer_plumbing(
+    ctx: &Context,
+    cells: &LocalCells,
+    operation: Ptr<Operation>,
+) -> 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,14 +533,15 @@ pub(super) fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr<Operation>]) -
         )?;
     }
 
+    let cells = LocalCells::new(ctx, all_ops);
     let writer_range = require_linear_protocol_order(
         ctx,
         "epilogue writer protocol",
         &[reusable_sync, store_fragment, ready_sync],
     )?;
     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<Operation>]) -
             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}`"
             )));

cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs

@@ -20,11 +20,12 @@ use pliron::{
 
 use super::{
     Result,
+    cfg::scope_has_any,
     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,
+        CarrierFact, LocalCells, OrdinaryFacts, branch_successor_operands,
+        is_semantic_pointer_carrier_cast, poison_missing_forwarding_sources,
         propagate_ordinary_plumbing, seed_unknown_integer_carriers,
     },
 };
@@ -101,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<MmaFact>,
@@ -191,6 +198,7 @@ pub(super) fn build_mma_facts(
     ctx: &Context,
     operations: &[Ptr<Operation>],
 ) -> OrdinaryFacts<MmaFact> {
+    let cells = LocalCells::new(ctx, operations);
     let mut facts = OrdinaryFacts::default();
     for operation in operations {
         let opid = Operation::get_opid(*operation, ctx);
@@ -233,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;
@@ -259,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;
         }
@@ -357,10 +370,13 @@ 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<MmaFact>,
-    key: &CarrierKey,
+    value: Value,
+    path: &[u32],
     fact: MmaFact,
     user: Ptr<Operation>,
     operand_index: usize,
@@ -372,14 +388,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 +406,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,22 +425,21 @@ 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);
         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(&key.path);
+            destination_path.extend_from_slice(path);
             return mma_fact_at_path(facts, root, &destination_path, fact);
         }
         return false;
@@ -434,13 +449,10 @@ 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) = 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);
@@ -452,13 +464,12 @@ 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;
     }
-    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 +478,37 @@ pub(super) fn verify_mma_use_closure(
     operations: &[Ptr<Operation>],
     facts: &OrdinaryFacts<MmaFact>,
 ) -> 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<Ptr<Operation>> = 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() {
+                if fact.is_opaque() {
                     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,
+                                &cells,
+                                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)
+                    )));
+                }
             }
         }
     }
@@ -518,6 +535,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<Operation>]) -> 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();