feat: optimize brillig ToRadix for common radices

acvm-repo/brillig_vm/Cargo.toml

@@ -22,6 +22,15 @@ num-bigint.workspace = true
 num-traits.workspace = true
 thiserror.workspace = true
 
+[dev-dependencies]
+criterion.workspace = true
+pprof.workspace = true
+rand.workspace = true
+
+[[bench]]
+name = "criterion"
+harness = false
+
 [features]
 bn254 = ["acir/bn254"]
 bls12_381 = ["acir/bls12_381"]

acvm-repo/brillig_vm/benches/README.md

@@ -0,0 +1,9 @@
+# Benchmarks 
+
+To generate flamegraphs for the execution of a specific benchmark, execute the following commands:
+
+```shell
+./scripts/benchmark_start.sh
+cargo bench -p brillig_vm --bench criterion -- --profile-time=30
+./scripts/benchmark_stop.sh
+```

acvm-repo/brillig_vm/benches/criterion.rs

@@ -0,0 +1,108 @@
+use acir::{
+    brillig::{self, MemoryAddress},
+    AcirField, FieldElement,
+};
+use brillig_vm::VMStatus;
+use criterion::{criterion_group, criterion_main, Criterion};
+
+use pprof::criterion::{Output, PProfProfiler};
+use rand::Rng;
+
+use std::time::Duration;
+
+fn byte_decomposition(c: &mut Criterion) {
+    let mut bytecode = vec![
+        // Radix
+        brillig::Opcode::Const {
+            destination: MemoryAddress::Direct(2),
+            bit_size: brillig::BitSize::Integer(brillig::IntegerBitSize::U32),
+            value: FieldElement::from(2u128),
+        },
+        // Output pointer
+        brillig::Opcode::Const {
+            destination: MemoryAddress::Direct(3),
+            bit_size: brillig::BitSize::Integer(brillig::IntegerBitSize::U32),
+            value: FieldElement::from(5u128),
+        },
+        // num_limbs
+        brillig::Opcode::Const {
+            destination: MemoryAddress::Direct(4),
+            bit_size: brillig::BitSize::Integer(brillig::IntegerBitSize::U32),
+            value: FieldElement::from(32u128),
+        },
+        // output_bits
+        brillig::Opcode::Const {
+            destination: MemoryAddress::Direct(5),
+            bit_size: brillig::BitSize::Integer(brillig::IntegerBitSize::U1),
+            value: FieldElement::from(false),
+        },
+        // calldata offset
+        brillig::Opcode::Const {
+            destination: MemoryAddress::Direct(0),
+            bit_size: brillig::BitSize::Integer(brillig::IntegerBitSize::U32),
+            value: FieldElement::from(0u128),
+        },
+        // calldata size
+        brillig::Opcode::Const {
+            destination: MemoryAddress::Direct(1),
+            bit_size: brillig::BitSize::Integer(brillig::IntegerBitSize::U32),
+            value: FieldElement::from(1u128),
+        },
+        brillig::Opcode::CalldataCopy {
+            destination_address: MemoryAddress::Direct(0),
+            size_address: MemoryAddress::Direct(1),
+            offset_address: MemoryAddress::Direct(0),
+        },
+    ];
+
+    for _ in 0..1 {
+        bytecode.push(brillig::Opcode::BlackBox(brillig::BlackBoxOp::ToRadix {
+            input: MemoryAddress::Direct(0),
+            radix: MemoryAddress::Direct(2),
+            output_pointer: MemoryAddress::Direct(3),
+            num_limbs: MemoryAddress::Direct(4),
+            output_bits: MemoryAddress::Direct(5),
+        }));
+    }
+
+    bench_bytecode(c, "byte_decomposition", &bytecode);
+}
+
+fn bench_bytecode<F: AcirField>(
+    c: &mut Criterion,
+    benchmark_name: &str,
+    bytecode: &[brillig::Opcode<F>],
+) {
+    c.bench_function(benchmark_name, |b| {
+        b.iter_batched(
+            || {
+                let input = rand::thread_rng().gen::<u128>();
+                let input_field = F::from(input);
+                let mut vm = brillig_vm::VM::new(
+                    vec![input_field],
+                    bytecode,
+                    &acvm_blackbox_solver::StubbedBlackBoxSolver(false),
+                    false,
+                );
+
+                // Process up to the last opcode which we want to benchmark
+                for _ in 0..bytecode.len() - 1 {
+                    vm.process_opcode();
+                }
+                vm
+            },
+            |mut vm| {
+                let status = vm.process_opcodes();
+                assert!(matches!(status, VMStatus::Finished { .. }))
+            },
+            criterion::BatchSize::SmallInput,
+        );
+    });
+}
+
+criterion_group! {
+    name = execution_benches;
+    config = Criterion::default().sample_size(20).measurement_time(Duration::from_secs(20)).with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
+    targets = byte_decomposition
+}
+criterion_main!(execution_benches);

acvm-repo/brillig_vm/src/black_box.rs

@@ -5,7 +5,6 @@
     sha256_compression, BigIntSolverWithId, BlackBoxFunctionSolver, BlackBoxResolutionError,
 };
 use num_bigint::BigUint;
-use num_traits::Zero;
 
 use crate::memory::MemoryValue;
 use crate::Memory;
