Chips in Halo2

You would not implement every method on the same class, would you?

So far, we have created multiple regions. We have created regions for addition and multiplication. We have been stringing them together having a grand old time. But things are getting a bit unwieldy: every method has been implemented on the TestCircuit struct.

Like writing Java with only a single class...

Throwing everything into the same struct also hinders reusability: nobody else can use the nice gadgets that we have created, without buying the whole TestCircuit struct. We need to find a way to make our gadgets more modular. The pattern we introduce here is that of "chips".

There is nothing special about a chip, it is just a way to structure our circuits. Like a class in object-oriented programming, a chip is a way to group related functionality together. Like a library.

I think we are also ready to graduate to multiple advice columns :)

An Arithmetic Chip

So let's dive into it.

We will create an ArithmeticChip that will implement a small library of arithmetic operations:

#[derive(Clone, Debug)]
struct ArithmeticChip<F: Field> {
    _ph: PhantomData<F>,
    q_mul: Selector,
    q_add: Selector,
    w0: Column<Advice>,
    w1: Column<Advice>,
    w2: Column<Advice>,
}

As you can see, the chip has:

  • A selector q_mul to enable the multiplication gate.
  • A selector q_add to enable the addition gate.
  • An advice column w0 to store the first "input".
  • An advice column w1 to store the second "input".
  • An advice column w2 to store the "output".

An instance of the ArithmeticChip can be created at configuration time for the circuit:

impl<F: Field> Circuit<F> for TestCircuit<F> {
    type Config = TestConfig<F>;
    type FloorPlanner = SimpleFloorPlanner;

    fn without_witnesses(&self) -> Self {
        TestCircuit {
            _ph: PhantomData,
            secret: Value::unknown(),
        }
    }

    fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
        // define advice columns
        let w0 = meta.advice_column();
        let w1 = meta.advice_column();
        let w2 = meta.advice_column();

        // enable equality constraints
        meta.enable_equality(w0);
        meta.enable_equality(w1);
        meta.enable_equality(w2);

        let arithmetic_chip = ArithmeticChip::configure(meta, w0, w1, w2);

        TestConfig {
            _ph: PhantomData,
            arithmetic_chip,
            w0,
            w1,
            w2,
        }
    }

Which in turn will invoke a method called configure on the ArithmeticChip. We will get to that in a bit. A natural question at this point arises:

Why do we define the columns in the configure of the TestCircuit?

Namely, why not define the Advice columns inside ArithmeticChip::configure?

The answer is that we could do so, but that it is often useful for different chips to share the same columns. This is because introducing new columns is relatively expensive: it increases the size of the proof and the verification time. Of course Selector columns are usually chip-specific: sharing them would cause different chips to interfere with each other.

Configuration of the ArithmeticChip should be relatively straightforward to the reader by now:

impl<F: Field> ArithmeticChip<F> {
    fn configure(
        meta: &mut ConstraintSystem<F>,
        w0: Column<Advice>,
        w1: Column<Advice>,
        w2: Column<Advice>,
    ) -> Self {
        let q_mul = meta.complex_selector();
        let q_add = meta.complex_selector();

        // define an addition gate:
        meta.create_gate("add", |meta| {
            let w0 = meta.query_advice(w0, Rotation::cur());
            let w1 = meta.query_advice(w1, Rotation::cur());
            let w2 = meta.query_advice(w2, Rotation::cur());
            let q_add = meta.query_selector(q_add);
            vec![q_add * (w0 + w1 - w2)]
        });

        // define a multiplication gate:
        meta.create_gate("mul", |meta| {
            let w0 = meta.query_advice(w0, Rotation::cur());
            let w1 = meta.query_advice(w1, Rotation::cur());
            let w2 = meta.query_advice(w2, Rotation::cur());
            let q_mul = meta.query_selector(q_mul);
            vec![q_mul * (w0 * w1 - w2)]
        });

        Self {
            _ph: PhantomData,
            q_mul,
            q_add,
            w0,
            w1,
            w2,
        }
    }

The only difference from previous examples is that the inputs/outputs are now stored next to each other in the Advice columns, rather than stacked on top of each other in the same column.

At this point we are ready to add some methods to our ArithmeticChip to create regions for addition and multiplication:

