From 479715d863805b0819a00d288ce71a0fef35de1f Mon Sep 17 00:00:00 2001 From: Ian Neal Date: Wed, 16 Aug 2023 12:02:12 -0400 Subject: [PATCH] Generalize mapped offset computation, add tests (#42) --- circom/tests/subcmps/mapped3.circom | 50 +++++++++++++++++ circom/tests/subcmps/mapped4.circom | 53 +++++++++++++++++++ circuit_passes/src/bucket_interpreter/mod.rs | 21 ++++---- .../src/bucket_interpreter/operations.rs | 52 ++++++++++++++++++ .../src/passes/mapped_to_indexed.rs | 18 +++---- 5 files changed, 172 insertions(+), 22 deletions(-) create mode 100644 circom/tests/subcmps/mapped3.circom create mode 100644 circom/tests/subcmps/mapped4.circom diff --git a/circom/tests/subcmps/mapped3.circom b/circom/tests/subcmps/mapped3.circom new file mode 100644 index 000000000..61b330457 --- /dev/null +++ b/circom/tests/subcmps/mapped3.circom @@ -0,0 +1,50 @@ +pragma circom 2.0.0; +// REQUIRES: circom +// RUN: rm -rf %t && mkdir %t && %circom --llvm -o %t %s | sed -n 's/.*Written successfully:.* \(.*\)/\1/p' | xargs cat | FileCheck %s + +template ArrayOp(q) { + signal input inp[15]; + signal output outp[15]; + + for (var i = 0; i < 15; i++) { + outp[i] <== inp[i] + q; + } +} + +//CHECK-LABEL: define void @ArrayOp_{{[0-9]+}}_build +//CHECK-SAME: ({ [0 x i256]*, i32 }* %{{.*}}) +//CHECK: alloca [30 x i256] +//CHECK: %[[DIM_REG:.*]] = getelementptr { [0 x i256]*, i32 }, { [0 x i256]*, i32 }* %0, i32 0, i32 1 +//CHECK: store i32 15, i32* %{{.*}}[[DIM_REG]] + +template Wrapper() { + signal input inp[15]; + signal output outp; + + component m[4]; + + for (var q = 0; q < 4; q++) { + // This test exhibits the behavior because the array of different subcomponents + // (differentiated by the template parameter changing) + m[q] = ArrayOp(q); + for (var i = 0; i < 15; i++) { + m[q].inp[i] <== inp[i]; + } + } + + outp <== m[2].outp[3]; +} + +component main = Wrapper(); + +//CHECK-LABEL: define void @Wrapper_{{[0-9]+}}_run +//CHECK-SAME: ([0 x i256]* %{{.*}}) +//CHECK: %lvars = alloca [2 x i256] +//COM: offset = (1 * (3 * 7)) + (2 * (7)) + (3) + 1 (since 0 is output) = 21 + 14 + 3 + 1 = 39 +//CHECK: store{{.*}}:{{.*}}; preds = %unrolled_loop{{.*}} +//CHECK: %[[SUB_PTR:.*]] = getelementptr [4 x { [0 x i256]*, i32 }], [4 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0 +//CHECK: %[[SUBCMP:.*]] = load [0 x i256]*, [0 x i256]** %[[SUB_PTR]] +//CHECK: %[[VAL_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %[[SUBCMP]], i32 0, i32 3 +//CHECK: %[[VAL:.*]] = load i256, i256* %[[VAL_PTR]] +//CHECK: %[[OUTP_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i32 0 +//CHECK: store i256 %[[VAL]], i256* %[[OUTP_PTR]] diff --git a/circom/tests/subcmps/mapped4.circom b/circom/tests/subcmps/mapped4.circom new file mode 100644 index 000000000..279ca8834 --- /dev/null +++ b/circom/tests/subcmps/mapped4.circom @@ -0,0 +1,53 @@ +pragma circom 2.0.0; +// REQUIRES: circom +// RUN: rm -rf %t && mkdir %t && %circom --llvm -o %t %s | sed -n 's/.*Written successfully:.* \(.*\)/\1/p' | xargs cat | FileCheck %s + +template MatrixOp(q) { + signal input inp[5][3]; + signal output outp[5][3]; + + for (var i = 0; i < 5; i++) { + for (var j = 0; j < 3; j++) { + outp[i][j] <== inp[i][j] + q; + } + } +} + +//CHECK-LABEL: define void @MatrixOp_{{[0-9]+}}_build +//CHECK-SAME: ({ [0 x i256]*, i32 }* %{{.*}}) +//CHECK: alloca [16 x i256] +//CHECK: %[[DIM_REG:.*]] = getelementptr { [0 x i256]*, i32 }, { [0 x i256]*, i32 }* %0, i32 0, i32 1 +//CHECK: store i32 15, i32* %{{.*}}[[DIM_REG]] + +template Wrapper() { + signal input inp[5][3]; + signal output outp; + + component m[4]; + + for (var q = 0; q < 4; q++) { + // This test exhibits the behavior because the array of different subcomponents + // (differentiated by the template parameter changing) + m[q] = MatrixOp(q); + for (var i = 0; i < 5; i++) { + for (var j = 0; j < 3; j++) { + m[q].inp[i][j] <== inp[i][j]; + } + } + } + + outp <== m[2].outp[1][2]; +} + +component main = Wrapper(); + +//CHECK-LABEL: define void @Wrapper_{{[0-9]+}}_run +//CHECK-SAME: ([0 x i256]* %{{.*}}) +//CHECK: %lvars = alloca [3 x i256] +//CHECK: store{{.*}}:{{.*}}; preds = %unrolled_loop{{.*}} +//CHECK: %[[SUB_PTR:.*]] = getelementptr [4 x { [0 x i256]*, i32 }], [4 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0 +//CHECK: %[[SUBCMP:.*]] = load [0 x i256]*, [0 x i256]** %[[SUB_PTR]] +//CHECK: %[[VAL_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %[[SUBCMP]], i32 0, i32 5 +//CHECK: %[[VAL:.*]] = load i256, i256* %[[VAL_PTR]] +//CHECK: %[[OUTP_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i32 0 +//CHECK: store i256 %[[VAL]], i256* %[[OUTP_PTR]] diff --git a/circuit_passes/src/bucket_interpreter/mod.rs b/circuit_passes/src/bucket_interpreter/mod.rs index a4df040e0..cb1c7d800 100644 --- a/circuit_passes/src/bucket_interpreter/mod.rs +++ b/circuit_passes/src/bucket_interpreter/mod.rs @@ -12,6 +12,7 @@ use compiler::num_bigint::BigInt; use observer::InterpreterObserver; use program_structure::constants::UsefulConstants; use crate::bucket_interpreter::env::Env; +use crate::bucket_interpreter::operations::compute_offset; use crate::bucket_interpreter::value::{JoinSemiLattice, Value}; use crate::bucket_interpreter::value::Value::{KnownBigInt, KnownU32, Unknown}; @@ -244,7 +245,8 @@ impl<'a> BucketInterpreter<'a> { }, LocationRule::Mapped { signal_code, indexes } => { let mut acc_env = env; - let map_access = self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code].offset; + let io_def = &self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code]; + let map_access = io_def.offset; if indexes.len() > 0 { let mut indexes_values = vec![]; for i in indexes { @@ -252,11 +254,8 @@ impl<'a> BucketInterpreter<'a> { indexes_values.push(val.expect("Mapped location must produce a value!").get_u32()); acc_env = new_env; } - if indexes.len() == 1 { - (map_access + indexes_values[0], acc_env) - } else { - todo!() - } + let offset = compute_offset(&indexes_values, &io_def.lengths); + (map_access + offset, acc_env) } else { (map_access, acc_env) } @@ -318,7 +317,8 @@ impl<'a> BucketInterpreter<'a> { LocationRule::Mapped { signal_code, indexes } => { let mut acc_env = env; let name = Some(acc_env.get_subcmp_name(addr).clone()); - let map_access = self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code].offset; + let io_def = &self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code]; + let map_access = io_def.offset; if indexes.len() > 0 { let mut indexes_values = vec![]; for i in indexes { @@ -326,11 +326,8 @@ impl<'a> BucketInterpreter<'a> { indexes_values.push(val.expect("Mapped location must produce a value!").get_u32()); acc_env = new_env; } - if indexes.len() == 1 { - (map_access + indexes_values[0], acc_env, name) - } else { - todo!() - } + let offset = compute_offset(&indexes_values, &io_def.lengths); + (map_access + offset, acc_env, name) } else { (map_access, acc_env, name) } diff --git a/circuit_passes/src/bucket_interpreter/operations.rs b/circuit_passes/src/bucket_interpreter/operations.rs index e4e21ca67..9f5848823 100644 --- a/circuit_passes/src/bucket_interpreter/operations.rs +++ b/circuit_passes/src/bucket_interpreter/operations.rs @@ -40,3 +40,55 @@ pub fn compute_operation(bucket: &ComputeBucket, stack: &Vec, p: &BigInt) }); computed_value } + +pub fn compute_offset(indexes: &Vec, lengths: &Vec) -> usize { + // Lengths are in order, i.e. arr[x][y] => [x, y], same with indices + // arr[x][y] is x arrays of length y, laid out sequentially + if indexes.len() != lengths.len() { + // I did check, both "indexes" and "indices" are valid plurals + panic!("must have the same number of indexes and array lengths!"); + } + let mut total_offset = indexes.last().copied().expect("must contain some indexes!"); + let mut size_multiplier = lengths.last().copied().expect("must contain some array lengths!"); + for i in (0..lengths.len()-1).rev() { + total_offset += indexes[i] * size_multiplier; + size_multiplier *= lengths[i]; + } + total_offset +} + +#[cfg(test)] +mod test { + use super::compute_offset; + + #[test] + fn test_expected_offset() { + let offset = compute_offset(&vec![1, 1], &vec![5, 3]); + assert_eq!(4, offset); + } + + #[test] + fn test_offsets() { + let lengths = vec![5, 3, 7]; + for i in 0..lengths[0] { + for j in 0..lengths[1] { + for k in 0..lengths[2] { + let offset = compute_offset(&vec![i, j, k], &lengths); + assert_eq!((i * 21) + (j * 7) + k, offset, "i={}, j={}, k={}", i, j, k); + } + } + } + } + + #[test] + fn test_increments() { + let lengths = vec![5, 7]; + for i in 0..lengths[0] { + for j in 0..lengths[1]-1 { + let offset = compute_offset(&vec![i, j], &lengths); + let next_offset = compute_offset(&vec![i, j + 1], &lengths); + assert_eq!(offset + 1, next_offset); + } + } + } +} \ No newline at end of file diff --git a/circuit_passes/src/passes/mapped_to_indexed.rs b/circuit_passes/src/passes/mapped_to_indexed.rs index 027ba59fa..7712f2b3b 100644 --- a/circuit_passes/src/passes/mapped_to_indexed.rs +++ b/circuit_passes/src/passes/mapped_to_indexed.rs @@ -6,6 +6,7 @@ use compiler::intermediate_representation::ir_interface::*; use compiler::intermediate_representation::{InstructionPointer, UpdateId}; use crate::bucket_interpreter::env::Env; use crate::bucket_interpreter::observer::InterpreterObserver; +use crate::bucket_interpreter::operations::compute_offset; use crate::bucket_interpreter::value::Value::KnownU32; use crate::passes::CircuitTransformationPass; use crate::passes::memory::PassMemory; @@ -35,7 +36,8 @@ impl MappedToIndexedPass { let mut acc_env = acc_env; let name = acc_env.get_subcmp_name(resolved_addr).clone(); - let map_access = mem.io_map[&acc_env.get_subcmp_template_id(resolved_addr)][signal_code].offset; + let io_def = &mem.io_map[&acc_env.get_subcmp_template_id(resolved_addr)][signal_code]; + let map_access = io_def.offset; if indexes.len() > 0 { let mut indexes_values = vec![]; for i in indexes { @@ -43,15 +45,11 @@ impl MappedToIndexedPass { indexes_values.push(val.expect("Mapped location must produce a value!").get_u32()); acc_env = new_env; } - if indexes.len() == 1 { - let value = map_access + indexes_values[0]; - let mut unused = vec![]; - LocationRule::Indexed { - location: KnownU32(value).to_value_bucket(&mut unused).allocate(), - template_header: Some(name), - } - } else { - todo!() + let offset = compute_offset(&indexes_values, &io_def.lengths); + let mut unused = vec![]; + LocationRule::Indexed { + location: KnownU32(map_access + offset).to_value_bucket(&mut unused).allocate(), + template_header: Some(name), } } else { let mut unused = vec![];