@@ -59,10 +58,10 @@
                 to_u8_vec(read_heap_array(memory, key)).try_into().map_err(|_| {
                     BlackBoxResolutionError::Failed(bb_func, "Invalid key length".to_string())
                 })?;
             let ciphertext = aes128_encrypt(&inputs, iv, key)?;
 
             memory.write(outputs.size, ciphertext.len().into());
             memory.write_slice(memory.read_ref(outputs.pointer), &to_value_vec(&ciphertext));
 
             Ok(())
         }
@@ -320,39 +319,112 @@
                 panic!("ToRadix opcode's output_bits size does not match expected bit size 1")
             };
 
-            let mut input = BigUint::from_bytes_be(&input.to_be_bytes());
-            let radix = BigUint::from_bytes_be(&radix.to_be_bytes());
-
-            let mut limbs: Vec<MemoryValue<F>> = vec![MemoryValue::default(); num_limbs];
-
             assert!(
-                radix >= BigUint::from(2u32) && radix <= BigUint::from(256u32),
+                (2..=256).contains(&radix),
                 "Radix out of the valid range [2,256]. Value: {}",
                 radix
             );
 
             assert!(
-                num_limbs >= 1 || input == BigUint::from(0u32),
+                num_limbs >= 1 || input.is_zero(),
                 "Input value {} is not zero but number of limbs is zero.",
                 input
             );
 
             assert!(
-                !output_bits || radix == BigUint::from(2u32),
+                !output_bits || radix == 2,
                 "Radix {} is not equal to 2 and bit mode is activated.",
                 radix
             );
 
-            for i in (0..num_limbs).rev() {
-                let limb = &input % &radix;
-                if output_bits {
-                    limbs[i] = MemoryValue::U1(!limb.is_zero());
-                } else {
-                    let limb: u8 = limb.try_into().unwrap();
-                    limbs[i] = MemoryValue::U8(limb);
-                };
-                input /= &radix;
-            }
+            // TODO: raise an error if we cannot fit `input` into `num_limbs`.
+            let limbs = match radix {
+                256 => input
+                    .to_be_bytes()
+                    .into_iter()
+                    .rev()
+                    .take(num_limbs)
+                    .rev()
+                    .map(|limb| MemoryValue::U8(limb))
+                    .collect(),
+                2 => input
+                    .to_be_bytes()
+                    .into_iter()
+                    .rev()
+                    .take(num_limbs.div_ceil(8))
+                    .rev()
+                    .flat_map(|limb| {
+                        [
+                            MemoryValue::U1(((limb & 0x80) >> 7) == 1),
+                            MemoryValue::U1(((limb & 0x40) >> 6) == 1),
+                            MemoryValue::U1(((limb & 0x20) >> 5) == 1),
+                            MemoryValue::U1(((limb & 0x10) >> 4) == 1),
+                            MemoryValue::U1(((limb & 0x08) >> 3) == 1),
+                            MemoryValue::U1(((limb & 0x04) >> 2) == 1),
+                            MemoryValue::U1(((limb & 0x02) >> 1) == 1),
+                            MemoryValue::U1((limb & 0x01) == 1),
+                        ]
+                        .into_iter()
+                    })
+                    .rev()
+                    .take(num_limbs)
+                    .collect::<Vec<MemoryValue<F>>>()
+                    .into_iter()
+                    .rev()
+                    .collect(),
+                16 => input
+                    .to_be_bytes()
+                    .into_iter()
+                    .rev()
+                    .take(num_limbs.div_ceil(2))
+                    .rev()
+                    .flat_map(|limb| {
+                        [MemoryValue::U8((limb & 0xf0) >> 4), MemoryValue::U8(limb & 0x0f)]
+                            .into_iter()
+                    })
+                    .rev()
+                    .take(num_limbs)
+                    .collect::<Vec<MemoryValue<F>>>()
+                    .into_iter()
+                    .rev()
+                    .collect(),
+                4 => input
+                    .to_be_bytes()
+                    .into_iter()
+                    .rev()
+                    .take(num_limbs.div_ceil(4))
+                    .rev()
+                    .flat_map(|limb| {
+                        [
+                            MemoryValue::U8((limb & 0xc0) >> 6),
+                            MemoryValue::U8((limb & 0x30) >> 4),
+                            MemoryValue::U8((limb & 0x0c) >> 2),
+                            MemoryValue::U8(limb & 0x03),
+                        ]
+                        .into_iter()
+                    })
+                    .rev()
+                    .take(num_limbs)
+                    .collect::<Vec<MemoryValue<F>>>()
+                    .into_iter()
+                    .rev()
+                    .collect(),
+
+                _ => {
+                    let mut input = BigUint::from_bytes_be(&input.to_be_bytes());
+                    let radix = BigUint::from(radix);
+
+                    let mut limbs: Vec<MemoryValue<F>> = vec![MemoryValue::default(); num_limbs];
+                    for i in (0..num_limbs).rev() {
+                        let limb = &input % &radix;
+                        let limb: u8 = limb.try_into().unwrap();
+                        limbs[i] = MemoryValue::U8(limb);
+
+                        input /= &radix;
+                    }
+                    limbs
+                }
+            };
 
             memory.write_slice(memory.read_ref(*output_pointer), &limbs);