Skip to content

Commit 06e6d8d

Browse files
georgwiesegeorgwiese
committed
Refactor
1 parent 60632a3 commit 06e6d8d

File tree

9 files changed

+148
-120
lines changed

9 files changed

+148
-120
lines changed

‎autoprecompiles/src/adapter.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use std::{fmt::Display, sync::Arc};
66
use powdr_number::FieldElement;
77
use serde::{Deserialize, Serialize};
88

9-
use crate::ExecutionStats;
9+
use crate::DataDrivenConstraints;
1010
use crate::{
1111
blocks::{BasicBlock, Instruction, Program},
1212
constraint_optimizer::IsBusStateful,
@@ -46,7 +46,7 @@ pub trait PgoAdapter {
4646
config: &PowdrConfig,
4747
vm_config: AdapterVmConfig<Self::Adapter>,
4848
labels: BTreeMap<u64, Vec<String>>,
49-
execution_stats: ExecutionStats,
49+
execution_stats: DataDrivenConstraints,
5050
) -> Vec<AdapterApcWithStats<Self::Adapter>> {
5151
let filtered_blocks = blocks
5252
.into_iter()
@@ -61,7 +61,7 @@ pub trait PgoAdapter {
6161
config: &PowdrConfig,
6262
vm_config: AdapterVmConfig<Self::Adapter>,
6363
labels: BTreeMap<u64, Vec<String>>,
64-
execution_stats: ExecutionStats,
64+
execution_stats: DataDrivenConstraints,
6565
) -> Vec<AdapterApcWithStats<Self::Adapter>>;
6666

6767
fn pc_execution_count(&self, _pc: u64) -> Option<u32> {

‎autoprecompiles/src/data_driven_constraints.rs‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
use std::collections::BTreeMap;
2+
3+
use serde::{Deserialize, Serialize};
4+
5+
use crate::{
6+
adapter::Adapter,
7+
blocks::BasicBlock,
8+
expression::{AlgebraicExpression, AlgebraicReference},
9+
SymbolicConstraint,
10+
};
11+
12+
/// "Constraints" that were inferred from execution statistics.
13+
#[derive(Serialize, Deserialize, Clone, Default)]
14+
pub struct DataDrivenConstraints {
15+
/// For each program counter, the range constraints for each column.
16+
/// The range might not hold in 100% of cases.
17+
pub column_ranges_by_pc: BTreeMap<u32, Vec<(u32, u32)>>,
18+
/// For each basic block (identified by its starting PC), the equivalence classes of columns.
19+
/// Each equivalence class is a list of (instruction index in block, column index).
20+
pub equivalence_classes_by_block: BTreeMap<u64, Vec<Vec<(usize, usize)>>>,
21+
}
22+
23+
/// Debug information mapping AIR ids to program counters and column names.
24+
#[derive(Serialize, Deserialize)]
25+
pub struct DebugInfo {
26+
/// Mapping from program counter to AIR id.
27+
pub air_id_by_pc: BTreeMap<u32, usize>,
28+
/// Mapping from AIR id to column names.
29+
pub column_names_by_air_id: BTreeMap<usize, Vec<String>>,
30+
}
31+
32+
#[derive(Serialize, Deserialize)]
33+
pub struct ExecutionStatsJson {
34+
pub execution_stats: DataDrivenConstraints,
35+
pub debug_info: DebugInfo,
36+
}
37+
38+
pub fn add_ai_constraints<A: Adapter>(
39+
execution_stats: &DataDrivenConstraints,
40+
subs: &[Vec<u64>],
41+
block: &BasicBlock<A::Instruction>,
42+
columns: impl Iterator<Item = AlgebraicReference>,
43+
) -> (
44+
Vec<SymbolicConstraint<A::PowdrField>>,
45+
Vec<SymbolicConstraint<A::PowdrField>>,
46+
) {
47+
let range_constraints = &execution_stats.column_ranges_by_pc;
48+
let equivalence_classes_by_block = &execution_stats.equivalence_classes_by_block;
49+
50+
let mut range_analyzer_constraints = Vec::new();
51+
let mut equivalence_analyzer_constraints = Vec::new();
52+
53+
// Mapping (instruction index, column index) -> AlgebraicReference
54+
let reverse_subs = subs
55+
.iter()
56+
.enumerate()
57+
.flat_map(|(instr_index, subs)| {
58+
subs.iter()
59+
.enumerate()
60+
.map(move |(col_index, &poly_id)| (poly_id, (instr_index, col_index)))
61+
})
62+
.collect::<BTreeMap<_, _>>();
63+
let algebraic_references = columns
64+
.map(|r| (reverse_subs.get(&r.id).unwrap().clone(), r.clone()))
65+
.collect::<BTreeMap<_, _>>();
66+
67+
for i in 0..block.statements.len() {
68+
let pc = (block.start_pc + (i * 4) as u64) as u32;
69+
let Some(range_constraints) = range_constraints.get(&pc) else {
70+
continue;
71+
};
72+
for (col_index, range) in range_constraints.iter().enumerate() {
73+
if range.0 == range.1 {
74+
let value = A::PowdrField::from(range.0 as u64);
75+
let Some(reference) = algebraic_references.get(&(i, col_index)).cloned() else {
76+
panic!(
77+
"Missing reference for (i: {}, col_index: {}, block_id: {})",
78+
i, col_index, block.start_pc
79+
);
80+
};
81+
let constraint =
82+
AlgebraicExpression::Reference(reference) - AlgebraicExpression::Number(value);
83+
84+
range_analyzer_constraints.push(SymbolicConstraint { expr: constraint });
85+
}
86+
}
87+
}
88+
89+
if let Some(equivalence_classes) = equivalence_classes_by_block.get(&block.start_pc) {
90+
for equivalence_class in equivalence_classes {
91+
let first = equivalence_class.first().unwrap();
92+
let Some(first_ref) = algebraic_references.get(first).cloned() else {
93+
// TODO: This fails in some blocks. For now, just return no extra constraints.
94+
println!(
95+
"Missing reference for (i: {}, col_index: {}, block_id: {})",
96+
first.0, first.1, block.start_pc
97+
);
98+
return (range_analyzer_constraints, vec![]);
99+
};
100+
for other in equivalence_class.iter().skip(1) {
101+
let Some(other_ref) = algebraic_references.get(other).cloned() else {
102+
// TODO: This fails in some blocks. For now, just return no extra constraints.
103+
println!(
104+
"Missing reference for (i: {}, col_index: {}, block_id: {})",
105+
other.0, other.1, block.start_pc
106+
);
107+
return (range_analyzer_constraints, vec![]);
108+
};
109+
let constraint = AlgebraicExpression::Reference(first_ref.clone())
110+
- AlgebraicExpression::Reference(other_ref.clone());
111+
equivalence_analyzer_constraints.push(SymbolicConstraint { expr: constraint });
112+
}
113+
}
114+
}
115+
116+
(range_analyzer_constraints, equivalence_analyzer_constraints)
117+
}

‎autoprecompiles/src/lib.rs‎

Lines changed: 4 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use crate::adapter::{Adapter, AdapterApc, AdapterVmConfig};
22
use crate::blocks::BasicBlock;
33
use crate::bus_map::{BusMap, BusType};
4+
use crate::data_driven_constraints::{add_ai_constraints, DataDrivenConstraints};
45
use crate::evaluation::AirStats;
56
use crate::expression_conversion::algebraic_to_grouped_expression;
67
use crate::symbolic_machine_generator::convert_machine_field_type;
@@ -12,7 +13,7 @@ use powdr_expression::{
1213
visitors::Children, AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicUnaryOperation,
1314
};
1415
use serde::{Deserialize, Serialize};
15-
use std::collections::{BTreeMap, BTreeSet};
16+
use std::collections::BTreeSet;
1617
use std::fmt::Display;
1718
use std::io::BufWriter;
1819
use std::iter::once;
@@ -26,6 +27,7 @@ pub mod adapter;
2627
pub mod blocks;
2728
pub mod bus_map;
2829
pub mod constraint_optimizer;
30+
pub mod data_driven_constraints;
2931
pub mod evaluation;
3032
pub mod execution_profile;
3133
pub mod expression;
@@ -416,111 +418,12 @@ impl ColumnAllocator {
416418
}
417419
}
418420

419-
#[derive(Serialize, Deserialize)]
420-
pub struct DebugInfo {
421-
pub air_id_by_pc: BTreeMap<u32, usize>,
422-
pub column_names_by_air_id: BTreeMap<usize, Vec<String>>,
423-
}
424-
425-
#[derive(Serialize, Deserialize, Clone)]
426-
pub struct ExecutionStats {
427-
pub column_ranges_by_pc: BTreeMap<u32, Vec<(u32, u32)>>,
428-
pub equivalence_classes_by_block: BTreeMap<u64, Vec<Vec<(usize, usize)>>>,
429-
}
430-
431-
#[derive(Serialize, Deserialize)]
432-
pub struct ExecutionStatsJson {
433-
pub execution_stats: ExecutionStats,
434-
pub debug_info: DebugInfo,
435-
}
436-
437-
fn add_ai_constraints<A: Adapter>(
438-
execution_stats: &ExecutionStats,
439-
subs: &[Vec<u64>],
440-
block: &BasicBlock<A::Instruction>,
441-
columns: impl Iterator<Item = AlgebraicReference>,
442-
) -> (
443-
Vec<SymbolicConstraint<A::PowdrField>>,
444-
Vec<SymbolicConstraint<A::PowdrField>>,
445-
) {
446-
let range_constraints = &execution_stats.column_ranges_by_pc;
447-
let equivalence_classes_by_block = &execution_stats.equivalence_classes_by_block;
448-
449-
let mut range_analyzer_constraints = Vec::new();
450-
let mut equivalence_analyzer_constraints = Vec::new();
451-
452-
// Mapping (instruction index, column index) -> AlgebraicReference
453-
let reverse_subs = subs
454-
.iter()
455-
.enumerate()
456-
.flat_map(|(instr_index, subs)| {
457-
subs.iter()
458-
.enumerate()
459-
.map(move |(col_index, &poly_id)| (poly_id, (instr_index, col_index)))
460-
})
461-
.collect::<BTreeMap<_, _>>();
462-
let algebraic_references = columns
463-
.map(|r| (reverse_subs.get(&r.id).unwrap().clone(), r.clone()))
464-
.collect::<BTreeMap<_, _>>();
465-
466-
for i in 0..block.statements.len() {
467-
let pc = (block.start_pc + (i * 4) as u64) as u32;
468-
let Some(range_constraints) = range_constraints.get(&pc) else {
469-
continue;
470-
};
471-
for (col_index, range) in range_constraints.iter().enumerate() {
472-
if range.0 == range.1 {
473-
let value = A::PowdrField::from(range.0 as u64);
474-
let Some(reference) = algebraic_references.get(&(i, col_index)).cloned() else {
475-
panic!(
476-
"Missing reference for (i: {}, col_index: {}, block_id: {})",
477-
i, col_index, block.start_pc
478-
);
479-
};
480-
let constraint =
481-
AlgebraicExpression::Reference(reference) - AlgebraicExpression::Number(value);
482-
483-
range_analyzer_constraints.push(SymbolicConstraint { expr: constraint });
484-
}
485-
}
486-
}
487-
488-
if let Some(equivalence_classes) = equivalence_classes_by_block.get(&block.start_pc) {
489-
for equivalence_class in equivalence_classes {
490-
let first = equivalence_class.first().unwrap();
491-
let Some(first_ref) = algebraic_references.get(first).cloned() else {
492-
// TODO: This fails in some blocks. For now, just return no extra constraints.
493-
println!(
494-
"Missing reference for (i: {}, col_index: {}, block_id: {})",
495-
first.0, first.1, block.start_pc
496-
);
497-
return (range_analyzer_constraints, vec![]);
498-
};
499-
for other in equivalence_class.iter().skip(1) {
500-
let Some(other_ref) = algebraic_references.get(other).cloned() else {
501-
// TODO: This fails in some blocks. For now, just return no extra constraints.
502-
println!(
503-
"Missing reference for (i: {}, col_index: {}, block_id: {})",
504-
other.0, other.1, block.start_pc
505-
);
506-
return (range_analyzer_constraints, vec![]);
507-
};
508-
let constraint = AlgebraicExpression::Reference(first_ref.clone())
509-
- AlgebraicExpression::Reference(other_ref.clone());
510-
equivalence_analyzer_constraints.push(SymbolicConstraint { expr: constraint });
511-
}
512-
}
513-
}
514-
515-
(range_analyzer_constraints, equivalence_analyzer_constraints)
516-
}
517-
518421
pub fn build<A: Adapter>(
519422
block: BasicBlock<A::Instruction>,
520423
vm_config: AdapterVmConfig<A>,
521424
degree_bound: DegreeBound,
522425
apc_candidates_dir_path: Option<&Path>,
523-
execution_stats: &ExecutionStats,
426+
execution_stats: &DataDrivenConstraints,
524427
) -> Result<AdapterApc<A>, crate::constraint_optimizer::Error> {
525428
let start = std::time::Instant::now();
526429

‎autoprecompiles/src/pgo/cell/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use crate::{
1313
blocks::BasicBlock,
1414
evaluation::EvaluationResult,
1515
pgo::cell::selection::parallel_fractional_knapsack,
16-
ExecutionStats, PowdrConfig,
16+
DataDrivenConstraints, PowdrConfig,
1717
};
1818

1919
mod selection;
@@ -92,7 +92,7 @@ impl<A: Adapter + Send + Sync, C: Candidate<A> + Send + Sync> PgoAdapter for Cel
9292
config: &PowdrConfig,
9393
vm_config: AdapterVmConfig<Self::Adapter>,
9494
labels: BTreeMap<u64, Vec<String>>,
95-
execution_stats: ExecutionStats,
95+
execution_stats: DataDrivenConstraints,
9696
) -> Vec<AdapterApcWithStats<Self::Adapter>> {
9797
tracing::info!(
9898
"Generating autoprecompiles with cell PGO for {} blocks",

‎autoprecompiles/src/pgo/instruction.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use crate::{
44
adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, PgoAdapter},
55
blocks::BasicBlock,
66
pgo::create_apcs_for_all_blocks,
7-
ExecutionStats, PowdrConfig,
7+
DataDrivenConstraints, PowdrConfig,
88
};
99

1010
pub struct InstructionPgo<A> {
@@ -30,7 +30,7 @@ impl<A: Adapter> PgoAdapter for InstructionPgo<A> {
3030
config: &PowdrConfig,
3131
vm_config: AdapterVmConfig<Self::Adapter>,
3232
_labels: BTreeMap<u64, Vec<String>>,
33-
execution_stats: ExecutionStats,
33+
execution_stats: DataDrivenConstraints,
3434
) -> Vec<AdapterApcWithStats<Self::Adapter>> {
3535
tracing::info!(
3636
"Generating autoprecompiles with instruction PGO for {} blocks",

‎autoprecompiles/src/pgo/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use strum::{Display, EnumString};
66
use crate::{
77
adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, ApcWithStats},
88
blocks::BasicBlock,
9-
ExecutionStats, PowdrConfig,
9+
DataDrivenConstraints, PowdrConfig,
1010
};
1111

1212
mod cell;
@@ -77,7 +77,7 @@ fn create_apcs_for_all_blocks<A: Adapter>(
7777
blocks: Vec<BasicBlock<A::Instruction>>,
7878
config: &PowdrConfig,
7979
vm_config: AdapterVmConfig<A>,
80-
execution_stats: ExecutionStats,
80+
execution_stats: DataDrivenConstraints,
8181
) -> Vec<AdapterApcWithStats<A>> {
8282
let n_acc = config.autoprecompiles as usize;
8383
tracing::info!("Generating {n_acc} autoprecompiles in parallel");

‎autoprecompiles/src/pgo/none.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use crate::{
44
adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, PgoAdapter},
55
blocks::BasicBlock,
66
pgo::create_apcs_for_all_blocks,
7-
ExecutionStats, PowdrConfig,
7+
DataDrivenConstraints, PowdrConfig,
88
};
99

1010
pub struct NonePgo<A> {
@@ -29,7 +29,7 @@ impl<A: Adapter> PgoAdapter for NonePgo<A> {
2929
config: &PowdrConfig,
3030
vm_config: AdapterVmConfig<Self::Adapter>,
3131
_labels: BTreeMap<u64, Vec<String>>,
32-
execution_stats: ExecutionStats,
32+
execution_stats: DataDrivenConstraints,
3333
) -> Vec<AdapterApcWithStats<Self::Adapter>> {
3434
// cost = number_of_original_instructions
3535
blocks.sort_by(|a, b| b.statements.len().cmp(&a.statements.len()));

0 commit comments

Comments
 (0)