    fn add(
        &self,
        layouter: &mut impl Layouter<F>,
        lhs: AssignedCell<F, F>,
        rhs: AssignedCell<F, F>,
    ) -> Result<AssignedCell<F, F>, Error> {
        layouter.assign_region(
            || "add",
            |mut region| {
                // enable the addition gate
                self.q_add.enable(&mut region, 0)?;

                // compute cell values
                let w0 = lhs.value().cloned();
                let w1 = rhs.value().cloned();
                let w2 = w0.and_then(|w0| w1.and_then(|w1| Value::known(w0 + w1)));

                // assign the values to the cells
                let w0 = region.assign_advice(|| "assign w0", self.w0, 0, || w0)?;
                let w1 = region.assign_advice(|| "assign w1", self.w1, 0, || w1)?;
                let w2 = region.assign_advice(|| "assign w2", self.w2, 0, || w2)?;

                // constrain the inputs
                region.constrain_equal(w0.cell(), lhs.cell())?;
                region.constrain_equal(w1.cell(), rhs.cell())?;

                Ok(w2)
            },
        )
    }
    fn mul(
        &self,
        layouter: &mut impl Layouter<F>,
        lhs: AssignedCell<F, F>,
        rhs: AssignedCell<F, F>,
    ) -> Result<AssignedCell<F, F>, Error> {
        layouter.assign_region(
            || "mul",
            |mut region| {
                // enable the multiplication gate
                self.q_mul.enable(&mut region, 0)?;

                // compute cell values
                let w0 = lhs.value().cloned();
                let w1 = rhs.value().cloned();
                let w2 = w0.and_then(|w0| w1.and_then(|w1| Value::known(w0 * w1)));

                // assign the values to the cells
                let w0 = region.assign_advice(|| "assign w0", self.w0, 0, || w0)?;
                let w1 = region.assign_advice(|| "assign w1", self.w1, 0, || w1)?;
                let w2 = region.assign_advice(|| "assign w2", self.w2, 0, || w2)?;

                // constrain the inputs
                region.constrain_equal(w0.cell(), lhs.cell())?;
                region.constrain_equal(w1.cell(), rhs.cell())?;

                Ok(w2)
            },
        )
    }

This is essentially just refactoring the code we saw in the previous sections.

Finally, we can use the ArithmeticChip in the TestCircuit during synthesize:

    fn synthesize(
        &self,
        config: Self::Config, //
        mut layouter: impl Layouter<F>,
    ) -> Result<(), Error> {
        let a1 = config
            .arithmetic_chip
            .free(&mut layouter, self.secret.clone())?;

        let a2 = config
            .arithmetic_chip
            .mul(&mut layouter, a1.clone(), a1.clone())?;

        let a3 = config
            .arithmetic_chip
            .add(&mut layouter, a1.clone(), a2.clone())?;

        Ok(())
    }

That's pretty much chips in a nutshell: they are simply a collection of functionality, with convenient methods to create regions. In the exercises we explore how to use a more complex chip along with custom types.

Exercises

So the arithmetic chip we have so far is nice and all, but pretty limited/inefficient. A common optimization in PlonKish (e.g. Halo2) is to have a single arithmetic super gate, that can do:

  • Addition
  • Multiplication
  • Bit range checks
  • Constant equality checks
  • And more!

While also allowing free multiplication/addition by constants.

This allows PlonK to recoup some of the "linear combination is free" advantages exhibited by alternative arithmetizations like R1CS (like Marlin or Groth16). Such a general arithmetic gate looks as follows:

\[ \text{c0} \cdot \text{w0} + \text{c1} \cdot \text{w1} + \text{c2} \cdot \text{w2} + \text{cm} \cdot (\text{w0} \cdot \text{w1}) + \text{cc} = 0 \]

Where:

