diff --git a/src/frontend/multi_pod/cost.rs b/src/frontend/multi_pod/cost.rs index 0bfaadb1..a73e593b 100644 --- a/src/frontend/multi_pod/cost.rs +++ b/src/frontend/multi_pod/cost.rs @@ -46,6 +46,9 @@ impl From<&CustomPredicateRef> for CustomPredicateId { /// op kind. #[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct OperationCost { + /// True for a statement copied directly from an external input POD. + #[serde(default)] + pub is_external_opening: bool, /// `Contains`-family proofs on a shallow tree (depth <= /// `max_depth_small`). May occupy either a small or medium state slot. pub merkle_proofs_small: usize, @@ -247,6 +250,9 @@ impl OperationCost { NativeOperation::PublicKeyFromEntries => { cost.public_key = 1; } + NativeOperation::OpenInputStatement => { + cost.is_external_opening = true; + } // Zero-cost ops (no per-POD resource consumed). NativeOperation::None | NativeOperation::EqualFromEntries @@ -259,8 +265,6 @@ impl OperationCost { | NativeOperation::ProductFromEntries | NativeOperation::MaxFromEntries | NativeOperation::HashFromEntries - // Tracked separately by the partitioner. - | NativeOperation::OpenInputStatement // Syntactic sugar variants (lowered before proving). | NativeOperation::GtEqFromEntries | NativeOperation::GtFromEntries @@ -361,6 +365,16 @@ mod tests { assert_eq!(pk.signed_by, 0); } + #[test] + fn open_input_is_marked_for_early_scheduling() { + let cost = OperationCost::from_operation( + &native_op(NativeOperation::OpenInputStatement, OperationAux::None), + &Params::default(), + ); + + assert!(cost.is_external_opening); + } + /// Contains on a tree at depth `<= max_depth_small` is small-eligible; /// at depth `> max_depth_small` it must use a medium slot. #[test] diff --git a/src/frontend/multi_pod/deps.rs b/src/frontend/multi_pod/deps.rs index ece2d775..f32a2a5f 100644 --- a/src/frontend/multi_pod/deps.rs +++ b/src/frontend/multi_pod/deps.rs @@ -4,16 +4,14 @@ use std::collections::HashMap; use crate::{ frontend::{Operation, OperationArg}, - middleware::{ - Hash, InputPodOpenStatement, NativeOperation, OperationAux, OperationType, Statement, - }, + middleware::{InputPodOpenStatement, NativeOperation, OperationAux, OperationType, Statement}, }; /// Reference to a statement sourced from an external input POD. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct ExternalDependency { - /// Hash of the external POD containing `statement` in its public set. - pub pod_hash: Hash, + /// Position of the concrete source POD in the builder's input list. + pub pod_index: usize, /// The statement value itself. pub statement: Statement, } @@ -36,14 +34,13 @@ pub struct DependencyGraph { } impl DependencyGraph { - /// Build the dependency graph from parallel `statements` and - /// `operations` arrays (where `operations[i]` produces `statements[i]`) - /// plus a `statement -> pod hash` map for recognising external - /// references. + /// Build dependencies from aligned statements and operations, where + /// `operations[i]` produces `statements[i]`, while retaining each external + /// reference's source POD. pub fn build( statements: &[Statement], operations: &[Operation], - external_pod_statements: &HashMap, + external_pod_statements: &HashMap, ) -> Self { let mut statement_deps = Vec::with_capacity(statements.len()); @@ -85,9 +82,9 @@ impl DependencyGraph { // it. Fall through to the external lookup below. } - if let Some(&pod_hash) = external_pod_statements.get(dep_stmt) { + if let Some(&pod_index) = external_pod_statements.get(dep_stmt) { deps.push(StatementSource::External(ExternalDependency { - pod_hash, + pod_index, statement: dep_stmt.clone(), })); } else { @@ -105,11 +102,11 @@ impl DependencyGraph { // accounts for the input-tree slot and the external-pod reference. if let OperationType::Native(NativeOperation::OpenInputStatement) = op.0 { if let OperationAux::OpenInputStatement(InputPodOpenStatement { - sts_root, .. + pod_index, .. }) = &op.2 { deps.push(StatementSource::External(ExternalDependency { - pod_hash: *sts_root, + pod_index: *pod_index, statement: statements[idx].clone(), })); } diff --git a/src/frontend/multi_pod/diagnostics.rs b/src/frontend/multi_pod/diagnostics.rs index 2dc3f506..35d7c13a 100644 --- a/src/frontend/multi_pod/diagnostics.rs +++ b/src/frontend/multi_pod/diagnostics.rs @@ -16,7 +16,8 @@ use std::{ use super::{ cost::{OperationCost, ResourceTotals}, - shape::{AbstractDep, InputShape, OutputShape}, + partition::SegmentChecker, + shape::{input_pod_slots, AbstractDep, InputShape, OutputShape}, }; use crate::middleware::Params; @@ -198,7 +199,8 @@ pub struct PodUtilization { pub num_statements: usize, pub resources: Vec, pub imports: UtilizationRow, - pub external_pods: UtilizationRow, + /// Number of input-POD slots occupied. See [`input_pod_slots`]. + pub input_pod_slots: UtilizationRow, /// Statements this POD contributes to a Merkle tree. For intermediate /// PODs: locally-proved statements consumed downstream. For the output /// POD: the fresh tree size (`|output_public_indices|`). The limit is @@ -215,9 +217,8 @@ pub struct SolutionBreakdown { } impl SolutionBreakdown { - /// Build a breakdown from an [`InputShape`] and its [`OutputShape`]. - /// Re-derives per-POD imports and external-pod references from the - /// dep graph and the partition; both are pure functions of the inputs. + /// Computes per-POD resource use from an [`InputShape`] and its + /// [`OutputShape`]. pub fn from_solution(input: &InputShape, output: &OutputShape) -> Self { let n = input.num_statements(); let pod_count = output.pod_count; @@ -298,15 +299,10 @@ impl SolutionBreakdown { } } - // Statement-table cap: each `OpenInputStatement` op - // produces a statement in the POD's statement table, so - // the "total statements" row reflects local statements - // PLUS chain and external imports: the same number - // `segment_feasible_with` checks against - // `max_statements`. + // Only chain imports add statements beyond assigned nodes. let total_imports = chain_imports.len() + external_imports.len(); if let Some(row) = resources.iter_mut().find(|r| r.name == "total statements") { - row.used += total_imports; + row.used += chain_imports.len(); } let imports_row = UtilizationRow { @@ -314,9 +310,9 @@ impl SolutionBreakdown { used: total_imports, limit: input.params.max_open_input_statement_ops, }; - let external_row = UtilizationRow { - name: "external pods", - used: external_pods.len(), + let input_pods_row = UtilizationRow { + name: "input pods", + used: input_pod_slots(pod_idx == 0, external_pods.len()), limit: input.params.max_input_pods, }; let publishes_row = UtilizationRow { @@ -331,7 +327,7 @@ impl SolutionBreakdown { num_statements: num_stmts, resources, imports: imports_row, - external_pods: external_row, + input_pod_slots: input_pods_row, publishes: publishes_row, } }) @@ -364,7 +360,7 @@ impl fmt::Display for SolutionBreakdown { for row in pod.resources .iter() - .chain([&pod.imports, &pod.external_pods, &pod.publishes]) + .chain([&pod.imports, &pod.input_pod_slots, &pod.publishes]) { if row.used > 0 { let pct = if row.limit > 0 { @@ -477,12 +473,13 @@ pub fn diagnose_failure(input: &InputShape) -> Option { // doesn't fit; if it doesn't fit on its own (Unsplittable above didn't // trigger but the segment-relative caps overflow), surface the // violation that broke the camel's back. + let mut checker = SegmentChecker::new(&ordering, input); let mut segment_start = 0_usize; let mut segment_index = 0_usize; let mut pos = 0_usize; while pos < n { let next_pos = pos + 1; - if super::partition::segment_feasible(&ordering, input, segment_start, next_pos) { + if checker.is_feasible(segment_start, next_pos, false) { pos = next_pos; continue; } @@ -547,10 +544,20 @@ fn identify_overflow( let s = ordering[a]; let c = &input.costs[s]; + let external_pods: HashSet = input.dep_edges[s] + .iter() + .filter_map(|dep| match dep { + AbstractDep::External { pod, .. } => Some(*pod), + AbstractDep::Internal(_) => None, + }) + .collect(); + let input_pods = input_pod_slots(segment_index == 0, external_pods.len()); + let state = ¶ms.containers.state_ops; let transition = ¶ms.containers.transition_ops; let categories: &[(&'static str, usize, usize)] = &[ ("total statements", 1, params.max_statements), + ("input pods", input_pods, params.max_input_pods), ( "merkle proofs (small)", c.merkle_proofs_small, @@ -632,6 +639,58 @@ mod tests { assert_eq!(bottleneck.min_pods(), Some(3)); // ceil(6/2) } + /// Reports the predecessor slot when another resource cap moves a + /// statement from the first POD to a later POD. + #[test] + fn diagnose_failure_names_the_input_pod_cap() { + let params = Params { + max_signed_by_ops: 1, + ..Params::default() + }; + let two_external_pods = vec![ + AbstractDep::External { + pod: 0, + statement: 0, + }, + AbstractDep::External { + pod: 1, + statement: 1, + }, + ]; + let signed_by_cost = OperationCost { + signed_by: 1, + ..OperationCost::default() + }; + let input = InputShape { + costs: vec![signed_by_cost.clone(), signed_by_cost], + dep_edges: vec![two_external_pods.clone(), two_external_pods], + output_public_indices: vec![1], + num_external_pods: 2, + statement_pod: vec![0, 1], + params, + }; + + assert!( + partition::partition(&input).is_none(), + "the second POD cannot hold two external pods alongside the chain slot" + ); + match diagnose_failure(&input).expect("failure must be diagnosed") { + CapViolation::Resource { + offending_stmt, + category, + used, + max_allowed, + .. + } => { + assert_eq!(category, "input pods"); + assert_eq!(offending_stmt, 1); + assert_eq!(used, 3); + assert_eq!(max_allowed, 2); + } + other => panic!("expected an input-pod cap violation, got {}", other), + } + } + #[test] fn solution_breakdown_reports_per_pod_utilisation() { let params = Params { diff --git a/src/frontend/multi_pod/mod.rs b/src/frontend/multi_pod/mod.rs index 0c5d6e2d..ef123beb 100644 --- a/src/frontend/multi_pod/mod.rs +++ b/src/frontend/multi_pod/mod.rs @@ -17,8 +17,8 @@ use std::{ use crate::{ frontend::{MainPod, MainPodBuilder, Operation}, middleware::{ - Hash, InputPodOpenStatement, MainPodProver, NativeOperation, OperationAux, OperationType, - Params, Statement, VDSet, Value, BASE_PARAMS, + InputPodOpenStatement, MainPodProver, NativeOperation, OperationAux, OperationType, Params, + Statement, VDSet, Value, BASE_PARAMS, }, }; @@ -33,6 +33,7 @@ mod shape; use cost::OperationCost; use deps::{DependencyGraph, ExternalDependency, StatementSource}; pub use diagnostics::{ResourceSummary, SolutionBreakdown}; +use shape::input_pod_slots; pub use shape::{AbstractDep, InputShape, OutputShape}; #[derive(Debug, thiserror::Error)] @@ -75,8 +76,9 @@ pub type Result = std::result::Result; /// [`Error::MilpUnavailable`]. #[derive(Copy, Clone, Debug, PartialEq, Eq)] pub enum SolverKind { - /// Production DP partitioner with bin-packing + random-priority Kahn - /// orderings. Polynomial time, no external solver dependency. + /// Production DP partitioner over deterministic and sampled + /// topological orderings. Polynomial time, with no external solver + /// dependency. Heuristic, /// MILP oracle (via `good_lp` / SCIP). Optimal K but non-linear time /// growth on hard instances. Intended for offline use against a @@ -93,26 +95,24 @@ impl fmt::Display for SolverKind { } } -/// Side table pairing an [`OutputShape`]'s positional external indices -/// with the concrete pod hashes and input statements they refer to. The -/// solver never sees concrete data; the build layer uses this index to -/// reattach hashes when materialising a [`MultiPodResult`] from a -/// partition. +/// Maps an [`OutputShape`]'s symbolic external indices back to concrete input +/// PODs and statements. This keeps the solver data-independent while allowing +/// the build layer to materialize concrete PODs from a partition. #[derive(Clone, Debug)] struct ExternalIndex { - pods: Vec, + input_pod_indices: Vec, statements: Vec, - /// Inverse of `pods` for O(1) hash → abstract-pod-index lookup. - pod_index_by_hash: HashMap, + /// Symbolic POD index keyed by builder input position. + abstract_index_by_input_pod: HashMap, } impl ExternalIndex { - fn new(pods: Vec, statements: Vec) -> Self { - let pod_index_by_hash = pods.iter().copied().zip(0..).collect(); + fn new(input_pod_indices: Vec, statements: Vec) -> Self { + let abstract_index_by_input_pod = input_pod_indices.iter().copied().zip(0..).collect(); Self { - pods, + input_pod_indices, statements, - pod_index_by_hash, + abstract_index_by_input_pod, } } } @@ -341,17 +341,6 @@ impl MultiPodBuilder { }); } - let input_pod_idx_by_abs: Vec = external_index - .pods - .iter() - .map(|h| { - self.input_pods - .iter() - .position(|p| p.statements_hash() == *h) - .expect("external pod referenced by user op is in input_pods") - }) - .collect(); - Ok(SolvedMultiPod { params: self.params, vd_set: self.vd_set, @@ -363,7 +352,6 @@ impl MultiPodBuilder { shape, output, external_index, - input_pod_idx_by_abs, public_sets, }) } @@ -382,10 +370,6 @@ pub struct SolvedMultiPod { shape: InputShape, output: OutputShape, external_index: ExternalIndex, - /// `external_index.pods[abs_pod]` is a hash; this maps that - /// abstract index to the matching POD's position in `input_pods`, - /// so `pod_inputs` can attach the right `MainPod` without scanning. - input_pod_idx_by_abs: Vec, /// Per-POD public sets. Computed once at `solve()` for the chain-tree /// capacity check and reused by `prove()`. public_sets: Vec>, @@ -466,15 +450,15 @@ impl SolvedMultiPod { // Staging-time aux carries a `pod_index` from the staging // builder's input slots; re-issue against this POD's // ext_slot mapping. - let OperationAux::OpenInputStatement(InputPodOpenStatement { sts_root, .. }) = + let OperationAux::OpenInputStatement(InputPodOpenStatement { pod_index, .. }) = &self.operations[s].2 else { unreachable!("OpenInputStatement op without InputPodOpenStatement aux"); }; let abs_pod = *self .external_index - .pod_index_by_hash - .get(sts_root) + .abstract_index_by_input_pod + .get(pod_index) .expect("staging OpenInputStatement's source pod is in external_index"); let slot = ext_slot[&abs_pod]; builder.open_input_st(public, slot, &self.statements[s])?; @@ -525,11 +509,14 @@ impl SolvedMultiPod { } } let mut ext_slot: HashMap = HashMap::new(); + let n_ext_pods = refs.len(); for abs_pod in refs { - let pod_idx = self.input_pod_idx_by_abs[abs_pod]; + let pod_idx = self.external_index.input_pod_indices[abs_pod]; ext_slot.insert(abs_pod, inputs.len()); inputs.push(self.input_pods[pod_idx].clone()); } + // Verify that solver slot accounting matches the inputs built here. + assert_eq!(inputs.len(), input_pod_slots(p == 0, n_ext_pods)); (inputs, ext_slot) } } @@ -644,12 +631,11 @@ fn intermediate_public_sets(shape: &InputShape, output: &OutputShape) -> Vec HashMap { +fn build_external_statement_map(input_pods: &[MainPod]) -> HashMap { let mut map = HashMap::new(); - for pod in input_pods { - let pod_hash = pod.statements_hash(); + for (pod_index, pod) in input_pods.iter().enumerate() { for stmt in pod.pod.pub_statements() { - map.insert(stmt, pod_hash); + map.entry(stmt).or_insert(pod_index); } } map @@ -680,8 +666,8 @@ fn build_shape_and_index( ) -> (InputShape, ExternalIndex) { let n_orig = operations.len(); - let mut external_pods: Vec = Vec::new(); - let mut pod_idx: HashMap = HashMap::new(); + let mut external_pods: Vec = Vec::new(); + let mut pod_idx: HashMap = HashMap::new(); let mut external_statements: Vec = Vec::new(); let mut external_statement_idx: HashMap = HashMap::new(); let mut statement_pod: Vec = Vec::new(); @@ -689,14 +675,14 @@ fn build_shape_and_index( for edges in &deps.statement_deps { for src in edges { if let StatementSource::External(ext) = src { - if let Entry::Vacant(e) = pod_idx.entry(ext.pod_hash) { + if let Entry::Vacant(e) = pod_idx.entry(ext.pod_index) { e.insert(external_pods.len()); - external_pods.push(ext.pod_hash); + external_pods.push(ext.pod_index); } if let Entry::Vacant(e) = external_statement_idx.entry(ext.clone()) { e.insert(external_statements.len()); external_statements.push(ext.clone()); - statement_pod.push(pod_idx[&ext.pod_hash]); + statement_pod.push(pod_idx[&ext.pod_index]); } } } @@ -727,6 +713,9 @@ fn build_shape_and_index( // here; the synthetic allocation loop below unions this with the // 2+ consumers rule. // + // This pre-pass cannot know which POD will host the statement. The + // partitioner performs the exact per-POD slot check. + // // Republishing any input statement forces chain use at the consumer // site, so the post-republish budget is // `(K - R) + 1 <= max_input_pods`, i.e. @@ -742,7 +731,7 @@ fn build_shape_and_index( match src { StatementSource::Internal(_) => has_internal = true, StatementSource::External(ext) => { - let pod = pod_idx[&ext.pod_hash]; + let pod = pod_idx[&ext.pod_index]; let statement = external_statement_idx[ext]; distinct_pods.insert(pod); statements_by_pod.entry(pod).or_default().push(statement); @@ -804,7 +793,10 @@ fn build_shape_and_index( .iter() .map(|op| OperationCost::from_operation(op, params)) .collect(); - costs.extend((0..n_synth).map(|_| OperationCost::default())); + costs.extend((0..n_synth).map(|_| OperationCost { + is_external_opening: true, + ..OperationCost::default() + })); // Augmented dep_edges. Original statements: External(pod, statement) // becomes Internal(synth_idx) when the input statement is being @@ -824,7 +816,7 @@ fn build_shape_and_index( AbstractDep::Internal(synth_idx) } else { AbstractDep::External { - pod: pod_idx[&ext.pod_hash], + pod: pod_idx[&ext.pod_index], statement: u, } } @@ -836,7 +828,7 @@ fn build_shape_and_index( for &u in &synthetic_to_statement { let ext = &external_statements[u]; dep_edges.push(vec![AbstractDep::External { - pod: pod_idx[&ext.pod_hash], + pod: pod_idx[&ext.pod_index], statement: u, }]); } @@ -857,10 +849,43 @@ fn build_shape_and_index( mod tests { use super::*; use crate::{ - backends::plonky2::mock::mainpod::MockProver, examples::MOCK_VD_SET, - frontend::Operation as FrontendOp, + backends::plonky2::{ + mock::mainpod::MockProver, primitives::ec::schnorr::SecretKey, signer::Signer, + }, + examples::MOCK_VD_SET, + frontend::{Operation as FrontendOp, SignedDict, SignedDictBuilder}, }; + /// Builds an external POD and returns its public `Lt(1, bound)` statement. + fn external_lt_pod( + bound: i64, + params: &Params, + vd_set: &VDSet, + prover: &MockProver, + ) -> (MainPod, Statement) { + let mut ext_builder = MainPodBuilder::new(params, vd_set); + ext_builder + .pub_op(FrontendOp::lt(1, bound)) + .expect("ext pub op"); + let pod = ext_builder.prove(prover).expect("ext prove"); + let stmt = pod + .pod + .pub_statements() + .into_iter() + .find(|s| !s.is_none()) + .expect("ext pod has a public statement"); + (pod, stmt) + } + + /// Builds a signed dictionary whose ID and signing key derive from `seed`. + fn signed_dict(seed: u32, params: &Params) -> SignedDict { + let mut signed_builder = SignedDictBuilder::new(params); + signed_builder.insert("id", seed as i64); + signed_builder + .sign(&Signer(SecretKey(seed.into()))) + .expect("sign dict") + } + #[test] fn end_to_end_solve_single_pod() { let params = Params::default(); @@ -1035,17 +1060,7 @@ mod tests { .expect("load module"); let batch = &module.batch; - let mut ext_builder = MainPodBuilder::new(¶ms, vd_set); - ext_builder - .pub_op(FrontendOp::lt(1, 2)) - .expect("ext pub op"); - let ext_pod = ext_builder.prove(&prover).expect("ext prove"); - let ext_stmt = ext_pod - .pod - .pub_statements() - .into_iter() - .find(|s| !s.is_none()) - .expect("ext pod has a public statement"); + let (ext_pod, ext_stmt) = external_lt_pod(2, ¶ms, vd_set, &prover); let mut builder = MultiPodBuilder::new(¶ms, vd_set); builder.add_pod(ext_pod).expect("add ext pod"); @@ -1066,6 +1081,69 @@ mod tests { } } + /// Later PODs reserve an input slot for their predecessor even when + /// they import no statements from it. + #[test] + fn later_pod_charges_chain_slot_against_max_input_pods() { + use crate::lang::load_module; + + let params = Params { + max_signed_by_ops: 1, + ..Params::default() + }; + let vd_set = &*MOCK_VD_SET; + let prover = MockProver {}; + + let module = load_module( + r#" + pred_ab(X, Y) = AND( + Lt(X, 2) + Lt(Y, 3) + ) + "#, + "test", + ¶ms, + &[], + ) + .expect("load module"); + let batch = &module.batch; + + let (pod_lt_2, lt_2) = external_lt_pod(2, ¶ms, vd_set, &prover); + let (pod_lt_3, lt_3) = external_lt_pod(3, ¶ms, vd_set, &prover); + + let mut builder = MultiPodBuilder::new(¶ms, vd_set); + builder.add_pod(pod_lt_2).expect("add first external pod"); + builder.add_pod(pod_lt_3).expect("add second external pod"); + builder + .priv_op(FrontendOp::dict_signed_by(&signed_dict(1, ¶ms))) + .expect("first signed_by"); + builder + .pub_op(FrontendOp::custom( + batch.predicate_ref_by_name("pred_ab").unwrap(), + [lt_2, lt_3], + )) + .expect("pred_ab over both external statements"); + builder + .priv_op(FrontendOp::dict_signed_by(&signed_dict(2, ¶ms))) + .expect("second signed_by"); + + let solved = builder.solve().expect("should solve"); + assert_eq!( + solved.solution().pod_count, + 2, + "2 SignedBy ops at max_signed_by_ops = 1 need 2 PODs" + ); + // The second POD uses one predecessor slot and one external-POD slot. + assert_eq!(solved.solution_breakdown().pods[1].input_pod_slots.used, 2); + + let result = solved.prove(&prover).expect("prove should succeed"); + for (i, pod) in result.pods.iter().enumerate() { + pod.pod + .verify() + .unwrap_or_else(|e| panic!("POD {} verification failed: {:?}", i, e)); + } + } + /// Captured `multi_pod_tree::InputShape` for zk-craft's /// `CraftRefineryCracked` action (episode-1 plugin). Used to /// stress-test the partitioner against a realistic, large input. @@ -1175,11 +1253,6 @@ mod tests { /// `max_statements`. #[test] fn signed_by_limit_forces_multi_pod_split() { - use crate::{ - backends::plonky2::{primitives::ec::schnorr::SecretKey, signer::Signer}, - frontend::SignedDictBuilder, - }; - let params = Params { max_signed_by_ops: 2, ..Params::default() @@ -1188,12 +1261,9 @@ mod tests { let prover = MockProver {}; let mut builder = MultiPodBuilder::new(¶ms, vd_set); - for i in 0..4_i64 { - let mut signed_builder = SignedDictBuilder::new(¶ms); - signed_builder.insert("id", i); - let signer = Signer(SecretKey((i as u32 + 1).into())); - let signed_dict = signed_builder.sign(&signer).expect("sign dict"); - if i == 3 { + for i in 1..=4_u32 { + let signed_dict = signed_dict(i, ¶ms); + if i == 4 { builder .pub_op(FrontendOp::dict_signed_by(&signed_dict)) .expect("pub signed_by"); @@ -1232,8 +1302,7 @@ mod tests { Operation as FrontendOp, }, middleware::{ - Hash, NativeOperation, OperationAux, OperationType, Params, RawValue, Statement, - Value, ValueRef, + NativeOperation, OperationAux, OperationType, Params, Statement, Value, ValueRef, }, }; @@ -1245,12 +1314,10 @@ mod tests { ) } - fn ext_statement(pod_seed: i64, val: i64) -> ExternalDependency { - // A unique pod hash per seed and a literal-Equal statement per - // value. The statement contents are arbitrary as long as - // different input statements hash differently. + fn ext_statement(pod_index: usize, val: i64) -> ExternalDependency { + // Distinct values keep the fixture statements distinguishable. ExternalDependency { - pod_hash: Hash::from(RawValue::from(pod_seed)), + pod_index, statement: Statement::Equal( ValueRef::Literal(Value::from(val)), ValueRef::Literal(Value::from(val)), @@ -1295,6 +1362,7 @@ mod tests { shape.dep_edges[synth_idx][0], AbstractDep::External { .. } )); + assert!(shape.costs[synth_idx].is_external_opening); } #[test] @@ -1377,5 +1445,23 @@ mod tests { assert_eq!(index.statements.len(), 2); assert_eq!(shape.num_external_pods, 1); } + + #[test] + fn concrete_input_pods_remain_distinct() { + let first = ext_statement(0, 1); + let second = ext_statement(1, 2); + let deps = DependencyGraph { + statement_deps: vec![ + vec![StatementSource::External(first)], + vec![StatementSource::External(second)], + ], + }; + let operations = vec![noop_op(), noop_op()]; + + let (shape, index) = build_shape_and_index(&operations, &deps, &[], &Params::default()); + + assert_eq!(shape.num_external_pods, 2); + assert_eq!(index.input_pod_indices, vec![0, 1]); + } } } diff --git a/src/frontend/multi_pod/partition.rs b/src/frontend/multi_pod/partition.rs index 1634dcc9..4965abb9 100644 --- a/src/frontend/multi_pod/partition.rs +++ b/src/frontend/multi_pod/partition.rs @@ -12,9 +12,10 @@ //! //! - **Picking an ordering**. Combinatorial: there's no realistic way //! to search every topological order. We sample a small set of -//! candidates: one bin-packing ordering ([`kahn_bin_packing`]), the -//! DFS-from-sinks ordering ([`build_dfs_topo_order`]), and ten -//! random-priority orderings ([`kahn_with_priority`]). +//! candidates: one bin-packing ordering ([`kahn_bin_packing`]), an +//! external-opening-first ordering, the DFS-from-sinks ordering +//! ([`build_dfs_topo_order`]), and ten random-priority orderings +//! ([`kahn_with_priority`]). //! //! - **Cutting the ordering into segments**. Once the order is fixed //! this collapses to a 1D problem: where do POD boundaries go? @@ -35,7 +36,7 @@ use rand_chacha::ChaCha20Rng; use super::{ cost::{CustomPredicateId, OperationCost, ResourceTotals}, - shape::{AbstractDep, InputShape, OutputShape}, + shape::{input_pod_slots, AbstractDep, InputShape, OutputShape}, }; use crate::middleware::Params; @@ -268,18 +269,17 @@ pub(super) fn kahn_with_priority(input: &InputShape, prio_of: &[usize]) -> Optio (ordering.len() == n).then_some(ordering) } -/// Per-segment cap check for input-tree imports (chain slot + external -/// slots) and input-pods. Tree imports are capped by -/// `params.max_open_input_statement_ops`; the chain slot counts as one -/// input-pod iff there are any prev-pod producers. +/// Checks a segment's statement-import and input-POD limits. Every POD +/// after the first reserves an input-POD slot for its predecessor. fn tree_imports_ok( + is_first_pod: bool, n_producers: usize, n_ext_imports: usize, n_ext_pods: usize, params: &Params, ) -> bool { n_producers + n_ext_imports <= params.max_open_input_statement_ops - && usize::from(n_producers > 0) + n_ext_pods <= params.max_input_pods + && input_pod_slots(is_first_pod, n_ext_pods) <= params.max_input_pods } /// Running state of the current segment used by greedy packing in @@ -420,13 +420,11 @@ impl GreedyState { let n_producers = self.prev_pod_producers.len() + self.scratch_new_producers.len(); let n_ext_imports = self.external_imports.len() + self.scratch_new_ext_imports.len(); let n_ext_pods = self.external_pods.len() + self.scratch_new_ext_pods.len(); - if !tree_imports_ok(n_producers, n_ext_imports, n_ext_pods, params) { + if !tree_imports_ok(self.a == 0, n_producers, n_ext_imports, n_ext_pods, params) { return None; } - // Statement-table cap: local segment statements plus chain and - // external imports together share `max_statements`, because each - // `OpenInputStatement` op produces a statement in the POD's table. - if tentative.num_operations + n_producers + n_ext_imports > params.max_statements { + // Only chain imports add statements beyond assigned nodes. + if tentative.num_operations + n_producers > params.max_statements { return None; } Some(self.scratch_new_producers.len() + self.scratch_new_ext_imports.len()) @@ -522,9 +520,46 @@ fn random_priority(rng: &mut ChaCha20Rng, n: usize) -> Vec { prio_of } +fn external_opening_pod(input: &InputShape, statement: usize) -> Option { + if !input.costs[statement].is_external_opening { + return None; + } + match input.dep_edges[statement].as_slice() { + [AbstractDep::External { pod, .. }] => Some(*pod), + _ => None, + } +} + +/// Prioritise external openings by source POD before unrelated work. +/// +/// The first POD can devote all of its input slots to external PODs, +/// while every later POD also needs its chain predecessor. Opening +/// statements there makes them available to later PODs through the +/// chain and can remove that additional input-slot pressure. +fn external_opening_priority(input: &InputShape) -> Option> { + let mut statements: Vec = (0..input.num_statements()).collect(); + if !statements + .iter() + .any(|&s| external_opening_pod(input, s).is_some()) + { + return None; + } + statements.sort_unstable_by_key(|&s| match external_opening_pod(input, s) { + Some(pod) => (0, pod, s), + None => (1, 0, s), + }); + + let mut priority = vec![0; statements.len()]; + for (rank, statement) in statements.into_iter().enumerate() { + priority[statement] = rank; + } + Some(priority) +} + /// Generate the candidate orderings the cutter will try. Bin-packing -/// goes first (strongest single seed on production-cap workloads), -/// then DFS-from-sinks, then the random-priority orderings for variety. +/// goes first (strongest single seed on production-cap workloads), then +/// the external-opening-first and DFS-from-sinks orderings, followed by +/// random-priority orderings for variety. fn candidate_orderings(input: &InputShape) -> Vec> { let n = input.num_statements(); let mut orderings: Vec> = Vec::new(); @@ -533,6 +568,11 @@ fn candidate_orderings(input: &InputShape) -> Vec> { if let Some(o) = kahn_bin_packing(input, &prio_id) { orderings.push(o); } + if let Some(prio) = external_opening_priority(input) { + if let Some(o) = kahn_with_priority(input, &prio) { + orderings.push(o); + } + } orderings.push(build_dfs_topo_order(input)); @@ -621,26 +661,46 @@ fn build_max_consumer_pos(consumers: &[Vec], pos_in_ordering: &[usize]) - mcp } -/// Non-terminal-only per-segment feasibility check, self-contained -/// (builds its own `pos_in_ordering` and `DpWorkspace`). Callers in a -/// hot loop should use `segment_feasible_with` to reuse allocations. -pub(super) fn segment_feasible(ordering: &[usize], input: &InputShape, a: usize, p: usize) -> bool { - let pos_in_ordering = build_pos_in_ordering(ordering); - let consumers = input.consumers(); - let max_consumer_pos = build_max_consumer_pos(&consumers, &pos_in_ordering); - let output_pub_set: HashSet = input.output_public_indices.iter().copied().collect(); - let mut ws = DpWorkspace::default(); - segment_feasible_with( - ordering, - &pos_in_ordering, - &max_consumer_pos, - &output_pub_set, - input, - a, - p, - false, - &mut ws, - ) +/// Reuses ordering metadata and scratch storage across segment checks. +pub(super) struct SegmentChecker<'a> { + ordering: &'a [usize], + input: &'a InputShape, + pos_in_ordering: Vec, + max_consumer_pos: Vec, + output_pub_set: HashSet, + workspace: DpWorkspace, +} + +impl<'a> SegmentChecker<'a> { + pub(super) fn new(ordering: &'a [usize], input: &'a InputShape) -> Self { + let pos_in_ordering = build_pos_in_ordering(ordering); + let consumers = input.consumers(); + let max_consumer_pos = build_max_consumer_pos(&consumers, &pos_in_ordering); + Self { + ordering, + input, + pos_in_ordering, + max_consumer_pos, + output_pub_set: input.output_public_indices.iter().copied().collect(), + workspace: DpWorkspace::default(), + } + } + + /// Returns whether `ordering[a..p]` fits in one POD. Terminal segments + /// use the output-POD rules. + pub(super) fn is_feasible(&mut self, a: usize, p: usize, is_terminal: bool) -> bool { + segment_feasible_with( + self.ordering, + &self.pos_in_ordering, + &self.max_consumer_pos, + &self.output_pub_set, + self.input, + a, + p, + is_terminal, + &mut self.workspace, + ) + } } /// Greedy partition of `ordering` into segments: extend each segment @@ -662,41 +722,15 @@ fn greedy_segments( if n == 0 { return Some(Vec::new()); } - let pos_in_ordering = build_pos_in_ordering(ordering); - let consumers = input.consumers(); - let max_consumer_pos = build_max_consumer_pos(&consumers, &pos_in_ordering); - let output_pub_set: HashSet = input.output_public_indices.iter().copied().collect(); - let mut ws = DpWorkspace::default(); + let mut checker = SegmentChecker::new(ordering, input); let mut segments: Vec = Vec::new(); let mut start = 0_usize; while start < n { - if !segment_feasible_with( - ordering, - &pos_in_ordering, - &max_consumer_pos, - &output_pub_set, - input, - start, - start + 1, - false, - &mut ws, - ) { + if !checker.is_feasible(start, start + 1, false) { return None; } let mut end = start + 1; - while end < n - && segment_feasible_with( - ordering, - &pos_in_ordering, - &max_consumer_pos, - &output_pub_set, - input, - start, - end + 1, - false, - &mut ws, - ) - { + while end < n && checker.is_feasible(start, end + 1, false) { end += 1; } segments.push(Segment { start, end }); @@ -706,17 +740,7 @@ fn greedy_segments( let last = *segments .last() .expect("non-empty ordering produced no segments"); - if !segment_feasible_with( - ordering, - &pos_in_ordering, - &max_consumer_pos, - &output_pub_set, - input, - last.start, - last.end, - true, - &mut ws, - ) { + if !checker.is_feasible(last.start, last.end, true) { return None; } } @@ -768,6 +792,7 @@ fn segment_feasible_with( // - external imports: external (input) statements (slots 1..N). // External-pod references are tracked separately for `max_input_pods`, // which is a per-slot cap, not a per-statement one. + let is_first_pod = start == 0; let mut totals = ResourceTotals::default(); workspace.distinct_cps.clear(); workspace.prev_pod_producers.clear(); @@ -798,12 +823,11 @@ fn segment_feasible_with( } } } - // Mid-loop bail: the chain-slot + external-slot tree-imports - // cap can only grow as we add more statements, so once it's - // busted there's no recovery. Saves the remaining statements' - // worth of inserts on infeasible segments. + // Import counts only increase as the segment grows, so stop once + // either cap is exceeded. if workspace.prev_pod_producers.len() + workspace.external_imports.len() > params.max_open_input_statement_ops + || input_pod_slots(is_first_pod, workspace.external_pods.len()) > params.max_input_pods { return false; } @@ -813,14 +837,14 @@ fn segment_feasible_with( let n_ext_pods = workspace.external_pods.len(); let n_chain_imports = workspace.prev_pod_producers.len(); - // Statement-table cap. Each `OpenInputStatement` op produces a statement - // in the POD's statement table, so the table holds `segment statements + - // imports`, capped by `max_statements`. - if segment.len() + n_chain_imports + n_ext_imports > params.max_statements { + // Only chain imports add statements beyond assigned nodes. + if segment.len() + n_chain_imports > params.max_statements { return false; } - if !tree_imports_ok(n_chain_imports, n_ext_imports, n_ext_pods, params) { + let imports_ok = + |n_chain: usize| tree_imports_ok(is_first_pod, n_chain, n_ext_imports, n_ext_pods, params); + if !imports_ok(n_chain_imports) { return false; } @@ -849,10 +873,10 @@ fn segment_feasible_with( } } let n_chain_imports_terminal = workspace.prev_pod_producers.len(); - if segment.len() + n_chain_imports_terminal + n_ext_imports > params.max_statements { + if segment.len() + n_chain_imports_terminal > params.max_statements { return false; } - tree_imports_ok(n_chain_imports_terminal, n_ext_imports, n_ext_pods, params) + imports_ok(n_chain_imports_terminal) } /// A half-open boundary range covering one POD's statements: @@ -940,10 +964,7 @@ struct DpEntry { #[allow(clippy::needless_range_loop)] fn run_dp(ordering: &[usize], input: &InputShape) -> Option> { let n = ordering.len(); - let pos_in_ordering = build_pos_in_ordering(ordering); - let consumers = input.consumers(); - let max_consumer_pos = build_max_consumer_pos(&consumers, &pos_in_ordering); - let output_pub_set: HashSet = input.output_public_indices.iter().copied().collect(); + let mut checker = SegmentChecker::new(ordering, input); // Each POD holds at most `max_statements` local statements, so any // candidate segment longer than that is infeasible. This bounds // the inner loop's start window per `end`. @@ -952,7 +973,6 @@ fn run_dp(ordering: &[usize], input: &InputShape) -> Option> { // The table has `n + 1` cells, one per boundary position (including // 0 and n). `dp[0]` is the base case: the empty prefix needs 0 PODs. let mut dp: Vec> = vec![None; n + 1]; - let mut workspace = DpWorkspace::default(); dp[0] = Some(DpEntry { pod_count: 0, prev_start: 0, @@ -965,23 +985,13 @@ fn run_dp(ordering: &[usize], input: &InputShape) -> Option> { let best_segment_ending_at = |end: usize, kind: PodKind, dp: &[Option], - workspace: &mut DpWorkspace| + checker: &mut SegmentChecker| -> Option { let window_start = end.saturating_sub(max_segment_len); let mut best: Option = None; for start in window_start..end { let Some(prev) = dp[start] else { continue }; - if !segment_feasible_with( - ordering, - &pos_in_ordering, - &max_consumer_pos, - &output_pub_set, - input, - start, - end, - kind.is_terminal(), - workspace, - ) { + if !checker.is_feasible(start, end, kind.is_terminal()) { continue; } let candidate = DpEntry { @@ -999,14 +1009,14 @@ fn run_dp(ordering: &[usize], input: &InputShape) -> Option> { // as a chain-extending POD (the output POD's special rules are // handled separately below). for end in 1..=n { - let entry = best_segment_ending_at(end, PodKind::ChainExtending, &dp, &mut workspace); + let entry = best_segment_ending_at(end, PodKind::ChainExtending, &dp, &mut checker); dp[end] = entry; } // Terminal scan: pick the cheapest output POD covering `start..n`, // using the output-POD feasibility flavour. Returns `None` if no // candidate is feasible. - let terminal = best_segment_ending_at(n, PodKind::Output, &dp, &mut workspace)?; + let terminal = best_segment_ending_at(n, PodKind::Output, &dp, &mut checker)?; // Backtrack: walk `prev_start` breadcrumbs from `terminal` back to // boundary 0 to recover the actual cut positions. @@ -1173,6 +1183,101 @@ mod tests { ); } + #[test] + fn external_open_node_consumes_one_statement_slot() { + use super::super::cost::OperationCost; + + let input = InputShape { + costs: vec![OperationCost::default(), OperationCost::default()], + dep_edges: vec![ + vec![AbstractDep::External { + pod: 0, + statement: 0, + }], + vec![AbstractDep::Internal(0)], + ], + output_public_indices: vec![1], + num_external_pods: 1, + statement_pod: vec![0], + params: Params { + max_statements: 2, + ..Params::default() + }, + }; + + let out = partition(&input).expect("Open plus consumer should fit"); + assert_eq!(out.pod_count, 1); + assert_eq!(out.pod_statements, vec![vec![0, 1]]); + } + + #[test] + fn external_openings_can_move_forward_to_save_a_pod() { + let params = Params { + max_statements: 5, + max_input_pods: 2, + max_signed_by_ops: 1, + max_public_statements: 3, + ..Params::default() + }; + let signed = || OperationCost { + signed_by: 1, + ..OperationCost::default() + }; + let opening = || OperationCost { + is_external_opening: true, + ..OperationCost::default() + }; + let input = InputShape { + costs: vec![ + signed(), + signed(), + opening(), + opening(), + OperationCost::default(), + ], + dep_edges: vec![ + vec![], + vec![], + vec![AbstractDep::External { + pod: 0, + statement: 0, + }], + vec![AbstractDep::External { + pod: 1, + statement: 1, + }], + vec![ + AbstractDep::Internal(0), + AbstractDep::Internal(2), + AbstractDep::Internal(3), + ], + ], + output_public_indices: vec![4], + num_external_pods: 2, + statement_pod: vec![0, 1], + params, + }; + + // Source order puts both external openings after the first + // resource-forced cut, where the predecessor takes a third slot. + let source_order: Vec = (0..input.num_statements()).collect(); + let source_out = + partition_with_ordering(&input, &source_order).expect("source order should partition"); + assert_eq!(source_out.pod_count, 3); + + let priority = external_opening_priority(&input).expect("input contains external openings"); + let opening_order = + kahn_with_priority(&input, &priority).expect("dependency graph should be acyclic"); + assert_eq!(opening_order, vec![2, 3, 0, 1, 4]); + assert_eq!(candidate_orderings(&input)[1], opening_order); + + let opening_out = partition_with_ordering(&input, &opening_order) + .expect("external-opening order should partition"); + assert_eq!(opening_out.pod_count, 2); + assert!(opening_out.pod_statements[0].contains(&2)); + assert!(opening_out.pod_statements[0].contains(&3)); + } + #[test] fn dependency_chain_respects_topo_order() { // 4 statements where each depends on the previous. With @@ -1226,7 +1331,7 @@ mod tests { use rand::SeedableRng; use rand_chacha::ChaCha20Rng; - use super::super::partition_milp::random_input; + use super::super::{partition_milp::random_input, SolutionBreakdown}; let param_variants: Vec<(&str, Params)> = vec![ ( @@ -1295,6 +1400,30 @@ mod tests { let input = random_input(&mut rng, n, params.clone()); let identity: Vec = (0..input.num_statements()).collect(); + // Recompute per-POD usage from the dependency graph to catch + // accounting errors shared by the DP and greedy partitioners. + if let Some(sol) = partition(&input) { + for pod in SolutionBreakdown::from_solution(&input, &sol).pods { + for row in pod.resources.iter().chain([ + &pod.imports, + &pod.input_pod_slots, + &pod.publishes, + ]) { + assert!( + row.used <= row.limit, + "POD {} over the {} cap [{} n={} trial={}]: {} > {}", + pod.pod_idx, + row.name, + label, + n, + trial, + row.used, + row.limit, + ); + } + } + } + let k_per_ord_dp = |o: &[usize]| -> Option { partition_with_ordering(&input, o).map(|s| s.pod_count) }; diff --git a/src/frontend/multi_pod/partition_milp.rs b/src/frontend/multi_pod/partition_milp.rs index f1b2e746..09b7ab2c 100644 --- a/src/frontend/multi_pod/partition_milp.rs +++ b/src/frontend/multi_pod/partition_milp.rs @@ -16,7 +16,7 @@ use good_lp::{ use super::{ cost::{CustomPredicateId, OperationCost, ResourceTotals}, - shape::{AbstractDep, InputShape, OutputShape}, + shape::{input_pod_slots, AbstractDep, InputShape, OutputShape}, }; use crate::middleware::Params; @@ -28,7 +28,6 @@ struct MilpVars { ext_import_from: Vec>, ext_used: Vec>, cp_used: Vec>, - uses_chain: Vec, } fn mk_binary_grid(vars: &mut ProblemVariables, rows: usize, cols: usize) -> Vec> { @@ -51,7 +50,6 @@ fn declare_vars( ext_import_from: mk_binary_grid(vars, num_ext_statements, k), ext_used: mk_binary_grid(vars, num_ext_pods, k), cp_used: mk_binary_grid(vars, num_cps, k), - uses_chain: (0..k).map(|_| vars.add(variable().binary())).collect(), } } @@ -229,18 +227,12 @@ pub fn solve_for_k(input: &InputShape, k: usize) -> Option { } } - // (3) Statement-table cap per POD. Each `OpenInputStatement` op - // produces a statement, so the table holds `local statements + - // chain imports + external-statement imports`, capped by - // `max_statements`. + // (3) Statement cap: only chain imports add nodes beyond assignments. for p in 0..k { let assign_sum: Expression = (0..n).map(|s| v.assign[s][p]).sum(); let chain_sum: Expression = (0..n).map(|d| v.import_from[d][p]).sum(); - let ext_sum: Expression = (0..num_ext_statements) - .map(|e_prem| v.ext_import_from[e_prem][p]) - .sum(); model.add_constraint(constraint!( - assign_sum + chain_sum + ext_sum <= input.params.max_statements as f64 + assign_sum + chain_sum <= input.params.max_statements as f64 )); } @@ -390,28 +382,17 @@ pub fn solve_for_k(input: &InputShape, k: usize) -> Option { } } - // (9) uses_chain[p] = OR of import_from[*][p]. - for p in 0..k { - for d in 0..n { - model.add_constraint(constraint!(v.uses_chain[p] >= v.import_from[d][p])); - } - // Upper bound: at most 1 if any import is set. We don't need a - // tight upper bound since the input-pod cap is the only thing - // that reads uses_chain, and a slack uses_chain = 1 is harmless - // when no imports are taken. - let sum: Expression = (0..n).map(|d| v.import_from[d][p]).sum(); - model.add_constraint(constraint!(v.uses_chain[p] <= sum)); - } - - // (10) Input-pod cap: uses_chain + external pods <= max_input_pods. + // (9) Input-POD cap: later PODs reserve one predecessor slot; each + // distinct external POD uses another. for p in 0..k { + let chain_pods = input_pod_slots(p == 0, 0) as f64; let ext_sum: Expression = (0..num_ext_pods).map(|e| v.ext_used[e][p]).sum(); model.add_constraint(constraint!( - v.uses_chain[p] + ext_sum <= input.params.max_input_pods as f64 + ext_sum + chain_pods <= input.params.max_input_pods as f64 )); } - // (11) POD-range preprocessing: each statement's assignable PODs are + // (10) POD-range preprocessing: each statement's assignable PODs are // bounded by `[min_pod(s), max_pod(s)]` derived from upstream and // downstream resource sums. Fix assigns outside this range to 0. // Pure constraint tightening; never rules out a feasible partition, @@ -613,6 +594,30 @@ mod tests { ); } + #[test] + fn external_open_node_consumes_one_statement_slot() { + let input = InputShape { + costs: vec![OperationCost::default(), OperationCost::default()], + dep_edges: vec![ + vec![AbstractDep::External { + pod: 0, + statement: 0, + }], + vec![AbstractDep::Internal(0)], + ], + output_public_indices: vec![1], + num_external_pods: 1, + statement_pod: vec![0], + params: Params { + max_statements: 2, + ..Params::default() + }, + }; + + let out = solve_for_k(&input, 1).expect("Open plus consumer should fit"); + assert_eq!(out.pod_statements, vec![vec![0, 1]]); + } + #[test] fn splits_into_two_pods_when_count_exceeds_cap() { let params = Params { diff --git a/src/frontend/multi_pod/shape.rs b/src/frontend/multi_pod/shape.rs index aff2b72f..b4ea73d1 100644 --- a/src/frontend/multi_pod/shape.rs +++ b/src/frontend/multi_pod/shape.rs @@ -32,6 +32,14 @@ pub enum AbstractDep { }, } +/// Returns the input-POD slots used: one per distinct external POD, plus +/// a predecessor slot for each POD after the first. The predecessor stays +/// in slot 0 to extend its public-statement tree even when no statements +/// are imported from it. +pub(super) fn input_pod_slots(is_first_pod: bool, num_external_pods: usize) -> usize { + usize::from(!is_first_pod) + num_external_pods +} + /// Symbolic input to the solver: the structure of a multi-POD problem in /// positional form. #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]