[VAN-676][VAN-670] flatten branches in extracted function by duplicat…

…ing the functions (#55) * Refactor StoreBucket to avoid generating call to "llvm.donothing"
Veridise · Nov 8, 2023 · 8f71b37 · 8f71b37
1 parent 632ebbb
commit 8f71b37
Show file tree

Hide file tree

Showing 11 changed files with 193 additions and 22 deletions.
diff --git a/circom/tests/loops/assign_in_loop.circom b/circom/tests/loops/assign_in_loop.circom
@@ -0,0 +1,186 @@
+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 --enable-var-scope
+// XFAIL:.* // panicked at 'not yet implemented', circuit_passes/src/passes/loop_unroll/loop_env_recorder.rs:149:44 (LocationRule::Mapped case)
+
+template Inner(i) {
+    signal input in;
+    signal output out;
+
+    out <-- (in >> i) & 1;
+}
+
+template Num2Bits(n) {
+    signal input in;
+    signal output out[n];
+
+    component c[n];
+    for (var i = 0; i < n; i++) {
+    	c[i] = Inner(i);
+    	c[i].in <-- in;
+    	out[i] <-- c[i].out;
+    }
+}
+
+component main = Num2Bits(3);
+
+//CHECK-LABEL: define void @..generated..loop.body.
+//CHECK-SAME: [[$F_ID_1:[0-9]+]]([0 x i256]* %lvars, [0 x i256]* %signals, i256* %subfix_[[X1:[0-9]+]], i256* %fix_[[X2:[0-9]+]], i256* %subfix_[[X3:[0-9]+]],
+//CHECK-SAME: [0 x i256]* %sub_[[X1]], i256* %subc_[[X1]], [0 x i256]* %sub_[[X3]], i256* %subc_[[X3]]){{.*}} {
+//CHECK-NEXT: ..generated..loop.body.[[$F_ID_1]]:
+//CHECK-NEXT:   br label %store1
+//CHECK-EMPTY: 
+//CHECK-NEXT: store1:
+//CHECK-NEXT:   %0 = getelementptr [0 x i256], [0 x i256]* %signals, i32 0, i32 3
+//CHECK-NEXT:   %1 = load i256, i256* %0, align 4
+//CHECK-NEXT:   %2 = getelementptr i256, i256* %subfix_[[X1]], i32 0
+//CHECK-NEXT:   store i256 %1, i256* %2, align 4
+//CHECK-NEXT:   br label %store2
+//CHECK-EMPTY: 
+//CHECK-NEXT: store2:
+//CHECK-NEXT:   %3 = getelementptr [0 x i256], [0 x i256]* %sub_[[X1]], i32 0
+//CHECK-NEXT:   call void @Inner_?_run([0 x i256]* %sub_[[X1]])                 //TODO: which function to call depends on which iteration of the loop
+//CHECK-NEXT:   br label %store3
+//CHECK-EMPTY: 
+//CHECK-NEXT: store3:
+//CHECK-NEXT:   %4 = getelementptr i256, i256* %subfix_[[X3]], i32 0
+//CHECK-NEXT:   %5 = load i256, i256* %4, align 4
+//CHECK-NEXT:   %6 = getelementptr i256, i256* %fix_[[X2]], i32 0
+//CHECK-NEXT:   store i256 %5, i256* %6, align 4
+//CHECK-NEXT:   br label %store4
+//CHECK-EMPTY: 
+//CHECK-NEXT: store4:
+//CHECK-NEXT:   %7 = getelementptr [0 x i256], [0 x i256]* %lvars, i32 0, i32 1
+//CHECK-NEXT:   %8 = load i256, i256* %7, align 4
+//CHECK-NEXT:   %call.fr_add = call i256 @fr_add(i256 %8, i256 1)
+//CHECK-NEXT:   %9 = getelementptr [0 x i256], [0 x i256]* %lvars, i32 0, i32 1
+//CHECK-NEXT:   store i256 %call.fr_add, i256* %9, align 4
+//CHECK-NEXT:   br label %return5
+//CHECK-EMPTY: 
+//CHECK-NEXT: return5:
+//CHECK-NEXT:   ret void
+//CHECK-NEXT: }
+//
+//CHECK-LABEL: define void @Inner_0_build({ [0 x i256]*, i32 }* %0){{.*}} {
+//
+//CHECK-LABEL: define void @Inner_0_run([0 x i256]* %0){{.*}} {
+//
+//CHECK-LABEL: define void @Inner_1_build({ [0 x i256]*, i32 }* %0){{.*}} {
+//
+//CHECK-LABEL: define void @Inner_1_run([0 x i256]* %0){{.*}} {
+//
+//CHECK-LABEL: define void @Inner_2_build({ [0 x i256]*, i32 }* %0){{.*}} {
+//
+//CHECK-LABEL: define void @Inner_2_run([0 x i256]* %0){{.*}} {
+//
+//CHECK-LABEL: define void @Num2Bits_3_build({ [0 x i256]*, i32 }* %0){{.*}} {
+//CHECK-NEXT: main:
+//CHECK-NEXT:   %1 = alloca [4 x i256], align 8
+//CHECK-NEXT:   %2 = getelementptr { [0 x i256]*, i32 }, { [0 x i256]*, i32 }* %0, i32 0, i32 1
+//CHECK-NEXT:   store i32 1, i32* %2, align 4
+//CHECK-NEXT:   %3 = getelementptr { [0 x i256]*, i32 }, { [0 x i256]*, i32 }* %0, i32 0, i32 0
+//CHECK-NEXT:   %4 = bitcast [4 x i256]* %1 to [0 x i256]*
+//CHECK-NEXT:   store [0 x i256]* %4, [0 x i256]** %3, align 8
+//CHECK-NEXT:   ret void
+//CHECK-NEXT: }
+//
+//CHECK-LABEL: define void @Num2Bits_3_run([0 x i256]* %0){{.*}} {
+//CHECK-NEXT: prelude:
+//CHECK-NEXT:   %lvars = alloca [2 x i256], align 8
+//CHECK-NEXT:   %subcmps = alloca [3 x { [0 x i256]*, i32 }], align 8
+//CHECK-NEXT:   br label %store1
+//CHECK-EMPTY: 
+//CHECK-NEXT: store1:
+//CHECK-NEXT:   %1 = getelementptr [2 x i256], [2 x i256]* %lvars, i32 0, i32 0
+//CHECK-NEXT:   store i256 3, i256* %1, align 4
+//CHECK-NEXT:   br label %create_cmp2
+//CHECK-EMPTY: 
+//CHECK-NEXT: create_cmp2:
+//CHECK-NEXT:   %2 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 0
+//CHECK-NEXT:   call void @Inner_0_build({ [0 x i256]*, i32 }* %2)
+//CHECK-NEXT:   br label %create_cmp3
+//CHECK-EMPTY: 
+//CHECK-NEXT: create_cmp3:
+//CHECK-NEXT:   %3 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 1
+//CHECK-NEXT:   call void @Inner_1_build({ [0 x i256]*, i32 }* %3)
+//CHECK-NEXT:   br label %create_cmp4
+//CHECK-EMPTY: 
+//CHECK-NEXT: create_cmp4:
+//CHECK-NEXT:   %4 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2
+//CHECK-NEXT:   call void @Inner_2_build({ [0 x i256]*, i32 }* %4)
+//CHECK-NEXT:   br label %store5
+//CHECK-EMPTY: 
+//CHECK-NEXT: store5:
+//CHECK-NEXT:   %5 = getelementptr [2 x i256], [2 x i256]* %lvars, i32 0, i32 1
+//CHECK-NEXT:   store i256 0, i256* %5, align 4
+//CHECK-NEXT:   br label %unrolled_loop6
+//CHECK-EMPTY: 
+//CHECK-NEXT: unrolled_loop6:
+//CHECK-NEXT:   %6 = bitcast [2 x i256]* %lvars to [0 x i256]*
+//CHECK-NEXT:   %7 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 0, i32 0
+//CHECK-NEXT:   %8 = load [0 x i256]*, [0 x i256]** %7, align 8
+//CHECK-NEXT:   %9 = getelementptr [0 x i256], [0 x i256]* %8, i32 0
+//CHECK-NEXT:   %10 = getelementptr [0 x i256], [0 x i256]* %9, i32 0, i256 1
+//CHECK-NEXT:   %11 = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i256 0
+//CHECK-NEXT:   %12 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 0, i32 0
+//CHECK-NEXT:   %13 = load [0 x i256]*, [0 x i256]** %12, align 8
+//CHECK-NEXT:   %14 = getelementptr [0 x i256], [0 x i256]* %13, i32 0
+//CHECK-NEXT:   %15 = getelementptr [0 x i256], [0 x i256]* %14, i32 0, i256 0
+//CHECK-NEXT:   %16 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 0, i32 0
+//CHECK-NEXT:   %17 = load [0 x i256]*, [0 x i256]** %16, align 8
+//CHECK-NEXT:   %18 = getelementptr [0 x i256], [0 x i256]* %17, i32 0
+//CHECK-NEXT:   %19 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 0, i32 1
+//CHECK-NEXT:   %20 = bitcast i32* %19 to i256*
+//CHECK-NEXT:   %21 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 0, i32 0
+//CHECK-NEXT:   %22 = load [0 x i256]*, [0 x i256]** %21, align 8
+//CHECK-NEXT:   %23 = getelementptr [0 x i256], [0 x i256]* %22, i32 0
+//CHECK-NEXT:   %24 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 0, i32 1
+//CHECK-NEXT:   %25 = bitcast i32* %24 to i256*
+//CHECK-NEXT:   call void @..generated..loop.body.[[$F_ID_1]]([0 x i256]* %6, [0 x i256]* %0, i256* %10, i256* %11, i256* %15, [0 x i256]* %18, i256* %20, [0 x i256]* %23, i256* %25)
+//CHECK-NEXT:   %26 = bitcast [2 x i256]* %lvars to [0 x i256]*
+//CHECK-NEXT:   %27 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 1, i32 0
+//CHECK-NEXT:   %28 = load [0 x i256]*, [0 x i256]** %27, align 8
+//CHECK-NEXT:   %29 = getelementptr [0 x i256], [0 x i256]* %28, i32 0
+//CHECK-NEXT:   %30 = getelementptr [0 x i256], [0 x i256]* %29, i32 0, i256 1
+//CHECK-NEXT:   %31 = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i256 1
+//CHECK-NEXT:   %32 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 1, i32 0
+//CHECK-NEXT:   %33 = load [0 x i256]*, [0 x i256]** %32, align 8
+//CHECK-NEXT:   %34 = getelementptr [0 x i256], [0 x i256]* %33, i32 0
+//CHECK-NEXT:   %35 = getelementptr [0 x i256], [0 x i256]* %34, i32 0, i256 0
+//CHECK-NEXT:   %36 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 1, i32 0
+//CHECK-NEXT:   %37 = load [0 x i256]*, [0 x i256]** %36, align 8
+//CHECK-NEXT:   %38 = getelementptr [0 x i256], [0 x i256]* %37, i32 0
+//CHECK-NEXT:   %39 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 1, i32 1
+//CHECK-NEXT:   %40 = bitcast i32* %39 to i256*
+//CHECK-NEXT:   %41 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 1, i32 0
+//CHECK-NEXT:   %42 = load [0 x i256]*, [0 x i256]** %41, align 8
+//CHECK-NEXT:   %43 = getelementptr [0 x i256], [0 x i256]* %42, i32 0
+//CHECK-NEXT:   %44 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 1, i32 1
+//CHECK-NEXT:   %45 = bitcast i32* %44 to i256*
+//CHECK-NEXT:   call void @..generated..loop.body.[[$F_ID_1]]([0 x i256]* %26, [0 x i256]* %0, i256* %30, i256* %31, i256* %35, [0 x i256]* %38, i256* %40, [0 x i256]* %43, i256* %45)
+//CHECK-NEXT:   %46 = bitcast [2 x i256]* %lvars to [0 x i256]*
+//CHECK-NEXT:   %47 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0
+//CHECK-NEXT:   %48 = load [0 x i256]*, [0 x i256]** %47, align 8
+//CHECK-NEXT:   %49 = getelementptr [0 x i256], [0 x i256]* %48, i32 0
+//CHECK-NEXT:   %50 = getelementptr [0 x i256], [0 x i256]* %49, i32 0, i256 1
+//CHECK-NEXT:   %51 = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i256 2
+//CHECK-NEXT:   %52 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0
+//CHECK-NEXT:   %53 = load [0 x i256]*, [0 x i256]** %52, align 8
+//CHECK-NEXT:   %54 = getelementptr [0 x i256], [0 x i256]* %53, i32 0
+//CHECK-NEXT:   %55 = getelementptr [0 x i256], [0 x i256]* %54, i32 0, i256 0
+//CHECK-NEXT:   %56 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0
+//CHECK-NEXT:   %57 = load [0 x i256]*, [0 x i256]** %56, align 8
+//CHECK-NEXT:   %58 = getelementptr [0 x i256], [0 x i256]* %57, i32 0
+//CHECK-NEXT:   %59 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 1
+//CHECK-NEXT:   %60 = bitcast i32* %59 to i256*
+//CHECK-NEXT:   %61 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0
+//CHECK-NEXT:   %62 = load [0 x i256]*, [0 x i256]** %61, align 8
+//CHECK-NEXT:   %63 = getelementptr [0 x i256], [0 x i256]* %62, i32 0
+//CHECK-NEXT:   %64 = getelementptr [3 x { [0 x i256]*, i32 }], [3 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 1
+//CHECK-NEXT:   %65 = bitcast i32* %64 to i256*
+//CHECK-NEXT:   call void @..generated..loop.body.[[$F_ID_1]]([0 x i256]* %46, [0 x i256]* %0, i256* %50, i256* %51, i256* %55, [0 x i256]* %58, i256* %60, [0 x i256]* %63, i256* %65)
+//CHECK-NEXT:   br label %prologue
+//CHECK-EMPTY: 
+//CHECK-NEXT: prologue:
+//CHECK-NEXT:   ret void
+//CHECK-NEXT: }
diff --git a/circom/tests/loops/assign_in_loop_1.circom b/circom/tests/loops/assign_in_loop_1.circom
@@ -44,7 +44,6 @@ component main = Num2Bits(3);
 //CHECK-NEXT:   br label %fold_true3
 //CHECK-EMPTY: 
 //CHECK-NEXT: fold_true3:
-//CHECK-NEXT:   call void @llvm.donothing()
 //CHECK-NEXT:   call void @Inner_0_run([0 x i256]* %sub_[[X3]])
 //CHECK-NEXT:   br label %store4
 //CHECK-EMPTY: 

diff --git a/circom/tests/subcmps/subcmps0A.circom b/circom/tests/subcmps/subcmps0A.circom
@@ -45,7 +45,6 @@ component main = SubCmps0A(2);
 //CHECK-NEXT:   br label %fold_true3
 //CHECK-EMPTY: 
 //CHECK-NEXT: fold_true3:
-//CHECK-NEXT:   call void @llvm.donothing()
 //CHECK-NEXT:   call void @IsZero_0_run([0 x i256]* %sub_[[X4]])
 //CHECK-NEXT:   br label %store4
 //CHECK-EMPTY: 

diff --git a/circom/tests/subcmps/subcmps0B.circom b/circom/tests/subcmps/subcmps0B.circom
@@ -46,7 +46,6 @@ component main = SubCmps0B(2);
 //CHECK-NEXT:   br label %fold_true3
 //CHECK-EMPTY: 
 //CHECK-NEXT: fold_true3:
-//CHECK-NEXT:   call void @llvm.donothing()
 //CHECK-NEXT:   call void @IsZero_0_run([0 x i256]* %sub_[[X5]])
 //CHECK-NEXT:   br label %store4
 //CHECK-EMPTY: 

diff --git a/circom/tests/subcmps/subcmps0C.circom b/circom/tests/subcmps/subcmps0C.circom
@@ -44,7 +44,6 @@ component main = SubCmps0C(2);
 //CHECK-NEXT:   br label %fold_true3
 //CHECK-EMPTY: 
 //CHECK-NEXT: fold_true3:
-//CHECK-NEXT:   call void @llvm.donothing()
 //CHECK-NEXT:   call void @IsZero_0_run([0 x i256]* %sub_[[X4]])
 //CHECK-NEXT:   br label %store4
 //CHECK-EMPTY: 

diff --git a/circom/tests/subcmps/subcmps0D.circom b/circom/tests/subcmps/subcmps0D.circom
@@ -62,7 +62,6 @@ component main = SubCmps0D(3);
 //CHECK-NEXT:   br label %fold_true6
 //CHECK-EMPTY: 
 //CHECK-NEXT: fold_true6:
-//CHECK-NEXT:   call void @llvm.donothing()
 //CHECK-NEXT:   call void @Add_0_run([0 x i256]* %sub_[[X6]])
 //CHECK-NEXT:   br label %store7
 //CHECK-EMPTY: 

diff --git a/circom/tests/subcmps/subcmps1.circom b/circom/tests/subcmps/subcmps1.circom
@@ -58,7 +58,6 @@ component main = SubCmps1(3);
 //CHECK-NEXT:   br label %fold_true3
 //CHECK-EMPTY: 
 //CHECK-NEXT: fold_true3:
-//CHECK-NEXT:   call void @llvm.donothing()
 //CHECK-NEXT:   call void @IsZero_0_run([0 x i256]* %sub_[[X4]])
 //CHECK-NEXT:   br label %store4
 //CHECK-EMPTY: 

diff --git a/circom/tests/subcmps/subcmps2.circom b/circom/tests/subcmps/subcmps2.circom
@@ -131,7 +131,6 @@ component main = Caller();
 //CHECK-NEXT:   br label %fold_true3
 //CHECK-EMPTY: 
 //CHECK-NEXT: fold_true3:
-//CHECK-NEXT:   call void @llvm.donothing()
 //CHECK-NEXT:   call void @Sum_0_run([0 x i256]* %sub_[[X3]])
 //CHECK-NEXT:   br label %store4
 //CHECK-EMPTY: 

diff --git a/circuit_passes/src/bucket_interpreter/mod.rs b/circuit_passes/src/bucket_interpreter/mod.rs
@@ -387,7 +387,6 @@ impl<'a: 'd, 'd> BucketInterpreter<'a, 'd> {
         env: Env<'env>,
         observe: bool,
     ) -> R<'env> {
