done

paradigmxyz · Apr 14, 2024 · 5cb946c · 5cb946c
1 parent ecfd10c
commit 5cb946c
Show file tree

Hide file tree

Showing 7 changed files with 158 additions and 67 deletions.
diff --git a/benches/benches/bench.rs b/benches/benches/bench.rs
@@ -73,12 +73,14 @@ fn run_bench(c: &mut Criterion, bench: &Bench) {
         r
     };
 
-    let jit_no_gas = jit.compile(Some(name), bytecode, SPEC_ID).unwrap();
+    let jit_no_gas = jit.translate(Some(name), bytecode, SPEC_ID).unwrap();
+    jit.set_disable_gas(false);
+    let jit_gas = jit.translate(Some(name), bytecode, SPEC_ID).unwrap();
+    let jit_no_gas = jit.jit_function(jit_no_gas).unwrap();
+    let jit_gas = jit.jit_function(jit_gas).unwrap();
+
     g.bench_function("revm-jit/no_gas", |b| b.iter(|| call_jit(jit_no_gas)));
 
-    unsafe { jit.free_all_functions() }.unwrap();
-    jit.set_disable_gas(false);
-    let jit_gas = jit.compile(Some(name), bytecode, SPEC_ID).unwrap();
     g.bench_function("revm-jit/gas", |b| b.iter(|| call_jit(jit_gas)));
 
     g.bench_function("revm-interpreter", |b| {

diff --git a/crates/revm-jit-backend/src/traits.rs b/crates/revm-jit-backend/src/traits.rs
@@ -131,10 +131,12 @@ pub trait Backend: BackendTypes + TypeMethods {
     >
     where
         Self: 'a;
-    type FuncId: Copy + Eq + fmt::Debug;
+    type FuncId: Copy + Eq + std::hash::Hash + fmt::Debug;
 
     fn ir_extension(&self) -> &'static str;
 
+    fn set_module_name(&mut self, name: &str);
+
     fn set_is_dumping(&mut self, yes: bool);
     fn set_debug_assertions(&mut self, yes: bool);
     fn opt_level(&self) -> OptimizationLevel;
@@ -151,8 +153,8 @@ pub trait Backend: BackendTypes + TypeMethods {
         linkage: Linkage,
     ) -> Result<(Self::Builder<'_>, Self::FuncId)>;
     fn verify_module(&mut self) -> Result<()>;
-    fn optimize_function(&mut self, id: Self::FuncId) -> Result<()>;
-    fn get_function(&mut self, id: Self::FuncId) -> Result<usize>;
+    fn optimize_module(&mut self) -> Result<()>;
+    fn jit_function(&mut self, id: Self::FuncId) -> Result<usize>;
     unsafe fn free_function(&mut self, id: Self::FuncId) -> Result<()>;
     unsafe fn free_all_functions(&mut self) -> Result<()>;
 }

diff --git a/crates/revm-jit-cranelift/src/lib.rs b/crates/revm-jit-cranelift/src/lib.rs
@@ -40,6 +40,7 @@ pub struct JitEvmCraneliftBackend {
 
     opt_level: OptimizationLevel,
     comments: CommentWriter,
+    functions: Vec<FuncId>,
 }
 
 #[allow(clippy::new_without_default)]
@@ -67,6 +68,7 @@ impl JitEvmCraneliftBackend {
             symbols,
             opt_level,
             comments: CommentWriter::new(),
+            functions: Vec::new(),
         }
     }
 }
@@ -112,6 +114,10 @@ impl Backend for JitEvmCraneliftBackend {
         "clif"
     }
 
+    fn set_module_name(&mut self, name: &str) {
+        let _ = name;
+    }
+
     fn set_is_dumping(&mut self, yes: bool) {
         self.ctx.set_disasm(yes);
     }
@@ -168,6 +174,7 @@ impl Backend for JitEvmCraneliftBackend {
             convert_linkage(linkage),
             &self.ctx.func.signature,
         )?;
+        self.functions.push(id);
         let bcx = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context);
         let builder = JitEvmCraneliftBuilder {
             module: &mut self.module,
@@ -183,13 +190,16 @@ impl Backend for JitEvmCraneliftBackend {
         Ok(())
     }
 
-    fn optimize_function(&mut self, id: Self::FuncId) -> Result<()> {
+    fn optimize_module(&mut self) -> Result<()> {
         // Define the function to jit. This finishes compilation, although
         // there may be outstanding relocations to perform. Currently, jit
         // cannot finish relocations until all functions to be called are
         // defined. For this toy demo for now, we'll just finalize the
         // function below.
-        self.module.define_function(id, &mut self.ctx)?;
+        for &id in &self.functions {
+            self.module.define_function(id, &mut self.ctx)?;
+        }
+        self.functions.clear();
 
         // Now that compilation is finished, we can clear out the context state.
         self.module.clear_context(&mut self.ctx);
@@ -203,7 +213,7 @@ impl Backend for JitEvmCraneliftBackend {
         Ok(())
     }
 
-    fn get_function(&mut self, id: Self::FuncId) -> Result<usize> {
+    fn jit_function(&mut self, id: Self::FuncId) -> Result<usize> {
         Ok(self.module.get_finalized_function(id) as usize)
     }
 

diff --git a/crates/revm-jit-llvm/src/lib.rs b/crates/revm-jit-llvm/src/lib.rs
@@ -49,7 +49,6 @@ pub struct JitEvmLlvmBackend<'ctx> {
 
     debug_assertions: bool,
     opt_level: OptimizationLevel,
-    function_counter: u32,
     function_names: FxHashMap<u32, String>,
 }
 
@@ -143,7 +142,6 @@ impl<'ctx> JitEvmLlvmBackend<'ctx> {
             ty_ptr,
             debug_assertions: cfg!(debug_assertions),
             opt_level,
-            function_counter: 0,
             function_names: FxHashMap::default(),
         })
     }
@@ -219,6 +217,10 @@ impl<'ctx> Backend for JitEvmLlvmBackend<'ctx> {
         "ll"
     }
 
+    fn set_module_name(&mut self, name: &str) {
+        self.module.set_name(name);
+    }
+
     fn set_is_dumping(&mut self, yes: bool) {
         self.machine.set_asm_verbosity(yes);
     }
@@ -260,8 +262,8 @@ impl<'ctx> Backend for JitEvmLlvmBackend<'ctx> {
         let entry = self.cx.append_basic_block(function, "entry");
         self.bcx.position_at_end(entry);
 
-        let id = self.function_counter;
-        self.function_counter += 1;
+        let id = self.function_names.len() as u32;
+        self.function_names.insert(id, name.to_string());
         let builder = JitEvmLlvmBuilder { backend: self, function };
         Ok((builder, id))
     }
@@ -270,7 +272,7 @@ impl<'ctx> Backend for JitEvmLlvmBackend<'ctx> {
         self.module.verify().map_err(error_msg)
     }
 
-    fn optimize_function(&mut self, _id: Self::FuncId) -> Result<()> {
+    fn optimize_module(&mut self) -> Result<()> {
         // From `opt --help`, `-passes`.
         let passes = match self.opt_level {
             OptimizationLevel::None => "default<O0>",
@@ -282,7 +284,7 @@ impl<'ctx> Backend for JitEvmLlvmBackend<'ctx> {
         self.module.run_passes(passes, &self.machine, opts).map_err(error_msg)
     }
 
-    fn get_function(&mut self, id: Self::FuncId) -> Result<usize> {
+    fn jit_function(&mut self, id: Self::FuncId) -> Result<usize> {
         let name = self.id_to_name(id);
         self.exec_engine.get_function_address(name).map_err(Into::into)
     }
@@ -291,6 +293,7 @@ impl<'ctx> Backend for JitEvmLlvmBackend<'ctx> {
         let name = self.id_to_name(id);
         let function = self.exec_engine.get_function_value(name)?;
         self.exec_engine.free_fn_machine_code(function);
+        self.function_names.clear();
         Ok(())
     }