  • cm is the multiplication constant (Column<Fixed>)
  • c0, c1, c2 are linear constants (Column<Fixed>)
  • cc is an additive constant (Column<Fixed>)
  • w0, w1, w2 are the advice cells (Column<Advice>)

Exercise

Update the ArithmeticChip to use the new gate for both addition/multiplication.

This requires introducing a number of new fixed columns.

Where do you think these should be defined?

Hint

Setting:

  • c0 = 1
  • c1 = 1
  • c2 = -1
  • cm = 0
  • cc = 0

Gives you an addition gate.

Exercise

Use the same gate to implement a function fixing a cell to a constant:

Create a function which takes a field element and returns an assigned cell which must be equal to the field element.

Exercise

Use the same gate to implement a function forcing equality between two cells.

Hint

Set:

  • c0 = 1
  • c1 = -1
  • c2 = 0
  • cm = 0
  • cc = 0

One of the main motivations for the new gate was to reduce the cost of multiplication/addition by constants. To achieve this we introduce a new struct, called Variable:

#[derive(Clone, Debug)]
struct Variable<F: Field> {
    mul: F,
    add: F,
    val: AssignedCell<F, F>,
}

impl<F: Field> Variable<F> {
    fn value(&self) -> Value<F> {
        self.val.value().map(|v| self.mul * v + self.add)
    }
}

The value of a Variable is mul * val + add where mul and add are fixed field elements and val is a cell of Column<Advice>. This is useful because we can add/multiply by constants without introducing new regions, simply by changing the mul and add constants:

impl<F: Field> Add<F> for Variable<F> {
    type Output = Self;

    fn add(self, rhs: F) -> Self {
        Self {
            mul: self.mul,
            add: self.add + rhs,
            val: self.val,
        }
    }
}

impl<F: Field> Mul<F> for Variable<F> {
    type Output = Self;

    fn mul(self, rhs: F) -> Self {
        Self {
            mul: self.mul * rhs,
            add: self.add * rhs,
            val: self.val,
        }
    }
}

Exercise

Update the ArithmeticChip to use the new Variable struct for addition/multiplication: accounting for the new multiplicative and additive constants.

Hint

Addition is easy, you just need to change c0, c1, cc appropriately.

Hint

Multiplication is a bit tricky, you need to expand the expression:

\[ (\text{mul}_1 \cdot \text{val}_1 + \text{add}_1) \cdot (\text{mul}_2 \cdot \text{val}_2 + \text{add}_2) \]

Then set c0, c1, c2, cm, cc appropriately to account for the cross terms. Good luck!

This gate can also be used to do other things, e.g. create a fresh variable restricted to a single bit (the prover must assign the variable either 0 or 1):

