Skip to content

Share bytecode IR between CPU and GPU trace gen #3756

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
qwang98 wants to merge 1 commit into
base: powdr-labs/cpu-tracegen-opt
Choose a base branch
from
Draft

Share bytecode IR between CPU and GPU trace gen #3756

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
230 changes: 0 additions & 230 deletions autoprecompiles/src/expression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
//! In this module, we instantiate `powdr_expression::AlgebraicExpression` using a
//! custom `AlgebraicReference` type.
use core::ops::{Add, Mul, Neg, Sub};
use powdr_expression::{AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicUnaryOperator};
use powdr_number::ExpressionConvertible;
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, hash::Hash, marker::PhantomData, sync::Arc};
Expand Down Expand Up @@ -253,232 +252,3 @@ where
(*self.witness.get(&algebraic_var.id).unwrap()).into()
}
}

/// Pre-compiled expression for fast per-row evaluation without recursive AST walking.
pub enum CompiledExpr<F> {
/// Expression is a compile-time constant. Zero cost per row.
Constant(F),
/// Expression is a single column read: `row[col_idx]`. One array access per row.
DirectLoad(usize),
/// Expression is `constant + sum(coeff * row[col_idx]) + sum(coeff * row[a] * row[b])`.
/// Covers both linear (degree-1) and quadratic (degree-2) expressions.
/// Linear expressions have an empty `products` vec.
Polynomial(Polynomial<F>),
}

/// Degree ≤ 2 is sufficient under the default powdr-openvm config:
/// `DEFAULT_DEGREE_BOUND.bus_interactions = 2` (see `powdr_openvm::DEFAULT_DEGREE_BOUND`,
/// derived from `openvm-stark-sdk::DEFAULT_APP_LOG_BLOWUP`). The cli-openvm-riscv
/// binary never overrides it.
pub struct Polynomial<F> {
pub constant: F,
/// Linear terms: (col_idx, coefficient).
pub linear: Vec<(usize, F)>,
/// Quadratic terms: (col_idx_a, col_idx_b, coefficient).
pub products: Vec<(usize, usize, F)>,
}

impl<F> CompiledExpr<F>
where
F: Add<Output = F> + Sub<Output = F> + Mul<Output = F> + Neg<Output = F> + Copy + PartialEq,
{
/// Compile an expression at build time into a fast-eval form.
/// `id_to_idx` is a dense Vec indexed directly by poly ID.
/// `zero` and `one` are the field's additive and multiplicative identities.
pub fn compile(expr: &AlgebraicExpression<F>, id_to_idx: &[usize], zero: F, one: F) -> Self {
let d = Self::decompose(expr, id_to_idx, zero, one);
if d.linear.is_empty() && d.products.is_empty() {
CompiledExpr::Constant(d.constant)
} else if d.linear.len() == 1
&& d.products.is_empty()
&& d.constant == zero
&& d.linear[0].1 == one
{
CompiledExpr::DirectLoad(d.linear[0].0)
} else {
CompiledExpr::Polynomial(d)
}
}

/// Evaluate the compiled expression against a row slice.
#[inline(always)]
pub fn eval(&self, row: &[F]) -> F {
match self {
CompiledExpr::Constant(c) => *c,
CompiledExpr::DirectLoad(idx) => row[*idx],
CompiledExpr::Polynomial(Polynomial {
constant,
linear,
products,
}) => {
let lin = linear
.iter()
.fold(*constant, |acc, &(idx, coeff)| acc + coeff * row[idx]);
products
.iter()
.fold(lin, |acc, &(a, b, coeff)| acc + coeff * row[a] * row[b])
}
}
}

/// Decompose an expression into constant + linear + quadratic components.
/// Panics if the expression is degree > 2.
fn decompose(
expr: &AlgebraicExpression<F>,
id_to_idx: &[usize],
zero: F,
one: F,
) -> Polynomial<F> {
use powdr_expression::AlgebraicExpression as AE;
match expr {
AE::Number(c) => Polynomial {
constant: *c,
linear: vec![],
products: vec![],
},
AE::Reference(r) => Polynomial {
constant: zero,
linear: vec![(id_to_idx[r.id as usize], one)],
products: vec![],
},
AE::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) => {
let l = Self::decompose(left, id_to_idx, zero, one);
let r = Self::decompose(right, id_to_idx, zero, one);
match op {
AlgebraicBinaryOperator::Add => Polynomial {
constant: l.constant + r.constant,
linear: l.linear.into_iter().chain(r.linear).collect(),
products: l.products.into_iter().chain(r.products).collect(),
},
AlgebraicBinaryOperator::Sub => Polynomial {
constant: l.constant - r.constant,
linear: l
.linear
.into_iter()
.chain(r.linear.into_iter().map(|(i, c)| (i, -c)))
.collect(),
products: l
.products
.into_iter()
.chain(r.products.into_iter().map(|(a, b, c)| (a, b, -c)))
.collect(),
},
AlgebraicBinaryOperator::Mul => Self::mul_decompositions(l, r, zero),
}
}
AE::UnaryOperation(powdr_expression::AlgebraicUnaryOperation { op, expr }) => {
match op {
AlgebraicUnaryOperator::Minus => {
let d = Self::decompose(expr, id_to_idx, zero, one);
Polynomial {
constant: -d.constant,
linear: d.linear.into_iter().map(|(i, c)| (i, -c)).collect(),
products: d.products.into_iter().map(|(a, b, c)| (a, b, -c)).collect(),
}
}
}
}
}
}

