Skip to content

Commit

Permalink
feat: decompose Instruction::Constrain into multiple more basic con…
Browse files Browse the repository at this point in the history 
…straints (#3892)

# Description

## Problem\*

Some experimentation related to #3782.

## Summary\*

If we switch to an unary constrain instruction then we will want to be
able to decompose an assertion on a composite type to perform assertions
directly on its fields. Otherwise we have to create a predicate for each
field and then collect these all together before constraining that.

For example consider this program which simulates an unary constrain
```
fn main(x: [Field; 2], y: [Field; 2]) {
    let equal = x == y;
    assert(equal);
}
```

Currently this produces:

```
acir fn main f0 {
  b0(v0: [Field; 2], v1: [Field; 2]):
    v66 = array_get v0, index Field 0
    v67 = array_get v1, index Field 0
    v68 = eq v66, v67
    v69 = array_get v0, index Field 1
    v70 = array_get v1, index Field 1
    v71 = eq v69, v70
    v72 = mul v68, v71
    constrain v72 == u1 1
    return 
}

Compiled ACIR for main (unoptimized):
current witness index : 10
public parameters indices : []
return value indices : []
EXPR [ (-1, _1) (1, _3) (-1, _5) 0 ]
BRILLIG: inputs: [Single(Expression { mul_terms: [], linear_combinations: [(1, Witness(5))], q_c: 0 })]
outputs: [Simple(Witness(6))]
[JumpIfNot { condition: RegisterIndex(0), location: 3 }, Const { destination: RegisterIndex(1), value: Value { inner: 1 } }, BinaryFieldOp { destination: RegisterIndex(0), op: Div, lhs: RegisterIndex(1), rhs: RegisterIndex(0) }, Stop]

EXPR [ (1, _5, _6) (1, _7) -1 ]
EXPR [ (1, _5, _7) 0 ]
EXPR [ (-1, _2) (1, _4) (-1, _8) 0 ]
BRILLIG: inputs: [Single(Expression { mul_terms: [], linear_combinations: [(1, Witness(8))], q_c: 0 })]
outputs: [Simple(Witness(9))]
[JumpIfNot { condition: RegisterIndex(0), location: 3 }, Const { destination: RegisterIndex(1), value: Value { inner: 1 } }, BinaryFieldOp { destination: RegisterIndex(0), op: Div, lhs: RegisterIndex(1), rhs: RegisterIndex(0) }, Stop]

EXPR [ (1, _8, _9) (1, _10) -1 ]
EXPR [ (1, _8, _10) 0 ]
EXPR [ (1, _7, _10) -1 ]
```

After this PR we can recover the current optimisations for `assert(x ==
y)`

```
acir fn main f0 {
  b0(v0: [Field; 2], v1: [Field; 2]):
    v66 = array_get v0, index Field 0
    v67 = array_get v1, index Field 0
    v68 = array_get v0, index Field 1
    v69 = array_get v1, index Field 1
    constrain v66 == v67
    constrain v68 == v69
    return 
}

Compiled ACIR for main (unoptimized):
current witness index : 4
public parameters indices : []
return value indices : []
EXPR [ (1, _1) (-1, _3) 0 ]
EXPR [ (1, _2) (-1, _4) 0 ]
```

## Additional Context



## Documentation\*

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[Exceptional Case]** Documentation to be submitted in a separate
PR.

# PR Checklist\*

- [ ] I have tested the changes locally.
- [ ] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

---------

Co-authored-by: jfecher <[email protected]>
  • Loading branch information
@TomAFrench @jfecher
TomAFrench and jfecher authored Jan 10, 2024
1 parent 857ff97 commit 51cf9d3
Show file tree
Hide file tree
Showing 3 changed files with 222 additions and 73 deletions.
30 changes: 25 additions & 5 deletions compiler/noirc_evaluator/src/ssa/ir/dfg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,11 +166,31 @@ impl DataFlowGraph {
SimplifiedToMultiple(simplification)
}
SimplifyResult::Remove => InstructionRemoved,
result @ (SimplifyResult::SimplifiedToInstruction(_) | SimplifyResult::None) => {
let instruction = result.instruction().unwrap_or(instruction);
let id = self.make_instruction(instruction, ctrl_typevars);
self.blocks[block].insert_instruction(id);
self.locations.insert(id, call_stack);
result @ (SimplifyResult::SimplifiedToInstruction(_)
| SimplifyResult::SimplifiedToInstructionMultiple(_)
| SimplifyResult::None) => {
let instructions = result.instructions().unwrap_or(vec![instruction]);

if instructions.len() > 1 {
// There's currently no way to pass results from one instruction in `instructions` on to the next.
// We then restrict this to only support multiple instructions if they're all `Instruction::Constrain`
// as this instruction type does not have any results.
assert!(
instructions.iter().all(|instruction| matches!(instruction, Instruction::Constrain(..))),
"`SimplifyResult::SimplifiedToInstructionMultiple` only supports `Constrain` instructions"
);
}

let mut last_id = None;

for instruction in instructions {
let id = self.make_instruction(instruction, ctrl_typevars.clone());
self.blocks[block].insert_instruction(id);
self.locations.insert(id, call_stack.clone());
last_id = Some(id);
}

let id = last_id.expect("There should be at least 1 simplified instruction");
InsertInstructionResult::Results(id, self.instruction_results(id))
}
}
Expand Down
202 changes: 134 additions & 68 deletions compiler/noirc_evaluator/src/ssa/ir/instruction.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -432,73 +432,11 @@ impl Instruction {
}
}
Instruction::Constrain(lhs, rhs, msg) => {
if dfg.resolve(*lhs) == dfg.resolve(*rhs) {
// Remove trivial case `assert_eq(x, x)`
SimplifyResult::Remove
let constraints = decompose_constrain(*lhs, *rhs, msg.clone(), dfg);
if constraints.is_empty() {
Remove
} else {
match (&dfg[dfg.resolve(*lhs)], &dfg[dfg.resolve(*rhs)]) {
(
Value::NumericConstant { constant, typ },
Value::Instruction { instruction, .. },
)
| (
Value::Instruction { instruction, .. },
Value::NumericConstant { constant, typ },
) if *typ == Type::bool() => {
match dfg[*instruction] {
Instruction::Binary(Binary {
lhs,
rhs,
operator: BinaryOp::Eq,
}) if constant.is_one() => {
// Replace an explicit two step equality assertion
//
// v2 = eq v0, u32 v1
// constrain v2 == u1 1
//
// with a direct assertion of equality between the two values
//
// v2 = eq v0, u32 v1
// constrain v0 == v1
//
// Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
// will likely be removed through dead instruction elimination.

SimplifiedToInstruction(Instruction::Constrain(
lhs,
rhs,
msg.clone(),
))
}
Instruction::Not(value) => {
// Replace an assertion that a not instruction is truthy
//
// v1 = not v0
// constrain v1 == u1 1
//
// with an assertion that the not instruction input is falsy
//
// v1 = not v0
// constrain v0 == u1 0
//
// Note that this doesn't remove the value `v1` as it may be used in other instructions, but it
// will likely be removed through dead instruction elimination.
let reversed_constant = FieldElement::from(!constant.is_one());
let reversed_constant =
dfg.make_constant(reversed_constant, Type::bool());
SimplifiedToInstruction(Instruction::Constrain(
value,
reversed_constant,
msg.clone(),
))
}

_ => None,
}
}

_ => None,
}
SimplifiedToInstructionMultiple(constraints)
}
}
Instruction::ArrayGet { array, index } => {
Expand Down Expand Up @@ -646,6 +584,129 @@ fn simplify_cast(value: ValueId, dst_typ: &Type, dfg: &mut DataFlowGraph) -> Sim
}
}

