Skip to content

Commit

Permalink
cross layer relay
Browse files Browse the repository at this point in the history
  • Loading branch information
@siq1
siq1 committed Nov 27, 2024
1 parent 9137c18 commit a06d6f6
Show file tree
Hide file tree
Showing 4 changed files with 107 additions and 42 deletions.
12 changes: 10 additions & 2 deletions expander_compiler/src/circuit/layered/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -383,15 +383,23 @@ impl<C: Config> Circuit<C> {
.iter()
.enumerate()
{
if i == l {
// if this is also the global input, it's always initialized
continue;
}
for j in 0..*len {
if input_mask[self.layer_ids[i]][l][j]
&& !output_mask[self.layer_ids[i - 1 - l]][j]
{
println!("{:?}", self.segments[17]);
println!("{:?}", self.segments[18]);
return Err(Error::InternalError(format!(
"circuit {} input {} not initialized by circuit {} output",
"circuit {} (layer {}) input {} not initialized by circuit {} (layer {}) output",
self.layer_ids[i],
i,
j,
self.layer_ids[i - 1 - l]
self.layer_ids[i - 1 - l],
i - 1 - l
)));
}
}
Expand Down
49 changes: 49 additions & 0 deletions expander_compiler/src/layering/compile.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,12 @@ pub struct IrContext<'a, C: Config> {
// it includes only variables mentioned in instructions, so internal variables in sub circuits are ignored here.
pub min_layer: Vec<usize>,
pub max_layer: Vec<usize>,
pub occured_layers: Vec<Vec<usize>>,
pub output_layer: usize,

// for each layer i, the minimum layer j that there exists gate j->i
pub min_used_layer: Vec<usize>,

pub output_order: HashMap<usize, usize>, // outputOrder[x] == y -> x is the y-th output

pub sub_circuit_loc_map: HashMap<usize, usize>,
Expand Down Expand Up @@ -161,7 +165,9 @@ impl<'a, C: Config> CompileContext<'a, C> {
num_sub_circuits: ns,
min_layer: Vec::new(),
max_layer: Vec::new(),
occured_layers: Vec::new(),
output_layer: 0,
min_used_layer: Vec::new(),
output_order: HashMap::new(),
sub_circuit_loc_map: HashMap::new(),
sub_circuit_insn_ids: Vec::new(),
Expand Down Expand Up @@ -433,6 +439,49 @@ impl<'a, C: Config> CompileContext<'a, C> {
}
}

// compute occured layers
ic.occured_layers = vec![Vec::new(); ic.max_layer.len()];
let outputs_set: HashSet<usize> = circuit.outputs.iter().cloned().collect();
for x in q.iter().cloned() {
let mut tmp = Vec::with_capacity(out_edges[x].len() + 1);
tmp.push(ic.min_layer[x]);
for y in out_edges[x].iter().cloned() {
tmp.push(ic.min_layer[y] - layer_advance[y]);
}
if outputs_set.contains(&x) {
tmp.push(ic.output_layer);
}
tmp.sort();
let mut tmp2 = Vec::with_capacity(tmp.len());
for &v in tmp.iter() {
if tmp2.is_empty() || *tmp2.last().unwrap() != v {
tmp2.push(v);
}
}
assert_eq!(tmp2[0], ic.min_layer[x]);
assert_eq!(*tmp2.last().unwrap(), ic.max_layer[x]);
ic.occured_layers[x] = tmp2;
}

// compute minUsedLayer
ic.min_used_layer = Vec::with_capacity(ic.output_layer + 1);
ic.min_used_layer.push(0);
ic.min_used_layer.extend(0..ic.output_layer);
for (i, sc) in ic.sub_circuit_insn_refs.iter().enumerate() {
let sub_circuit = &self.circuits[&sc.sub_circuit_id];
let input_layer = ic.sub_circuit_start_layer[i];
for j in 0..=sub_circuit.output_layer {
ic.min_used_layer[j + input_layer] = ic.min_used_layer[j + input_layer]
.min(sub_circuit.min_used_layer[j] + input_layer);
}
}
for x in q.iter().cloned() {
let t = &ic.occured_layers[x];
for (u, v) in t.iter().zip(t.iter().skip(1)) {
ic.min_used_layer[*v] = ic.min_used_layer[*v].min(*u);
}
}

self.circuits.insert(circuit_id, ic);
}
}
5 changes: 4 additions & 1 deletion expander_compiler/src/layering/layer_layout.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,10 @@ impl<'a, C: Config> CompileContext<'a, C> {
ic.lcs[ic.output_layer].vars.add(v);
}
for i in 1..ic.num_var {
for j in ic.min_layer[i]..=ic.max_layer[i] {
/*for j in ic.min_layer[i]..=ic.max_layer[i] {
ic.lcs[j].vars.add(&i);
}*/
for j in ic.occured_layers[i].iter().cloned() {
ic.lcs[j].vars.add(&i);
}
}
Expand Down
83 changes: 44 additions & 39 deletions expander_compiler/src/layering/wire.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,13 +147,7 @@ impl<'a, C: Config> CompileContext<'a, C> {
.map(|x| self.layout_query(x, ic.lcs[x.layer].vars.vec()))
.collect::<Vec<_>>();

for (i, (lc, lq)) in ic
.lcs
.iter()
.zip(lqs.iter())
.take(ic.output_layer)
.enumerate()
{
for (lc, lq) in ic.lcs.iter().zip(lqs.iter()).take(ic.output_layer) {
for x in lc.vars.vec() {
if !lq.var_pos.contains_key(x) {
panic!("unexpected situation");
Expand Down Expand Up @@ -207,9 +201,12 @@ impl<'a, C: Config> CompileContext<'a, C> {
}

let mut ress: Vec<Segment<C>> = Vec::new();
for (a, b) in layouts.iter().zip(layouts.iter().skip(1)) {
for (i, b) in layouts.iter().enumerate().skip(1) {
ress.push(Segment {
num_inputs: vec![a.size],
num_inputs: (ic.min_used_layer[i]..i)
.rev()
.map(|j| layouts[j].size)
.collect(),
num_outputs: b.size,
..Default::default()
});
Expand All @@ -228,10 +225,15 @@ impl<'a, C: Config> CompileContext<'a, C> {
.map(|x| sub_layouts_of_layer[x][insn_id].id)
.collect::<Vec<_>>();
let segment_ids = self.connect_wires(&cur_sub_layout_ids);
let sub_c = &self.circuits[&sub_id];

for (i, segment_id) in segment_ids.iter().enumerate() {
let alloc_min_layer = sub_c.min_used_layer[i + 1] + input_layer;
let al = Allocation {
input_offset: vec![sub_layouts_of_layer[input_layer + i][insn_id].offset],
input_offset: (alloc_min_layer..=input_layer + i)
.rev()
.map(|x| sub_layouts_of_layer[x][insn_id].offset)
.collect::<Vec<_>>(),
output_offset: sub_layouts_of_layer[input_layer + i + 1][insn_id].offset,
};
let mut found = false;
Expand All @@ -255,45 +257,28 @@ impl<'a, C: Config> CompileContext<'a, C> {
}

// connect self variables
for (cur_layer, ((lc, bq), aq)) in ic
.lcs
.iter()
.zip(lqs.iter())
.skip(1)
.zip(lqs.iter())
.enumerate()
{
let next_layer = cur_layer + 1;
let res = &mut ress[cur_layer];
for x in lc.vars.vec().iter() {
// only consider real variables
if *x >= ic.num_var {
continue;
}
let pos = if let Some(p) = bq.var_pos.get(x) {
for x in 0..ic.num_var {
// connect first occurance
if ic.min_layer[x] != 0 {
let next_layer = ic.min_layer[x];
let cur_layer = next_layer - 1;
let res = &mut ress[cur_layer];
let aq = &lqs[cur_layer];
let bq = &lqs[next_layer];
let pos = if let Some(p) = bq.var_pos.get(&x) {
*p
} else {
assert_eq!(cur_layer + 1, ic.output_layer);
//assert!(!ic.output_order.contains_key(x));
continue;
};
// if it's not the first layer, just relay it
if ic.min_layer[*x] != next_layer {
res.gate_adds.push(GateAdd {
inputs: [Input::new(0, aq.var_pos[x])],
output: pos,
coef: Coef::Constant(C::CircuitField::one()),
});
continue;
}
if let Some(value) = ic.constant_like_variables.get(x) {
if let Some(value) = ic.constant_like_variables.get(&x) {
res.gate_consts.push(GateConst {
inputs: [],
output: pos,
coef: value.clone(),
});
} else if ic.internal_variable_expr.contains_key(x) {
for term in ic.internal_variable_expr[x].iter() {
} else if ic.internal_variable_expr.contains_key(&x) {
for term in ic.internal_variable_expr[&x].iter() {
match &term.vars {
VarSpec::Const => {
res.gate_consts.push(GateConst {
Expand Down Expand Up @@ -341,6 +326,26 @@ impl<'a, C: Config> CompileContext<'a, C> {
}
}
}
// connect relays (this may generate cross layer connections)
for (cur_layer, next_layer) in ic.occured_layers[x]
.iter()
.zip(ic.occured_layers[x].iter().skip(1))
{
let res = &mut ress[next_layer - 1];
let aq = &lqs[*cur_layer];
let bq = &lqs[*next_layer];
let pos = if let Some(p) = bq.var_pos.get(&x) {
*p
} else {
assert_eq!(*next_layer, ic.output_layer);
continue;
};
res.gate_adds.push(GateAdd {
inputs: [Input::new(next_layer - cur_layer - 1, aq.var_pos[&x])],
output: pos,
coef: Coef::Constant(C::CircuitField::one()),
});
}
}

// also combined output variables
Expand Down

0 comments on commit a06d6f6

Please sign in to comment.