/// Multiply two decompositions. Supports constant * anything and
/// linear * linear (producing quadratic products).
/// Panics if the result would be degree > 2.
fn mul_decompositions(l: Polynomial<F>, r: Polynomial<F>, zero: F) -> Polynomial<F> {
// If either side is purely constant, scale the other.
if l.linear.is_empty() && l.products.is_empty() {
return Polynomial {
constant: l.constant * r.constant,
linear: r
.linear
.into_iter()
.map(|(i, c)| (i, c * l.constant))
.collect(),
products: r
.products
.into_iter()
.map(|(a, b, c)| (a, b, c * l.constant))
.collect(),
};
}
if r.linear.is_empty() && r.products.is_empty() {
return Polynomial {
constant: l.constant * r.constant,
linear: l
.linear
.into_iter()
.map(|(i, c)| (i, c * r.constant))
.collect(),
products: l
.products
.into_iter()
.map(|(a, b, c)| (a, b, c * r.constant))
.collect(),
};
}
// Both sides have variable terms. Products of quadratic with anything → degree > 2.
assert!(
l.products.is_empty() && r.products.is_empty(),
"Bus interaction expression is degree > 2. This is unexpected."
);
// (c_l + sum(a_i * x_i)) * (c_r + sum(b_j * y_j))
// = c_l*c_r + c_l*sum(b_j*y_j) + c_r*sum(a_i*x_i) + sum(a_i*b_j * x_i*y_j)
let mut linear = Vec::new();
let mut products = Vec::new();

// c_l * linear_r
if l.constant != zero {
linear.extend(r.linear.iter().map(|&(i, c)| (i, c * l.constant)));
}
// c_r * linear_l
if r.constant != zero {
linear.extend(l.linear.iter().map(|&(i, c)| (i, c * r.constant)));
}
// linear_l * linear_r → quadratic products
for &(i, ci) in &l.linear {
for &(j, cj) in &r.linear {
products.push((i, j, ci * cj));
}
}

Polynomial {
constant: l.constant * r.constant,
linear,
products,
}
}
}

/// Pre-compiled bus interaction for fast per-row evaluation.
pub struct CompiledBusInteraction<F> {
pub id: u64,
pub mult: CompiledExpr<F>,
pub args: Vec<CompiledExpr<F>>,
}

impl<F> CompiledBusInteraction<F>
where
F: Add<Output = F> + Sub<Output = F> + Mul<Output = F> + Neg<Output = F> + Copy + PartialEq,
{
/// Compile all bus interactions from symbolic form.
/// `zero` and `one` are the field's additive and multiplicative identities.
pub fn compile_all(
interactions: &[SymbolicBusInteraction<F>],
id_to_idx: &[usize],
zero: F,
one: F,
) -> Vec<Self> {
interactions
.iter()
.map(|bi| CompiledBusInteraction {
id: bi.id,
mult: CompiledExpr::compile(&bi.mult, id_to_idx, zero, one),
args: bi
.args
.iter()
.map(|a| CompiledExpr::compile(a, id_to_idx, zero, one))
.collect(),
})
.collect()
}
}
Loading
Loading