Fix precompiled assembly mem reads (#328)

* Fix precompiled assembly mem reads

Author: fractasy

core/src/zisk_rom.rs

@@ -3286,16 +3286,13 @@ impl ZiskRom {
         for j in 0..2 {
             // For every indirection
             for i in 0..indirections_count {
-                if i >= load_count {
-                    break;
-                }
                 // Store next aligned address value in mem_reads, and advance it
                 *s += &format!(
                     "\tmov {}, [{} + {}*8] /* value = mem[address+{}] */\n",
                     REG_VALUE, REG_ADDRESS, i, i
                 );
 
-                // During the first iteration, store the indirectionread value in mem_reads
+                // During the first iteration, store the indirection read value in mem_reads
                 if j == 0 {
                     *s += &format!(
                         "\tmov [{} + {}*8 + {}*8], {} /* mem_reads[@+size*8+ind*8] = ind */\n",
@@ -3307,6 +3304,12 @@ impl ZiskRom {
                 // During the second iteration, store the first load_count iterations
                 // load_size elements in mem_reads
                 if j == 1 {
+                    // Only store the first load_count indirections
+                    if i >= load_count {
+                        break;
+                    }
+
+                    // For each chunk of the indirection, store it in mem_reads
                     for l in 0..load_size {
                         *s += &format!(
                             "\tmov {}, [{} + {}*8] /* aux = mem[ind+{}] */\n",

emulator-asm/src/main.c

@@ -1226,36 +1226,36 @@ void log_trace(void)
 
         // Log current chunk start state
         printf("\tStart state:\n");
-        printf("\t\tpc=0x%lx:\n", chunk[i]);
+        printf("\t\tpc=0x%lx\n", chunk[i]);
         i++;
-        printf("\t\tsp=0x%lx:\n", chunk[i]);
+        printf("\t\tsp=0x%lx\n", chunk[i]);
         i++;
-        printf("\t\tc=0x%lx:\n", chunk[i]);
+        printf("\t\tc=0x%lx\n", chunk[i]);
         i++;
-        printf("\t\tstep=%ld:\n", chunk[i]);
+        printf("\t\tstep=%ld\n", chunk[i]);
         i++;
         for (uint64_t r=1; r<34; r++)
         {
-            printf("\t\tregister[%ld]=0x%ld:\n", r, chunk[i]);
+            printf("\t\tregister[%ld]=0x%lx\n", r, chunk[i]);
             i++;
         }
 
         // Log current chunk last state
         printf("\tLast state:\n");
-        printf("\t\tc=0x%lx:\n", chunk[i]);
+        printf("\t\tc=0x%lx\n", chunk[i]);
         i++;
 
         // Log current chunk end
         printf("\tEnd:\n");
-        printf("\t\tend=%ld:\n", chunk[i]);
+        printf("\t\tend=%ld\n", chunk[i]);
         i++;
 
         // Log current chunk steps
         printf("\tSteps:\n");
-        printf("\t\tsteps=%ld:\n", chunk[i]);
+        printf("\t\tsteps=%ld\n", chunk[i]);
         i++;
         uint64_t mem_reads_size = chunk[i];
-        printf("\t\tmem_reads_size=%ld:\n", mem_reads_size);
+        printf("\t\tmem_reads_size=%ld\n", mem_reads_size);
         i++;
         if (mem_reads_size > 10000000)
         {
@@ -1266,7 +1266,7 @@ void log_trace(void)
         {
             for (uint64_t m=0; m<mem_reads_size; m++)
             {
-                printf("\t\tchunk[%ld].mem_reads[%ld]=%08lx:\n", c, m, chunk[i]);
+                printf("\t\tchunk[%ld].mem_reads[%ld]=%08lx\n", c, m, chunk[i]);
                 i++;
             }
         }

emulator/src/emu.rs

@@ -1184,7 +1184,7 @@ impl<'a> Emu<'a> {
         callback: Option<impl Fn(EmuTrace)>,
     ) {
         // Context, where the state of the execution is stored and modified at every execution step
-        self.ctx = self.create_emu_context(inputs);
+        self.ctx = self.create_emu_context(inputs.clone());
 
         // Check that callback is provided if trace_steps is specified
         if options.trace_steps.is_some() {
@@ -1213,6 +1213,35 @@ impl<'a> Emu<'a> {
         if options.is_fast() {
             return self.run_fast(options);
         }
+        if options.generate_minimal_traces {
+            let par_emu_options =
+                ParEmuOptions { num_steps: 1024 * 1024, num_threads: 1, thread_id: 0 };
+            let minimal_trace = self.run_gen_trace(inputs, options, &par_emu_options);
+
+            for (c, chunk) in minimal_trace.iter().enumerate() {
+                println!("Chunk {}:", c);
+                println!("\tStart state:");
+                println!("\t\tpc=0x{:x}", chunk.start_state.pc);
+                println!("\t\tsp=0x{:x}", chunk.start_state.sp);
+                println!("\t\tc=0x{:x}", chunk.start_state.c);
+                println!("\t\tstep={}", chunk.start_state.step);
+                for i in 1..32 {
+                    // TODO: 34
+                    println!("\t\tregister[{}]=0x{:x}", i, chunk.start_state.regs[i]);
+                }
+                println!("\tLast state:");
+                println!("\t\tc=0x{:x}", chunk.last_c);
+                println!("\tEnd:");
+                println!("\t\tend={}", if chunk.end { 1 } else { 0 });
+                println!("\tSteps:");
+                println!("\t\tsteps={}", chunk.steps);
+                println!("\t\tmem_reads_size={}", chunk.mem_reads.len());
+                for i in 0..chunk.mem_reads.len() {
+                    println!("\t\tchunk[{}].mem_reads[{}]={:08x}", c, i, chunk.mem_reads[i]);
+                }
+            }
+            return;
+        }
         //println!("Emu::run() full-equipe");
 
         // Store the stats option into the emulator context
@@ -1291,7 +1320,7 @@ impl<'a> Emu<'a> {
     }
 
     /// Run the whole program
-    pub fn par_run<F: PrimeField>(
+    pub fn par_run(
         &mut self,
         inputs: Vec<u8>,
         options: &EmuOptions,
@@ -1335,7 +1364,8 @@ impl<'a> Emu<'a> {
                         end: false,
                     });
                 }
-                self.par_step_my_block::<F>(emu_traces.last_mut().unwrap());
+
+                self.par_step_my_block(emu_traces.last_mut().unwrap());
 
                 if self.ctx.inst_ctx.step >= options.max_steps {
                     panic!("Emu::par_run() reached max_steps");
@@ -1346,6 +1376,55 @@ impl<'a> Emu<'a> {
         emu_traces
     }
 
+    /// Run the whole program
+    pub fn run_gen_trace(
+        &mut self,
+        inputs: Vec<u8>,
+        options: &EmuOptions,
+        par_options: &ParEmuOptions,
+    ) -> Vec<EmuTrace> {
+        // Context, where the state of the execution is stored and modified at every execution step
+        self.ctx = self.create_emu_context(inputs);
+
+        // Init pc to the rom entry address
+        self.ctx.trace.start_state.pc = ROM_ENTRY;
+
+        // Store the stats option into the emulator context
+        self.ctx.do_stats = options.stats;
+
+        // Set emulation mode
+        self.ctx.inst_ctx.emulation_mode = EmulationMode::GenerateMemReads;
+
+        let mut emu_traces = Vec::new();
+
+        while !self.ctx.inst_ctx.end {
+            // Check if is the first step of a new block
+            if self.ctx.inst_ctx.step % par_options.num_steps as u64 == 0 {
+                emu_traces.push(EmuTrace {
+                    start_state: EmuTraceStart {
+                        pc: self.ctx.inst_ctx.pc,
+                        sp: self.ctx.inst_ctx.sp,
+                        c: self.ctx.inst_ctx.c,
+                        step: self.ctx.inst_ctx.step,
+                        regs: self.ctx.inst_ctx.regs,
+                    },
+                    last_c: 0,
+                    steps: 0,
+                    mem_reads: Vec::with_capacity(par_options.num_steps),
+                    end: false,
+                });
+            }
+
+            self.par_step_my_block(emu_traces.last_mut().unwrap());
+
+            if self.ctx.inst_ctx.step >= options.max_steps {
+                panic!("Emu::par_run() reached max_steps");
+            }
+        }
+
+        emu_traces
+    }
+
     /// Performs one single step of the emulation
     #[inline(always)]
     #[allow(unused_variables)]
@@ -1454,7 +1533,7 @@ impl<'a> Emu<'a> {
 
     /// Performs one single step of the emulation
     #[inline(always)]
-    pub fn par_step_my_block<F: PrimeField>(&mut self, emu_full_trace_vec: &mut EmuTrace) {
+    pub fn par_step_my_block(&mut self, emu_full_trace_vec: &mut EmuTrace) {
         let instruction = self.rom.get_instruction(self.ctx.inst_ctx.pc);
         // Build the 'a' register value  based on the source specified by the current instruction
         self.source_a_mem_reads_generate(instruction, &mut emu_full_trace_vec.mem_reads);

emulator/src/emu_options.rs

@@ -61,6 +61,9 @@ pub struct EmuOptions {
     /// Generates statistics about opcodes and memory usage.  Enabled with `-x`.
     #[clap(short = 'x', long, value_name = "STATS", default_value = "false")]
     pub stats: bool,
+    /// Generates minimal traces.  Enabled with `-g`.
+    #[clap(short = 'g', long, value_name = "MINIMAL_TRACES", default_value = "false")]
+    pub generate_minimal_traces: bool,
 }
 
 impl Default for EmuOptions {
@@ -81,6 +84,7 @@ impl Default for EmuOptions {
             log_metrics: false,
             tracerv: false,
             stats: false,
+            generate_minimal_traces: false,
         }
     }
 }
@@ -101,7 +105,8 @@ impl fmt::Display for EmuOptions {
         writeln!(f, "METRICS: {:?}", self.log_metrics)?;
         writeln!(f, "STATS: {:?}", self.stats)?;
         writeln!(f, "TRACERV: {:?}", self.tracerv)?;
-        writeln!(f, "LOG_STEP: {:?}", self.log_step)
+        writeln!(f, "LOG_STEP: {:?}", self.log_step)?;
+        writeln!(f, "MINIMAL_TRACES: {:?}", self.generate_minimal_traces)
     }
 }
 
@@ -115,5 +120,6 @@ impl EmuOptions {
             && !self.verbose
             && !self.tracerv
             && !self.stats
+            && !self.generate_minimal_traces
     }
 }

emulator/src/emulator.rs

@@ -177,7 +177,7 @@ impl ZiskEmulator {
     /// First phase of the witness computation
     /// 8 threads in waterfall (# threads to be re-calibrated after memory reads refactor)
     /// Must be fast
-    pub fn compute_minimal_traces<F: PrimeField>(
+    pub fn compute_minimal_traces(
         rom: &ZiskRom,
         inputs: &[u8],
         options: &EmuOptions,
@@ -191,7 +191,7 @@ impl ZiskEmulator {
 
             // Run the emulation
             let mut emu = Emu::new(rom);
-            let result = emu.par_run::<F>(inputs.to_owned(), options, &par_emu_options);
+            let result = emu.par_run(inputs.to_owned(), options, &par_emu_options);
 
             if !emu.terminated() {
                 panic!("Emulation did not complete");

executor/src/executor.rs

@@ -197,7 +197,7 @@ impl<F: PrimeField64> ZiskExecutor<F> {
             ..EmuOptions::default()
         };
 
-        let min_traces = ZiskEmulator::compute_minimal_traces::<F>(
+        let min_traces = ZiskEmulator::compute_minimal_traces(
             &self.zisk_rom,
             &input_data,
             &emu_options,