/// Try to decompose this constrain instruction. This constraint will be broken down such that it instead constrains
/// all the values which are used to compute the values which were being constrained.
fn decompose_constrain(
lhs: ValueId,
rhs: ValueId,
msg: Option<String>,
dfg: &mut DataFlowGraph,
) -> Vec<Instruction> {
let lhs = dfg.resolve(lhs);
let rhs = dfg.resolve(rhs);

if lhs == rhs {
// Remove trivial case `assert_eq(x, x)`
Vec::new()
} else {
match (&dfg[lhs], &dfg[rhs]) {
(Value::NumericConstant { constant, typ }, Value::Instruction { instruction, .. })
| (Value::Instruction { instruction, .. }, Value::NumericConstant { constant, typ })
if *typ == Type::bool() =>
{
match dfg[*instruction] {
Instruction::Binary(Binary { lhs, rhs, operator: BinaryOp::Eq })
if constant.is_one() =>
{
// Replace an explicit two step equality assertion
//
// v2 = eq v0, u32 v1
// constrain v2 == u1 1
//
// with a direct assertion of equality between the two values
//
// v2 = eq v0, u32 v1
// constrain v0 == v1
//
// Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
// will likely be removed through dead instruction elimination.

vec![Instruction::Constrain(lhs, rhs, msg)]
}

Instruction::Binary(Binary { lhs, rhs, operator: BinaryOp::Mul })
if constant.is_one() && dfg.type_of_value(lhs) == Type::bool() =>
{
// Replace an equality assertion on a boolean multiplication
//
// v2 = mul v0, v1
// constrain v2 == u1 1
//
// with a direct assertion that each value is equal to 1
//
// v2 = mul v0, v1
// constrain v0 == 1
// constrain v1 == 1
//
// This is due to the fact that for `v2` to be 1 then both `v0` and `v1` are 1.
//
// Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
// will likely be removed through dead instruction elimination.
let one = FieldElement::one();
let one = dfg.make_constant(one, Type::bool());

[
decompose_constrain(lhs, one, msg.clone(), dfg),
decompose_constrain(rhs, one, msg, dfg),
]
.concat()
}

Instruction::Binary(Binary { lhs, rhs, operator: BinaryOp::Or })
if constant.is_zero() =>
{
// Replace an equality assertion on an OR
//
// v2 = or v0, v1
// constrain v2 == u1 0
//
// with a direct assertion that each value is equal to 0
//
// v2 = or v0, v1
// constrain v0 == 0
// constrain v1 == 0
//
// This is due to the fact that for `v2` to be 0 then both `v0` and `v1` are 0.
//
// Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
// will likely be removed through dead instruction elimination.
let zero = FieldElement::zero();
let zero = dfg.make_constant(zero, dfg.type_of_value(lhs));

[
decompose_constrain(lhs, zero, msg.clone(), dfg),
decompose_constrain(rhs, zero, msg, dfg),
]
.concat()
}

Instruction::Not(value) => {
// Replace an assertion that a not instruction is truthy
//
// v1 = not v0
// constrain v1 == u1 1
//
// with an assertion that the not instruction input is falsy
//
// v1 = not v0
// constrain v0 == u1 0
//
// Note that this doesn't remove the value `v1` as it may be used in other instructions, but it
// will likely be removed through dead instruction elimination.
let reversed_constant = FieldElement::from(!constant.is_one());
let reversed_constant = dfg.make_constant(reversed_constant, Type::bool());
decompose_constrain(value, reversed_constant, msg, dfg)
}

_ => vec![Instruction::Constrain(lhs, rhs, msg)],
}
}

_ => vec![Instruction::Constrain(lhs, rhs, msg)],
}
}
}