    /// Allocate a bit-constrained variable.
    fn bit(
        &self,
        layouter: &mut impl Layouter<F>,
        value: Value<bool>,
    ) -> Result<Variable<F>, Error> {

Exercise

Implement the bit function above.

Exercise

Implement a function enforcing equality between two Variable's.

Solution

use std::{
    marker::PhantomData,
    ops::{Add, Mul},
};

use halo2_proofs::{
    circuit::{AssignedCell, Layouter, SimpleFloorPlanner, Value},
    dev::MockProver,
    plonk::{
        Advice,
        Circuit,
        Column,
        ConstraintSystem, //
        Error,
        Expression,
        Fixed,
        Selector,
    },
    poly::Rotation,
};

use ff::{Field, PrimeField};

struct TestCircuit<F: Field> {
    _ph: PhantomData<F>,
    secret: Value<F>,
}

// ANCHOR: variable
#[derive(Clone, Debug)]
struct Variable<F: Field> {
    mul: F,
    add: F,
    val: AssignedCell<F, F>,
}

impl<F: Field> Variable<F> {
    fn value(&self) -> Value<F> {
        self.val.value().map(|v| self.mul * v + self.add)
    }
}
// ANCHOR_END: variable

// ANCHOR: add-mul-const
impl<F: Field> Add<F> for Variable<F> {
    type Output = Self;

    fn add(self, rhs: F) -> Self {
        Self {
            mul: self.mul,
            add: self.add + rhs,
            val: self.val,
        }
    }
}

impl<F: Field> Mul<F> for Variable<F> {
    type Output = Self;

    fn mul(self, rhs: F) -> Self {
        Self {
            mul: self.mul * rhs,
            add: self.add * rhs,
            val: self.val,
        }
    }
}
// ANCHOR_END: add-mul-const

#[derive(Clone, Debug)]
struct ArithmeticChip<F: Field> {
    _ph: PhantomData<F>,
    q_arith: Selector,
    cm: Column<Fixed>,
    c0: Column<Fixed>,
    c1: Column<Fixed>,
    c2: Column<Fixed>,
    cc: Column<Fixed>,
    w0: Column<Advice>,
    w1: Column<Advice>,
    w2: Column<Advice>,
}

impl<F: Field> ArithmeticChip<F> {
    fn configure(
        meta: &mut ConstraintSystem<F>,
        w0: Column<Advice>,
        w1: Column<Advice>,
        w2: Column<Advice>,
        c0: Column<Fixed>,
        c1: Column<Fixed>,
        c2: Column<Fixed>,
        cm: Column<Fixed>,
        cc: Column<Fixed>,
    ) -> Self {
        let q_arith = meta.complex_selector();

        // define arithmetic gate
        meta.create_gate("arith", |meta| {
            let w0 = meta.query_advice(w0, Rotation::cur());
            let w1 = meta.query_advice(w1, Rotation::cur());
            let w2 = meta.query_advice(w2, Rotation::cur());

            let c0 = meta.query_fixed(c0, Rotation::cur());
            let c1 = meta.query_fixed(c1, Rotation::cur());
            let c2 = meta.query_fixed(c2, Rotation::cur());

            let cm = meta.query_fixed(cm, Rotation::cur());
            let cc = meta.query_fixed(cc, Rotation::cur());

            let q_arith = meta.query_selector(q_arith);

            // define the arithmetic expression
            //
            // w0 * c0 + w1 * c1 + w2 * c2 + cm * (w0 * w1) + cc
            let expr = Expression::Constant(F::ZERO);
            let expr = expr + c0 * w0.clone();
            let expr = expr + c1 * w1.clone();
            let expr = expr + c2 * w2.clone();
            let expr = expr + cm * (w0 * w1);
            let expr = expr + cc;
            vec![q_arith * expr]
        });

        Self {
            _ph: PhantomData,
            q_arith,
            cm,
            c0,
            c1,
            c2,
            cc,
            w0,
            w1,
            w2,
        }
    }

    /// Multiply two variables
    fn mul(
        &self,
        layouter: &mut impl Layouter<F>,
        lhs: Variable<F>,
        rhs: Variable<F>,
    ) -> Result<Variable<F>, Error> {
        layouter.assign_region(
            || "mul",
            |mut region| {
                // turn on the arithmetic gate
                self.q_arith.enable(&mut region, 0)?;

                // (c0 * w0 + cc1) * (c1 * w1 + cc2)
                // c0 * c1 * (w0 * w1) + c0 * cc2 * w0 + c1 * cc1 * w1 + cc1 * cc2
                lhs.val.copy_advice(|| "lhs", &mut region, self.w0, 0)?;
                rhs.val.copy_advice(|| "rhs", &mut region, self.w1, 0)?;

                let val =
                    region.assign_advice(|| "res", self.w2, 0, || lhs.value() * rhs.value())?;

                region.assign_fixed(|| "c0", self.c0, 0, || Value::known(lhs.mul * rhs.add))?;
                region.assign_fixed(|| "c1", self.c1, 0, || Value::known(rhs.mul * lhs.add))?;
                region.assign_fixed(|| "c2", self.c2, 0, || Value::known(-F::ONE))?;
                region.assign_fixed(|| "cc", self.cc, 0, || Value::known(lhs.add * rhs.add))?;
                region.assign_fixed(|| "cm", self.cm, 0, || Value::known(lhs.mul * rhs.mul))?;

                Ok(Variable {
                    mul: F::ONE,
                    add: F::ZERO,
                    val,
                })
            },
        )
    }

    /// Add two variables
    fn add(
        &self,
        layouter: &mut impl Layouter<F>,
        lhs: Variable<F>,
        rhs: Variable<F>,
    ) -> Result<Variable<F>, Error> {
        layouter.assign_region(
            || "add",
            |mut region| {
                // turn on the arithmetic gate
                self.q_arith.enable(&mut region, 0)?;

                lhs.val.copy_advice(|| "lhs", &mut region, self.w0, 0)?;
                rhs.val.copy_advice(|| "rhs", &mut region, self.w1, 0)?;

                let val =
                    region.assign_advice(|| "res", self.w2, 0, || lhs.value() + rhs.value())?;

                region.assign_fixed(|| "c0", self.c0, 0, || Value::known(lhs.mul))?;
                region.assign_fixed(|| "c1", self.c1, 0, || Value::known(rhs.mul))?;
                region.assign_fixed(|| "c2", self.c2, 0, || Value::known(-F::ONE))?;
                region.assign_fixed(|| "cc", self.cc, 0, || Value::known(lhs.add + rhs.add))?;
                region.assign_fixed(|| "cm", self.cm, 0, || Value::known(F::ZERO))?;

                Ok(Variable {
                    mul: F::ONE,
                    add: F::ZERO,
                    val,
                })
            },
        )
    }

    /// Allocate a free variable.
    fn free(&self, layouter: &mut impl Layouter<F>, value: Value<F>) -> Result<Variable<F>, Error> {
        layouter.assign_region(
            || "free",
            |mut region| {
                // no need to turn on anything
                let val = region.assign_advice(|| "free", self.w0, 0, || value)?;
                region.assign_advice(|| "junk1", self.w1, 0, || Value::known(F::ZERO))?;
                region.assign_advice(|| "junk2", self.w2, 0, || Value::known(F::ZERO))?;
                Ok(Variable {
                    mul: F::ONE,
                    add: F::ZERO,
                    val,
                })
            },
        )
    }

    /// Assert equal
    fn eq_consant(
        &self,
        layouter: &mut impl Layouter<F>,
        constant: F,
        variable: Variable<F>,
    ) -> Result<(), Error> {
        layouter.assign_region(
            || "eq_constant",
            |mut region| {
                // turn on the arithmetic gate
                self.q_arith.enable(&mut region, 0)?;

                variable
                    .val
                    .copy_advice(|| "val", &mut region, self.w0, 0)?;

                let delta = variable.add - constant;

                region.assign_advice(|| "junk1", self.w1, 0, || Value::known(F::ZERO))?;
                region.assign_advice(|| "junk2", self.w2, 0, || Value::known(F::ZERO))?;

                region.assign_fixed(|| "c0", self.c0, 0, || Value::known(variable.mul))?;
                region.assign_fixed(|| "c1", self.c1, 0, || Value::known(F::ZERO))?;
                region.assign_fixed(|| "c2", self.c2, 0, || Value::known(F::ZERO))?;
                region.assign_fixed(|| "cc", self.cc, 0, || Value::known(delta))?;
                region.assign_fixed(|| "cm", self.cm, 0, || Value::known(F::ZERO))?;

                Ok(())
            },
        )
    }

    // ANCHOR: bit
    /// Allocate a bit-constrained variable.
    fn bit(
        &self,
        layouter: &mut impl Layouter<F>,
        value: Value<bool>,
    ) -> Result<Variable<F>, Error> {
        // ANCHOR_END: bit
        layouter.assign_region(
            || "bit",
            |mut region| {
                // turn on the arithmetic gate
                self.q_arith.enable(&mut region, 0)?;

                // (v1 - 1) * v1 = v1^2 - v1
                let w0 = region.assign_advice(
                    || "bit0",
                    self.w0,
                    0,
                    || value.map(|b| if b { F::ONE } else { F::ZERO }),
                )?;

                let w1 = region.assign_advice(
                    || "bit1",
                    self.w1,
                    0,
                    || value.map(|b| if b { F::ONE } else { F::ZERO }),
                )?;

                region.assign_advice(|| "junk", self.w2, 0, || Value::known(F::ZERO))?;

                region.constrain_equal(w0.cell(), w1.cell())?;

                region.assign_fixed(|| "c0", self.c0, 0, || Value::known(F::ZERO))?;
                region.assign_fixed(|| "c1", self.c0, 0, || Value::known(-F::ONE))?;
                region.assign_fixed(|| "c2", self.c0, 0, || Value::known(F::ZERO))?;
                region.assign_fixed(|| "cc", self.cc, 0, || Value::known(F::ZERO))?;
                region.assign_fixed(|| "cm", self.cm, 0, || Value::known(F::ONE))?;

                Ok(Variable {
                    mul: F::ONE,
                    add: F::ZERO,
                    val: w0,
                })
            },
        )
    }
}

#[derive(Clone, Debug)]
struct TestConfig<F: Field + Clone> {
    _ph: PhantomData<F>,
    arithmetic_chip: ArithmeticChip<F>,
}

impl<F: PrimeField> Circuit<F> for TestCircuit<F> {
    type Config = TestConfig<F>;
    type FloorPlanner = SimpleFloorPlanner;

    fn without_witnesses(&self) -> Self {
        TestCircuit {
            _ph: PhantomData,
            secret: Value::unknown(),
        }
    }

    fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
        // let q_enable = meta.fixed_column();
        let w0 = meta.advice_column();
        let w1 = meta.advice_column();
        let w2 = meta.advice_column();

        let c0 = meta.fixed_column();
        let c1 = meta.fixed_column();
        let c2 = meta.fixed_column();
        let cc = meta.fixed_column();
        let cm = meta.fixed_column();

        // enable equality constraints
        meta.enable_equality(w0);
        meta.enable_equality(w1);
        meta.enable_equality(w2);

        let arithmetic_chip = ArithmeticChip::configure(meta, w0, w1, w2, c0, c1, c2, cc, cm);

        TestConfig {
            _ph: PhantomData,
            arithmetic_chip,
        }
    }

    fn synthesize(
        &self,
        config: Self::Config, //
        mut layouter: impl Layouter<F>,
    ) -> Result<(), Error> {
        let a1 = config
            .arithmetic_chip
            .free(&mut layouter, self.secret.clone())?;

        let a2 = config
            .arithmetic_chip
            .add(&mut layouter, a1.clone(), a1.clone())?;

        let a3 = config
            .arithmetic_chip
            .mul(&mut layouter, a1.clone(), a2.clone())?;

        config
            .arithmetic_chip
            .eq_consant(&mut layouter, F::from_u128(1337 * (1337 + 1337)), a3)?;

        Ok(())
    }
}

fn main() {
    use halo2_proofs::halo2curves::bn256::Fr;

    // run the MockProver
    let circuit = TestCircuit::<Fr> {
        _ph: PhantomData,
        secret: Value::known(Fr::from(1337u64)),
    };
    let prover = MockProver::run(8, &circuit, vec![]).unwrap();
    prover.verify().unwrap();
}