improve apc interface

autoprecompiles/src/adapter.rs

@@ -6,6 +6,7 @@ use std::{fmt::Display, sync::Arc};
 use powdr_number::FieldElement;
 use serde::{Deserialize, Serialize};
 
+use crate::evaluation::{ApcPerformanceReport, ApcStats};
 use crate::{
     blocks::{BasicBlock, Instruction, Program},
     constraint_optimizer::IsBusStateful,
@@ -15,29 +16,23 @@ use crate::{
 };
 
 #[derive(Serialize, Deserialize)]
-pub struct ApcWithStats<F, I, S> {
+pub struct ApcWithReport<F, I, S> {
     apc: Arc<Apc<F, I>>,
-    stats: Option<S>,
+    report: ApcPerformanceReport<S>,
 }
-impl<F, I, S> ApcWithStats<F, I, S> {
-    pub fn with_stats(mut self, stats: S) -> Self {
-        self.stats = Some(stats);
-        self
+impl<F, I, S> ApcWithReport<F, I, S> {
+    pub fn new(apc: Arc<Apc<F, I>>, report: ApcPerformanceReport<S>) -> Self {
+        Self { apc, report }
     }
 
-    pub fn into_parts(self) -> (Arc<Apc<F, I>>, Option<S>) {
-        (self.apc, self.stats)
-    }
-}
-
-impl<F, I, S> From<Arc<Apc<F, I>>> for ApcWithStats<F, I, S> {
-    fn from(apc: Arc<Apc<F, I>>) -> Self {
-        Self { apc, stats: None }
+    pub fn into_parts(self) -> (Arc<Apc<F, I>>, ApcPerformanceReport<S>) {
+        (self.apc, self.report)
     }
 }
 
 pub trait PgoAdapter {
     type Adapter: Adapter;
+    type Air: ApcArithmetization<Self::Adapter>;
 
     fn filter_blocks_and_create_apcs_with_pgo(
         &self,
@@ -66,10 +61,18 @@ pub trait PgoAdapter {
     }
 }
 
+pub trait ApcArithmetization<A: Adapter>: Send + Sync {
+    /// Given an apc circuit and a degree bound, return the stats when compiling this apc using this arithmetization
+    fn get_metrics(apc: Arc<AdapterApc<A>>, max_constraint_degree: usize) -> A::ApcStats;
+}
+
 pub trait Adapter: Sized
 where
-    Self::InstructionHandler:
-        InstructionHandler<Field = Self::Field, Instruction = Self::Instruction>,
+    Self::InstructionHandler: InstructionHandler<
+        Field = Self::Field,
+        Instruction = Self::Instruction,
+        ApcStats = Self::ApcStats,
+    >,
 {
     type Field: Serialize + for<'de> Deserialize<'de> + Send + Sync + Clone;
     type PowdrField: FieldElement;
@@ -86,7 +89,7 @@ where
         V,
     >;
     type CustomBusTypes: Clone + Display + Sync + Eq + PartialEq;
-    type ApcStats: Send + Sync;
+    type ApcStats: ApcStats;
     type AirId: Eq + Hash + Send + Sync;
 
     fn into_field(e: Self::PowdrField) -> Self::Field;
@@ -99,8 +102,7 @@ where
 }
 
 pub type AdapterApcWithStats<A> =
-    ApcWithStats<<A as Adapter>::Field, <A as Adapter>::Instruction, <A as Adapter>::ApcStats>;
-pub type ApcStats<A> = <A as Adapter>::ApcStats;
+    ApcWithReport<<A as Adapter>::Field, <A as Adapter>::Instruction, <A as Adapter>::ApcStats>;
 pub type AdapterApc<A> = Apc<<A as Adapter>::Field, <A as Adapter>::Instruction>;
 pub type AdapterVmConfig<'a, A> = VmConfig<
     'a,

autoprecompiles/src/evaluation.rs

@@ -1,6 +1,9 @@
-use std::{fmt::Display, iter::Sum, ops::Add};
+use std::{fmt::Display, iter::Sum, ops::Add, sync::Arc};
 
-use crate::{blocks::Instruction, InstructionHandler, SymbolicMachine};
+use crate::{
+    adapter::{Adapter, AdapterApc, ApcArithmetization},
+    InstructionHandler,
+};
 
 use serde::{Deserialize, Serialize};
 
@@ -17,16 +20,6 @@ pub struct AirStats {
     pub bus_interactions: usize,
 }
 
-impl AirStats {
-    pub fn new<F: Clone + Ord + std::fmt::Display>(machine: &SymbolicMachine<F>) -> Self {
-        Self {
-            main_columns: machine.main_columns().count(),
-            constraints: machine.constraints.len(),
-            bus_interactions: machine.bus_interactions.len(),
-        }
-    }
-}
-
 impl Add for AirStats {
     type Output = AirStats;
     fn add(self, rhs: AirStats) -> AirStats {
@@ -46,38 +39,58 @@ impl Sum<AirStats> for AirStats {
 
 #[derive(Clone, Copy, Serialize, Deserialize)]
 /// Evaluation result of an APC evaluation
-pub struct EvaluationResult {
+pub struct ApcPerformanceReport<S> {
     /// Statistics before optimizations, i.e., the sum of the AIR stats
     /// of all AIRs that *would* be involved in proving this basic block
     /// if it was run in software.
-    pub before: AirStats,
+    pub before: S,
     /// The AIR stats of the APC.
-    pub after: AirStats,
+    pub after: S,
+}
+
+impl<S: ApcStats> From<ApcPerformanceReport<S>> for ApcPerformanceReport<AirStats> {
+    fn from(report: ApcPerformanceReport<S>) -> Self {
+        ApcPerformanceReport {
+            before: report.before.into(),
+            after: report.after.into(),
+        }
+    }
+}
+
+pub trait ApcStats:
+    Into<AirStats> + Clone + Copy + Serialize + for<'a> Deserialize<'a> + Send + Sync + Sum
+{
+    fn cells_per_call(&self) -> usize;
+}
+
+impl<S: ApcStats> ApcPerformanceReport<S> {
+    pub fn cells_saved_per_call(&self) -> usize {
+        self.before.cells_per_call() - self.after.cells_per_call()
+    }
 }
 
 /// Evaluate an APC by comparing its cost to the cost of executing the
 /// basic block in software.
-pub fn evaluate_apc<IH>(
-    basic_block: &[IH::Instruction],
-    instruction_handler: &IH,
-    machine: &SymbolicMachine<impl Clone + Ord + std::fmt::Display>,
-) -> EvaluationResult
-where
-    IH: InstructionHandler,
-    IH::Field: Clone + Ord + std::fmt::Display,
-    IH::Instruction: Instruction<IH::Field>,
-{
-    let before = basic_block
+pub fn evaluate_apc<A: Adapter, Air: ApcArithmetization<A>>(
+    apc: Arc<AdapterApc<A>>,
+    instruction_handler: &A::InstructionHandler,
+    max_constraint_degree: usize,
+) -> ApcPerformanceReport<A::ApcStats> {
+    let before = apc
+        .block
+        .statements
         .iter()
-        .map(|instruction| instruction_handler.get_instruction_air_stats(instruction))
+        .map(|instruction| instruction_handler.get_instruction_air_metrics(instruction))
         .sum();
-    let after = AirStats::new(machine);
-    EvaluationResult { before, after }
+    let after = Air::get_metrics(apc, max_constraint_degree);
+    ApcPerformanceReport { before, after }
 }
 
-impl Display for EvaluationResult {
+impl<S: ApcStats> Display for ApcPerformanceReport<S> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let EvaluationResult { before, after } = self;
+        let ApcPerformanceReport { before, after } = self;
+        let before: AirStats = (*before).into();
+        let after: AirStats = (*after).into();
         write!(
             f,
             "APC advantage:\n  - Main columns: {}\n  - Bus interactions: {}\n  - Constraints: {}",

autoprecompiles/src/lib.rs

@@ -1,7 +1,6 @@
 use crate::adapter::{Adapter, AdapterApc, AdapterVmConfig};
 use crate::blocks::BasicBlock;
 use crate::bus_map::{BusMap, BusType};
-use crate::evaluation::AirStats;
 use crate::expression_conversion::algebraic_to_grouped_expression;
 use crate::symbolic_machine_generator::convert_machine_field_type;
 use expression::{AlgebraicExpression, AlgebraicReference};
@@ -300,6 +299,7 @@ pub trait InstructionHandler {
     type Field;
     type Instruction;
     type AirId;
+    type ApcStats;
 
     /// Returns the AIR for the given instruction.
     fn get_instruction_air_and_id(
@@ -308,7 +308,7 @@ pub trait InstructionHandler {
     ) -> (Self::AirId, &SymbolicMachine<Self::Field>);
 
     /// Returns the AIR stats for the given instruction.
-    fn get_instruction_air_stats(&self, instruction: &Self::Instruction) -> AirStats;
+    fn get_instruction_air_metrics(&self, instruction: &Self::Instruction) -> Self::ApcStats;
 
     /// Returns whether the given instruction is allowed in an autoprecompile.
     fn is_allowed(&self, instruction: &Self::Instruction) -> bool;