/// The possible return values for Instruction::return_types
pub(crate) enum InstructionResultType {
/// The result type of this instruction matches that of this operand
Expand Down Expand Up @@ -1134,6 +1195,10 @@ pub(crate) enum SimplifyResult {
/// Replace this function with an simpler but equivalent instruction.
SimplifiedToInstruction(Instruction),

/// Replace this function with a set of simpler but equivalent instructions.
/// This is currently only to be used for [`Instruction::Constrain`].
SimplifiedToInstructionMultiple(Vec<Instruction>),

/// Remove the instruction, it is unnecessary
Remove,

Expand All @@ -1142,9 +1207,10 @@ pub(crate) enum SimplifyResult {
}

impl SimplifyResult {
pub(crate) fn instruction(self) -> Option<Instruction> {
pub(crate) fn instructions(self) -> Option<Vec<Instruction>> {
match self {
SimplifyResult::SimplifiedToInstruction(instruction) => Some(instruction),
SimplifyResult::SimplifiedToInstruction(instruction) => Some(vec![instruction]),
SimplifyResult::SimplifiedToInstructionMultiple(instructions) => Some(instructions),
_ => None,
}
}
Expand Down
63 changes: 63 additions & 0 deletions compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -445,4 +445,67 @@ mod test {

assert_eq!(instruction, &Instruction::Cast(v0, Type::unsigned(32)));
}

#[test]
fn constraint_decomposition() {
// fn main f0 {
// b0(v0: u1, v1: u1, v2: u1):
// v3 = mul v0 v1
// v4 = not v2
// v5 = mul v3 v4
// constrain v4 u1 1
// }
//
// When constructing this IR, we should automatically decompose the constraint to be in terms of `v0`, `v1` and `v2`.
//
// The mul instructions are retained and will be removed in the dead instruction elimination pass.
let main_id = Id::test_new(0);

// Compiling main
let mut builder = FunctionBuilder::new("main".into(), main_id, RuntimeType::Acir);
let v0 = builder.add_parameter(Type::bool());
let v1 = builder.add_parameter(Type::bool());
let v2 = builder.add_parameter(Type::bool());

let v3 = builder.insert_binary(v0, BinaryOp::Mul, v1);
let v4 = builder.insert_not(v2);
let v5 = builder.insert_binary(v3, BinaryOp::Mul, v4);

// This constraint is automatically decomposed when it is inserted.
let v_true = builder.numeric_constant(true, Type::bool());
builder.insert_constrain(v5, v_true, None);

let v_false = builder.numeric_constant(false, Type::bool());

// Expected output:
//
// fn main f0 {
// b0(v0: u1, v1: u1, v2: u1):
// v3 = mul v0 v1
// v4 = not v2
// v5 = mul v3 v4
// constrain v0 u1 1
// constrain v1 u1 1
// constrain v2 u1 0
// }

let ssa = builder.finish();
let main = ssa.main();
let instructions = main.dfg[main.entry_block()].instructions();

assert_eq!(instructions.len(), 6);

assert_eq!(
main.dfg[instructions[0]],
Instruction::Binary(Binary { lhs: v0, operator: BinaryOp::Mul, rhs: v1 })
);
assert_eq!(main.dfg[instructions[1]], Instruction::Not(v2));
assert_eq!(
main.dfg[instructions[2]],
Instruction::Binary(Binary { lhs: v3, operator: BinaryOp::Mul, rhs: v4 })
);
assert_eq!(main.dfg[instructions[3]], Instruction::Constrain(v0, v_true, None));
assert_eq!(main.dfg[instructions[4]], Instruction::Constrain(v1, v_true, None));
assert_eq!(main.dfg[instructions[5]], Instruction::Constrain(v2, v_false, None));
}
}

0 comments on commit 51cf9d3

Please sign in to comment.