-
Notifications
You must be signed in to change notification settings - Fork 82
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.
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
Prepare to call jit from block machine. #2098
Open
chriseth
wants to merge
27
commits into
main
Choose a base branch
from
call_jit_from_block
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+305
−65
Open
Changes from all commits
Commits
Show all changes
27 commits
Select commit
Hold shift + click to select a range
fbdabca
Use contiguous data for finalized rows.
chriseth 71d2ebd
Remove last row.
chriseth f66c5c2
Fix extend.
chriseth 934ebb1
clippy
chriseth 19d544f
remove unused functiion
chriseth 6b83811
Allow non-sorted column IDs.
chriseth 80c24f9
Jit driver and machine.
chriseth 5528a4f
reserve block
chriseth 4d4111c
revie
chriseth 0076d2c
review
chriseth 3c577c1
review
chriseth 4e83830
Add "remove_last_row"
chriseth 9f3130a
Re-added one case.
chriseth 1aadfe0
Merge branch 'improve_finalizable' into call_jit_from_block
chriseth 056a904
jitedijit
chriseth f2fee46
jitjit
chriseth 9415cc3
Remove try_remove_last_row.
chriseth ec5d110
compact data ref.
chriseth 7cbbb8d
Remove unwrap.
chriseth 4490dcd
Cleanup
chriseth 7ffdad9
Merge remote-tracking branch 'origin/improve_finalizable' into call_j…
chriseth 09863e8
cleanup
chriseth cf51774
cleanup
chriseth 825d0eb
Merge remote-tracking branch 'origin/main' into call_jit_from_block
chriseth d470fd7
Remove type constraints.
chriseth b8864c8
Avoid dyn iter.
chriseth ca38415
Correct typo.
chriseth File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,6 +3,7 @@ use std::{ | |
ops::{Index, IndexMut}, | ||
}; | ||
|
||
use auto_enums::auto_enum; | ||
use bit_vec::BitVec; | ||
use itertools::Itertools; | ||
use powdr_ast::analyzed::{PolyID, PolynomialType}; | ||
|
@@ -13,7 +14,7 @@ use crate::witgen::rows::Row; | |
/// Sequence of rows of field elements, stored in a compact form. | ||
/// Optimized for contiguous column IDs, but works with any combination. | ||
#[derive(Clone)] | ||
struct CompactData<T: FieldElement> { | ||
pub struct CompactData<T> { | ||
/// The ID of the first column used in the table. | ||
first_column_id: u64, | ||
/// The length of a row in the table. | ||
|
@@ -26,7 +27,7 @@ struct CompactData<T: FieldElement> { | |
|
||
impl<T: FieldElement> CompactData<T> { | ||
/// Creates a new empty compact data storage. | ||
fn new(column_ids: &[PolyID]) -> Self { | ||
pub fn new(column_ids: &[PolyID]) -> Self { | ||
let col_id_range = column_ids.iter().map(|id| id.id).minmax(); | ||
let (first_column_id, last_column_id) = col_id_range.into_option().unwrap(); | ||
Self { | ||
|
@@ -37,28 +38,28 @@ impl<T: FieldElement> CompactData<T> { | |
} | ||
} | ||
|
||
fn is_empty(&self) -> bool { | ||
pub fn is_empty(&self) -> bool { | ||
self.data.is_empty() | ||
} | ||
|
||
/// Returns the number of stored rows. | ||
fn len(&self) -> usize { | ||
pub fn len(&self) -> usize { | ||
self.data.len() / self.column_count | ||
} | ||
|
||
/// Truncates the data to `len` rows. | ||
fn truncate(&mut self, len: usize) { | ||
pub fn truncate(&mut self, len: usize) { | ||
self.data.truncate(len * self.column_count); | ||
self.known_cells.truncate(len * self.column_count); | ||
} | ||
|
||
fn clear(&mut self) { | ||
pub fn clear(&mut self) { | ||
self.data.clear(); | ||
self.known_cells.clear(); | ||
} | ||
|
||
/// Appends a non-finalized row to the data, turning it into a finalized row. | ||
fn push(&mut self, row: Row<T>) { | ||
pub fn push(&mut self, row: Row<T>) { | ||
self.data.reserve(self.column_count); | ||
self.known_cells.reserve(self.column_count); | ||
for col_id in self.first_column_id..(self.first_column_id + self.column_count as u64) { | ||
|
@@ -75,11 +76,69 @@ impl<T: FieldElement> CompactData<T> { | |
} | ||
} | ||
|
||
fn get(&self, row: usize, col: u64) -> (T, bool) { | ||
pub fn append_new_rows(&mut self, count: usize) { | ||
self.data | ||
.resize(self.data.len() + count * self.column_count, T::zero()); | ||
self.known_cells | ||
.grow(self.known_cells.len() + count * self.column_count, false); | ||
} | ||
|
||
fn index(&self, row: usize, col: u64) -> usize { | ||
let col = col - self.first_column_id; | ||
let idx = row * self.column_count + col as usize; | ||
row * self.column_count + col as usize | ||
} | ||
|
||
pub fn get(&self, row: usize, col: u64) -> (T, bool) { | ||
let idx = self.index(row, col); | ||
(self.data[idx], self.known_cells[idx]) | ||
} | ||
|
||
pub fn set(&mut self, row: usize, col: u64, value: T) { | ||
let idx = self.index(row, col); | ||
assert!(!self.known_cells[idx] || self.data[idx] == value); | ||
self.data[idx] = value; | ||
self.known_cells.set(idx, true); | ||
} | ||
|
||
pub fn known_values_in_row(&self, row: usize) -> impl Iterator<Item = (u64, &T)> { | ||
(0..self.column_count).filter_map(move |i| { | ||
let idx = row * self.column_count + i; | ||
self.known_cells[idx].then(|| { | ||
let col_id = self.first_column_id + i as u64; | ||
(col_id, &self.data[idx]) | ||
}) | ||
}) | ||
} | ||
} | ||
|
||
/// A mutable reference into CompactData that is meant to be used | ||
/// only for a certain block of rows, starting from row index zero. | ||
/// It allows negative row indices as well. | ||
pub struct CompactDataRef<'a, T> { | ||
data: &'a mut CompactData<T>, | ||
row_offset: usize, | ||
} | ||
|
||
impl<'a, T: FieldElement> CompactDataRef<'a, T> { | ||
/// Creates a new reference to the data, supplying the offset of the row | ||
/// that is supposed to be "row zero". | ||
pub fn new(data: &'a mut CompactData<T>, row_offset: usize) -> Self { | ||
Self { data, row_offset } | ||
} | ||
|
||
pub fn get(&self, row: i32, col: u32) -> T { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Not sure if the 32 bit values here provide a performance advantage. Any opinions? |
||
let (v, known) = self.data.get(self.inner_row(row), col as u64); | ||
assert!(known); | ||
v | ||
} | ||
|
||
pub fn set(&mut self, row: i32, col: u32, value: T) { | ||
self.data.set(self.inner_row(row), col as u64, value); | ||
} | ||
|
||
fn inner_row(&self, row: i32) -> usize { | ||
(row + self.row_offset as i32) as usize | ||
} | ||
} | ||
|
||
/// A data structure that stores witness data. | ||
|
@@ -215,6 +274,41 @@ impl<T: FieldElement> FinalizableData<T> { | |
} | ||
} | ||
|
||
/// Returns an iterator over the values known in that row together with the PolyIDs. | ||
#[auto_enum(Iterator)] | ||
pub fn known_values_in_row( | ||
&self, | ||
row: usize, | ||
) -> impl Iterator<Item = (PolyID, T)> + use<'_, T> { | ||
match self.location_of_row(row) { | ||
Location::PreFinalized(local) => { | ||
let row = &self.pre_finalized_data[local]; | ||
self.column_ids | ||
.iter() | ||
.filter_map(move |id| row.value(id).map(|v| (*id, v))) | ||
} | ||
Location::Finalized(local) => { | ||
self.finalized_data | ||
.known_values_in_row(local) | ||
.map(|(id, v)| { | ||
( | ||
PolyID { | ||
id, | ||
ptype: PolynomialType::Committed, | ||
}, | ||
*v, | ||
) | ||
}) | ||
} | ||
Location::PostFinalized(local) => { | ||
let row = &self.post_finalized_data[local]; | ||
self.column_ids | ||
.iter() | ||
.filter_map(move |id| row.value(id).map(|v| (*id, v))) | ||
} | ||
} | ||
} | ||
|
||
pub fn last(&self) -> Option<&Row<T>> { | ||
match self.location_of_last_row()? { | ||
Location::PreFinalized(local) => self.pre_finalized_data.get(local), | ||
|
@@ -283,6 +377,18 @@ impl<T: FieldElement> FinalizableData<T> { | |
} | ||
} | ||
|
||
/// Appends a given amount of new finalized rows set to zero and "unknown". | ||
/// Returns a `CompactDataRef` that is built so that its "row zero" is the | ||
/// first newly appended row. | ||
/// | ||
/// Panics if there are any non-finalized rows at the end. | ||
pub fn append_new_finalized_rows(&mut self, count: usize) -> CompactDataRef<'_, T> { | ||
assert!(self.post_finalized_data.is_empty()); | ||
let row_zero = self.finalized_data.len(); | ||
self.finalized_data.append_new_rows(count); | ||
CompactDataRef::new(&mut self.finalized_data, row_zero) | ||
} | ||
|
||
/// Takes all data out of the [FinalizableData] and returns it as a list of columns. | ||
/// Columns are represented as a tuple of: | ||
/// - A list of values | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,63 @@ | ||
use bit_vec::BitVec; | ||
use powdr_number::FieldElement; | ||
|
||
use crate::witgen::{ | ||
data_structures::finalizable_data::CompactDataRef, | ||
machines::{LookupCell, MachineParts}, | ||
util::try_to_simple_poly, | ||
EvalError, FixedData, MutableState, QueryCallback, | ||
}; | ||
|
||
pub struct JitProcessor<'a, T: FieldElement> { | ||
_fixed_data: &'a FixedData<'a, T>, | ||
parts: MachineParts<'a, T>, | ||
_block_size: usize, | ||
latch_row: usize, | ||
} | ||
|
||
impl<'a, T: FieldElement> JitProcessor<'a, T> { | ||
pub fn new( | ||
fixed_data: &'a FixedData<'a, T>, | ||
parts: MachineParts<'a, T>, | ||
block_size: usize, | ||
latch_row: usize, | ||
) -> Self { | ||
JitProcessor { | ||
_fixed_data: fixed_data, | ||
parts, | ||
_block_size: block_size, | ||
latch_row, | ||
} | ||
} | ||
|
||
pub fn can_answer_lookup(&self, _identity_id: u64, _known_inputs: &BitVec) -> bool { | ||
// TODO call the JIT compiler here. | ||
false | ||
} | ||
|
||
pub fn process_lookup_direct<'b, 'c, 'd, Q: QueryCallback<T>>( | ||
&self, | ||
_mutable_state: &'b mut MutableState<'a, 'b, T, Q>, | ||
connection_id: u64, | ||
values: Vec<LookupCell<'c, T>>, | ||
mut data: CompactDataRef<'d, T>, | ||
) -> Result<bool, EvalError<T>> { | ||
// Transfer inputs. | ||
let right = self.parts.connections[&connection_id].right; | ||
for (e, v) in right.expressions.iter().zip(&values) { | ||
match v { | ||
LookupCell::Input(&v) => { | ||
let col = try_to_simple_poly(e).unwrap(); | ||
data.set(self.latch_row as i32, col.poly_id.id as u32, v); | ||
} | ||
LookupCell::Output(_) => {} | ||
} | ||
} | ||
|
||
// Just some code here to avoid "unused" warnings. | ||
// This code will not be called as long as `can_answer_lookup` returns false. | ||
data.get(self.latch_row as i32, 0); | ||
|
||
unimplemented!(); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
pub mod jit_processor; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this mean that we sometimes set the same value twice?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Or that there's a default value?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It's OK to set the value twice if it's the same value. We can maybe make this more strict later, though.
The "default" value is "not known", internally it will be represented by zero.