-        // println!("Interpreter executing {:?}", bucket);
         let (src, env) = self.execute_instruction(&bucket.src, env, observe);
         let src = src.expect("src instruction in StoreBucket must produce a value!");
         let env =

diff --git a/circuit_passes/src/passes/loop_unroll/body_extractor.rs b/circuit_passes/src/passes/loop_unroll/body_extractor.rs
@@ -412,12 +412,6 @@ impl LoopBodyExtractor {
                 } else {
                     //Since SubcmpSignal is always added above, this should be unreachable.
                     unreachable!()
-                    // bucket_to_args.insert(
-                    //     *b,
-                    //     ArgIndex::SubCmp { signal: next_idx, arena: arena_idx, counter: counter_idx,
-                    //     },
-                    // );
-                    // next_idx += 1;
                 }
             }
         }

diff --git a/compiler/src/intermediate_representation/store_bucket.rs b/compiler/src/intermediate_representation/store_bucket.rs
@@ -112,8 +112,7 @@ impl StoreBucket {
             Some(name) => {
                 assert_eq!(1, context.size, "unhandled array store");
                 if name == LLVM_DONOTHING_FN_NAME {
-                    //LLVM equivalent of a "nop" instruction
-                    create_call(producer, LLVM_DONOTHING_FN_NAME, &[])
+                    None
                 } else {
                     let arr_ptr = match &dest_address_type {
                         AddressType::Variable => producer.body_ctx().get_variable_array(producer),
@@ -128,11 +127,11 @@ impl StoreBucket {
                     }
                     .into_pointer_value();
                     let arr_ptr = pointer_cast(producer, arr_ptr, array_ptr_ty(producer));
-                    create_call(
+                    Some(create_call(
                         producer,
                         name.as_str(),
                         &[arr_ptr.into(), dest_index.into(), source.into_int_value().into()],
-                    )
+                    ))
                 }
             }
             None => {
@@ -175,18 +174,18 @@ impl StoreBucket {
                             };
                         }
                     }
-                    create_call(
+                    Some(create_call(
                         producer,
                         FR_ARRAY_COPY_FN_NAME,
                         &[
                             source.into_pointer_value().into(),
                             dest_gep.into(),
                             create_literal_u32(producer, context.size as u64).into(),
                         ],
-                    )
+                    ))
                 } else {
                     // In the scalar case, just produce a store from the source value that was given
-                    create_store(producer, dest_gep, source)
+                    Some(create_store(producer, dest_gep, source))
                 }
             }
         };
@@ -245,7 +244,7 @@ impl StoreBucket {
                 }
             }
         }
-        Some(store)
+        store
     }
 }