diff --git a/autoprecompiles/src/expression.rs b/autoprecompiles/src/expression.rs index d50bdc0544..ab1683a6bf 100644 --- a/autoprecompiles/src/expression.rs +++ b/autoprecompiles/src/expression.rs @@ -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}; @@ -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 { - /// 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), -} - -/// 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 { - 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 CompiledExpr -where - F: Add + Sub + Mul + Neg + 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, 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, - id_to_idx: &[usize], - zero: F, - one: F, - ) -> Polynomial { - 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, r: Polynomial, zero: F) -> Polynomial { - // 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 { - pub id: u64, - pub mult: CompiledExpr, - pub args: Vec>, -} - -impl CompiledBusInteraction -where - F: Add + Sub + Mul + Neg + 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], - id_to_idx: &[usize], - zero: F, - one: F, - ) -> Vec { - 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() - } -} diff --git a/openvm/src/bytecode.rs b/openvm/src/bytecode.rs new file mode 100644 index 0000000000..d6629da592 --- /dev/null +++ b/openvm/src/bytecode.rs @@ -0,0 +1,223 @@ +//! Stack-machine bytecode for algebraic expression evaluation. +//! +//! Shared between the GPU kernel (`openvm/cuda/src/expr_eval.cuh`) and the CPU +//! trace generator. The bytecode is a flat `Vec` of opcodes interleaved +//! with inline operands; `ExprSpan { off, len }` indexes contiguous expression +//! programs within the buffer. +//! +//! The emitter (`emit_expr`, `compile_bus_to_bytecode`) is field-generic. The +//! CPU evaluator (`eval_expr`) requires `PrimeField32` so it can reconstruct +//! field elements from the inline `u32` constants. + +use std::collections::BTreeMap; + +use openvm_stark_backend::p3_field::{integers::QuotientMap, PrimeField32}; +use powdr_autoprecompiles::{ + expression::AlgebraicExpression, symbolic_machine::SymbolicBusInteraction, +}; +use powdr_expression::{AlgebraicBinaryOperator, AlgebraicUnaryOperator}; + +/// Stack-machine opcode tags. Each opcode word is followed by 0 or 1 operand +/// words: `PushApc`/`PushConst` consume one, all others consume zero. +#[repr(u32)] +pub enum OpCode { + /// Push `apc_trace[base + row]` onto the stack. Operand: `base` (column + /// offset for GPU column-major layout, `col_idx` directly for CPU when + /// `apc_height = 1`). + PushApc = 0, + /// Push the inline constant onto the stack. Operand: field element as `u32`. + PushConst = 1, + /// `a + b` where `b` is the top, `a` is below. + Add = 2, + /// `a - b` where `b` is the top, `a` is below. + Sub = 3, + /// `a * b`. + Mul = 4, + /// Negate the top of the stack. + Neg = 5, + /// Replace top with its inverse, or zero if it was zero. + InvOrZero = 6, +} + +/// `(offset, length)` of a compiled expression program inside a shared +/// bytecode buffer. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ExprSpan { + /// Offset (in `u32` words) into the bytecode buffer where this expression + /// program starts. + pub off: u32, + /// Length (in `u32` words) of the expression program. + pub len: u32, +} + +/// Per-bus-interaction metadata produced by `compile_bus_to_bytecode`. +/// `#[repr(C)]` so it can also be passed directly to the GPU kernel. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct BusMeta { + /// Bus id this interaction targets (matches periphery chip bus id). + pub bus_id: u32, + /// Number of argument expressions for this interaction. + pub num_args: u32, + /// Starting index into the `ExprSpan` array for this interaction. The + /// layout per interaction is `[multiplicity_span, arg0, arg1, ...]`. + pub args_index_off: u32, +} + +/// Emit bytecode for `expr` into `bc`. References use +/// `id_to_apc_index[&ref.id] * apc_height` as the inline `PushApc` operand. +/// On GPU, `apc_height = trace_height` (column-major byte offset); on CPU, +/// `apc_height = 1` so the operand is just the column index. +pub fn emit_expr( + bc: &mut Vec, + expr: &AlgebraicExpression, + id_to_apc_index: &BTreeMap, + apc_height: usize, +) { + match expr { + AlgebraicExpression::Number(c) => { + bc.push(OpCode::PushConst as u32); + bc.push(c.as_canonical_u32()); + } + AlgebraicExpression::Reference(r) => { + let idx = (id_to_apc_index[&r.id] * apc_height) as u32; + bc.push(OpCode::PushApc as u32); + bc.push(idx); + } + AlgebraicExpression::UnaryOperation(u) => { + emit_expr(bc, &u.expr, id_to_apc_index, apc_height); + match u.op { + AlgebraicUnaryOperator::Minus => bc.push(OpCode::Neg as u32), + } + } + AlgebraicExpression::BinaryOperation(b) => { + emit_expr(bc, &b.left, id_to_apc_index, apc_height); + emit_expr(bc, &b.right, id_to_apc_index, apc_height); + match b.op { + AlgebraicBinaryOperator::Add => bc.push(OpCode::Add as u32), + AlgebraicBinaryOperator::Sub => bc.push(OpCode::Sub as u32), + AlgebraicBinaryOperator::Mul => bc.push(OpCode::Mul as u32), + } + } + } +} + +/// Append bytecode for `expr` and return its `ExprSpan`. +pub fn emit_expr_span( + bc: &mut Vec, + expr: &AlgebraicExpression, + id_to_apc_index: &BTreeMap, + apc_height: usize, +) -> ExprSpan { + let off = bc.len() as u32; + emit_expr(bc, expr, id_to_apc_index, apc_height); + let len = (bc.len() as u32) - off; + ExprSpan { off, len } +} + +/// Compile a list of symbolic bus interactions into shared bytecode. +/// +/// Returns `(metadata, spans, bytecode)`. The `spans` array is laid out per +/// interaction as `[multiplicity_span, arg0_span, ..., argN_span]`; an +/// interaction's `args_index_off` field indexes the start of its run. +pub fn compile_bus_to_bytecode( + bus_interactions: &[SymbolicBusInteraction], + apc_poly_id_to_index: &BTreeMap, + apc_height: usize, +) -> (Vec, Vec, Vec) { + let mut interactions = Vec::with_capacity(bus_interactions.len()); + let mut arg_spans = Vec::new(); + let mut bytecode = Vec::new(); + + for bus_interaction in bus_interactions { + let args_index_off = arg_spans.len() as u32; + let mult_span = emit_expr_span( + &mut bytecode, + &bus_interaction.mult, + apc_poly_id_to_index, + apc_height, + ); + arg_spans.push(mult_span); + + for arg in &bus_interaction.args { + let span = emit_expr_span(&mut bytecode, arg, apc_poly_id_to_index, apc_height); + arg_spans.push(span); + } + + interactions.push(BusMeta { + bus_id: bus_interaction.id as u32, + num_args: bus_interaction.args.len() as u32, + args_index_off, + }); + } + + (interactions, arg_spans, bytecode) +} + +const STACK_CAPACITY: usize = 16; + +/// Evaluate the bytecode program at `[span.off, span.off + span.len)` against +/// a row slice. Mirrors `eval_expr` in `openvm/cuda/src/expr_eval.cuh`. +/// +/// SAFETY: assumes the bytecode is well-formed (balanced stack effects, no +/// out-of-range row index, stack depth bounded by `STACK_CAPACITY`). The +/// emitter guarantees all three for any bus interaction up to degree 2. +#[inline(always)] +pub fn eval_expr>( + bytecode: &[u32], + span: ExprSpan, + row: &[F], +) -> F { + let mut stack: [F; STACK_CAPACITY] = [F::ZERO; STACK_CAPACITY]; + let mut sp: usize = 0; + let mut ip = span.off as usize; + let end = (span.off + span.len) as usize; + while ip < end { + let op = bytecode[ip]; + ip += 1; + unsafe { + match op { + x if x == OpCode::PushApc as u32 => { + let base = bytecode[ip] as usize; + ip += 1; + *stack.get_unchecked_mut(sp) = *row.get_unchecked(base); + sp += 1; + } + x if x == OpCode::PushConst as u32 => { + let u = bytecode[ip]; + ip += 1; + // SAFETY: emitter wrote `F::as_canonical_u32(c)`, which is + // guaranteed canonical (< p), satisfying the precondition + // of `from_canonical_unchecked`. + *stack.get_unchecked_mut(sp) = F::from_canonical_unchecked(u); + sp += 1; + } + x if x == OpCode::Add as u32 => { + let r = *stack.get_unchecked(sp - 1); + let l = *stack.get_unchecked(sp - 2); + *stack.get_unchecked_mut(sp - 2) = l + r; + sp -= 1; + } + x if x == OpCode::Sub as u32 => { + let r = *stack.get_unchecked(sp - 1); + let l = *stack.get_unchecked(sp - 2); + *stack.get_unchecked_mut(sp - 2) = l - r; + sp -= 1; + } + x if x == OpCode::Mul as u32 => { + let r = *stack.get_unchecked(sp - 1); + let l = *stack.get_unchecked(sp - 2); + *stack.get_unchecked_mut(sp - 2) = l * r; + sp -= 1; + } + x if x == OpCode::Neg as u32 => { + let v = *stack.get_unchecked(sp - 1); + *stack.get_unchecked_mut(sp - 1) = -v; + } + _ => debug_assert!(false, "Unknown opcode {op}"), + } + } + } + stack[0] +} diff --git a/openvm/src/cuda_abi.rs b/openvm/src/cuda_abi.rs index ee09e09306..59fa49a38d 100644 --- a/openvm/src/cuda_abi.rs +++ b/openvm/src/cuda_abi.rs @@ -134,39 +134,8 @@ pub fn apc_apply_derived_expr( } } -/// OpCode enum for the GPU stack machine bus evaluator. -#[repr(u32)] -pub enum OpCode { - PushApc = 0, // Push the APC value onto the stack. Must be followed by the index of the value in the APC device buffer. - PushConst = 1, // Push a constant value onto the stack. Must be followed by the constant value. - Add = 2, // Add the top two values on the stack. - Sub = 3, // Subtract the top two values on the stack. - Mul = 4, // Multiply the top two values on the stack. - Neg = 5, // Negate the top value on the stack. - InvOrZero = 6, // Invert the top value on the stack if it is not zero, otherwise pop and push zero. -} - -/// GPU device representation of a bus interaction. -#[repr(C)] -#[derive(Clone, Copy)] -pub struct DevInteraction { - /// Bus id this interaction targets (matches periphery chip bus id) - pub bus_id: u32, - /// Number of argument expressions for this interaction - pub num_args: u32, - /// Starting index into the `ExprSpan` array for this interaction's args - /// Layout: [ multiplicity span, arg0, arg1, ... ] - pub args_index_off: u32, -} - -#[repr(C)] -#[derive(Clone, Copy)] -pub struct ExprSpan { - /// Offset (in u32 words) into `bytecode` where this arg expression starts - pub off: u32, - /// Length (instruction count) of this arg expression - pub len: u32, -} +// Bytecode IR shared with the CPU evaluator; see `crate::bytecode`. +pub use crate::bytecode::{BusMeta as DevInteraction, ExprSpan, OpCode}; /// High-level safe wrapper for `_apc_apply_bus`. Applies bus interactions on the GPU, /// updating periphery histograms in-place. diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs index c18f116c86..ac7cad1536 100644 --- a/openvm/src/lib.rs +++ b/openvm/src/lib.rs @@ -51,6 +51,7 @@ use crate::extraction_utils::{get_air_metrics, AirWidths, OriginalVmConfig}; use crate::powdr_extension::{PowdrExtensionExecutor, PowdrPrecompile}; mod air_builder; +pub mod bytecode; pub mod cuda_abi; pub mod empirical_constraints; pub mod extraction_utils; diff --git a/openvm/src/powdr_extension/trace_generator/cpu/mod.rs b/openvm/src/powdr_extension/trace_generator/cpu/mod.rs index 204d63a371..2d7b640772 100644 --- a/openvm/src/powdr_extension/trace_generator/cpu/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/cpu/mod.rs @@ -153,24 +153,15 @@ impl PowdrTraceGeneratorCpu { &self.apc, ); - // Build dense Vec indexed by poly ID for O(1) column lookups in the hot loop. - // Poly IDs may be sparse (gaps between IDs), so the Vec is sized to max_id + 1. let width = apc_poly_id_to_index.len(); - let max_poly_id = apc_poly_id_to_index.keys().last().copied().unwrap_or(0) as usize; - let apc_poly_id_to_index: Vec = (0..=max_poly_id) - .map(|id| apc_poly_id_to_index.get(&(id as u64)).copied().unwrap_or(0)) - .collect(); - // Compile bus interactions once before the hot loop - let compiled_interactions = { - use powdr_autoprecompiles::expression::CompiledBusInteraction; - CompiledBusInteraction::compile_all( - &self.apc.machine().bus_interactions, - &apc_poly_id_to_index, - BabyBear::ZERO, - BabyBear::ONE, - ) - }; + // Compile bus interactions to shared bytecode (apc_height = 1 so the + // PushApc operand is the column index directly). + let (bus_meta, arg_spans, bytecode) = crate::bytecode::compile_bus_to_bytecode( + &self.apc.machine().bus_interactions, + &apc_poly_id_to_index, + 1, + ); // allocate for apc trace let height = next_power_of_two_or_zero(num_apc_calls); @@ -198,35 +189,45 @@ impl PowdrTraceGeneratorCpu { // Fill in the columns we have to compute from other columns // (these are either new columns or for example the "is_valid" column). for derived_column in columns_to_compute { - let col_index = apc_poly_id_to_index[derived_column.variable.id as usize]; + let col_index = apc_poly_id_to_index[&derived_column.variable.id]; row_slice[col_index] = match &derived_column.computation_method { ComputationMethod::Constant(c) => *c, ComputationMethod::QuotientOrZero(e1, e2) => { use powdr_number::ExpressionConvertible; let divisor_val = e2.to_expression(&|n| *n, &|column_ref| { - row_slice[apc_poly_id_to_index[column_ref.id as usize]] + row_slice[apc_poly_id_to_index[&column_ref.id]] }); if divisor_val.is_zero() { BabyBear::ZERO } else { divisor_val.inverse() * e1.to_expression(&|n| *n, &|column_ref| { - row_slice[apc_poly_id_to_index[column_ref.id as usize]] + row_slice[apc_poly_id_to_index[&column_ref.id]] }) } } }; } - // Evaluate bus interactions using compiled expressions. - // Periphery chips use AtomicU32 counters — thread-safe. - for ci in &compiled_interactions { - let mult = ci.mult.eval(row_slice); + // Evaluate bus interactions via shared stack-machine interpreter + // (same bytecode IR as the GPU kernel in `expr_eval.cuh`). + for meta in &bus_meta { + let base = meta.args_index_off as usize; + let mult_span = arg_spans[base]; + let mult = crate::bytecode::eval_expr(&bytecode, mult_span, row_slice); + let num_args = meta.num_args as usize; periphery_real.apply( - ci.id as u16, + meta.bus_id as u16, mult.as_canonical_u32(), - ci.args.iter().map(|a| a.eval(row_slice).as_canonical_u32()), + (0..num_args).map(|i| { + crate::bytecode::eval_expr( + &bytecode, + arg_spans[base + 1 + i], + row_slice, + ) + .as_canonical_u32() + }), periphery_bus_ids, ); } diff --git a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs index 9464a452f3..c51360538c 100644 --- a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs @@ -21,7 +21,8 @@ use powdr_constraint_solver::constraint_system::{ComputationMethod, DerivedVaria use powdr_expression::{AlgebraicBinaryOperator, AlgebraicUnaryOperator}; use crate::{ - cuda_abi::{self, DerivedExprSpec, DevInteraction, ExprSpan, OpCode, OriginalAir, Subst}, + bytecode::{compile_bus_to_bytecode, emit_expr, ExprSpan, OpCode}, + cuda_abi::{self, DerivedExprSpec, OriginalAir, Subst}, extraction_utils::{OriginalAirs, OriginalVmConfig}, isa::{IsaApc, OpenVmISA}, powdr_extension::{chip::PowdrChipGpu, executor::OriginalArenas}, @@ -36,66 +37,6 @@ pub use periphery::{ PowdrPeripheryInstancesGpu, SharedPeripheryChipsGpu, SharedPeripheryChipsGpuProverExt, }; -/// Encodes an algebraic expression into GPU stack-machine bytecode. -/// -/// Appends instructions to `bc` representing `expr` using the opcodes in `OpCode`. -/// References are encoded as `PushApc` with a column-major offset computed from -/// `id_to_apc_index` and `apc_height` (offset = apc_col_index * apc_height). -/// Constants are encoded as `PushConst` followed by the field element as `u32`. -/// Unary minus and binary operations map to `Neg`, `Add`, `Sub`, and `Mul`. -/// -/// Note: This function does not track or enforce the evaluation stack depth, -/// which is done in device code. -fn emit_expr( - bc: &mut Vec, - expr: &AlgebraicExpression, - id_to_apc_index: &BTreeMap, - apc_height: usize, -) { - match expr { - AlgebraicExpression::Number(c) => { - bc.push(OpCode::PushConst as u32); - bc.push(c.as_canonical_u32()); - } - AlgebraicExpression::Reference(r) => { - let idx = (id_to_apc_index[&r.id] * apc_height) as u32; - bc.push(OpCode::PushApc as u32); - bc.push(idx); - } - AlgebraicExpression::UnaryOperation(u) => { - emit_expr(bc, &u.expr, id_to_apc_index, apc_height); - match u.op { - AlgebraicUnaryOperator::Minus => bc.push(OpCode::Neg as u32), - } - } - AlgebraicExpression::BinaryOperation(b) => { - emit_expr(bc, &b.left, id_to_apc_index, apc_height); - emit_expr(bc, &b.right, id_to_apc_index, apc_height); - match b.op { - AlgebraicBinaryOperator::Add => bc.push(OpCode::Add as u32), - AlgebraicBinaryOperator::Sub => bc.push(OpCode::Sub as u32), - AlgebraicBinaryOperator::Mul => bc.push(OpCode::Mul as u32), - } - } - } -} - -/// Given the current bytecode, appends bytecode for the expression `expr` and returns the associated span -fn emit_expr_span( - bc: &mut Vec, - expr: &AlgebraicExpression, - id_to_apc_index: &BTreeMap, - apc_height: usize, -) -> ExprSpan { - // The span starts where the bytecode currently ends - let off = bc.len() as u32; - // Append the bytecode for `expr` - emit_expr(bc, expr, id_to_apc_index, apc_height); - // Calculate the length of the span - let len = (bc.len() as u32) - off; - ExprSpan { off, len } -} - /// Compile derived columns to GPU bytecode according to input order. fn compile_derived_to_gpu( derived_columns: &[DerivedVariable< @@ -140,41 +81,8 @@ fn compile_derived_to_gpu( (specs, bytecode) } -pub fn compile_bus_to_gpu( - bus_interactions: &[SymbolicBusInteraction], - apc_poly_id_to_index: &BTreeMap, - apc_height: usize, -) -> (Vec, Vec, Vec) { - let mut interactions = Vec::with_capacity(bus_interactions.len()); - let mut arg_spans = Vec::new(); - let mut bytecode = Vec::new(); - - for bus_interaction in bus_interactions { - // multiplicity as first arg span - let args_index_off = arg_spans.len() as u32; - let mult_span = emit_expr_span( - &mut bytecode, - &bus_interaction.mult, - apc_poly_id_to_index, - apc_height, - ); - arg_spans.push(mult_span); - - // args - for arg in &bus_interaction.args { - let span = emit_expr_span(&mut bytecode, arg, apc_poly_id_to_index, apc_height); - arg_spans.push(span); - } - - interactions.push(DevInteraction { - bus_id: (bus_interaction.id as u32), - num_args: bus_interaction.args.len() as u32, - args_index_off, - }); - } - - (interactions, arg_spans, bytecode) -} +/// Compile bus interactions to GPU bytecode (column-major encoding of refs). +pub use crate::bytecode::compile_bus_to_bytecode as compile_bus_to_gpu; pub struct PowdrTraceGeneratorGpu { pub apc: IsaApc,