Allow use of pinned heap registers

Author: shravanrn

benchmarks/lucet-benchmarks/src/context.rs

@@ -9,7 +9,8 @@ fn context_init(c: &mut Criterion) {
 
     c.bench_function("context_init", move |b| {
         b.iter(|| {
-            ContextHandle::create_and_init(&mut *stack, f as usize, &[]).unwrap();
+            ContextHandle::create_and_init(&mut *stack, f as usize, &[], std::ptr::null_mut())
+                .unwrap();
         })
     });
 }
@@ -22,7 +23,13 @@ fn context_swap_return(c: &mut Criterion) {
         b.iter_batched(
             || {
                 let mut stack = vec![0u64; 1024].into_boxed_slice();
-                let child = ContextHandle::create_and_init(&mut *stack, f as usize, &[]).unwrap();
+                let child = ContextHandle::create_and_init(
+                    &mut *stack,
+                    f as usize,
+                    &[],
+                    std::ptr::null_mut(),
+                )
+                .unwrap();
                 (stack, child)
             },
             |(stack, mut child)| unsafe {
@@ -44,8 +51,13 @@ fn context_init_swap_return(c: &mut Criterion) {
             || vec![0u64; 1024].into_boxed_slice(),
             |mut stack| {
                 let mut parent = ContextHandle::new();
-                let mut child =
-                    ContextHandle::create_and_init(&mut *stack, f as usize, &[]).unwrap();
+                let mut child = ContextHandle::create_and_init(
+                    &mut *stack,
+                    f as usize,
+                    &[],
+                    std::ptr::null_mut(),
+                )
+                .unwrap();
                 unsafe { Context::swap(&mut parent, &mut child) };
                 stack
             },
@@ -332,8 +344,13 @@ fn context_init_swap_return_many_args(c: &mut Criterion) {
             || vec![0u64; 1024].into_boxed_slice(),
             |mut stack| {
                 let mut parent = ContextHandle::new();
-                let mut child =
-                    ContextHandle::create_and_init(&mut *stack, f as usize, &args).unwrap();
+                let mut child = ContextHandle::create_and_init(
+                    &mut *stack,
+                    f as usize,
+                    &args,
+                    std::ptr::null_mut(),
+                )
+                .unwrap();
                 unsafe { Context::swap(&mut parent, &mut child) };
                 stack
             },

lucet-module/src/module_data.rs

@@ -59,6 +59,8 @@ pub struct ModuleFeatures {
     pub lzcnt: bool,
     pub popcnt: bool,
     pub instruction_count: bool,
+    pub pinned_heap: bool,
+    pub pinned_heap_register: u16,
     _hidden: (),
 }
 
@@ -75,6 +77,8 @@ impl ModuleFeatures {
             lzcnt: false,
             popcnt: false,
             instruction_count: false,
+            pinned_heap: false,
+            pinned_heap_register: 0,
             _hidden: (),
         }
     }

lucet-runtime/lucet-runtime-internals/src/alloc/tests.rs

@@ -656,6 +656,7 @@ macro_rules! alloc_tests {
                     inst.alloc_mut().stack_u64_mut(),
                     heap_touching_child as usize,
                     &[Val::CPtr(heap_ptr)],
+                    heap_ptr,
                 )
                 .expect("context init succeeds");
                 Context::swap(&mut parent, &mut child);
@@ -705,6 +706,7 @@ macro_rules! alloc_tests {
                     inst.alloc_mut().stack_u64_mut(),
                     stack_pattern_child as usize,
                     &[Val::CPtr(heap_ptr)],
+                    heap_ptr,
                 )
                 .expect("context init succeeds");
                 Context::swap(&mut parent, &mut child);

lucet-runtime/lucet-runtime-internals/src/context/mod.rs

@@ -208,9 +208,10 @@ impl ContextHandle {
         stack: &mut [u64],
         fptr: usize,
         args: &[Val],
+        heap: *mut core::ffi::c_void,
     ) -> Result<ContextHandle, Error> {
         let mut child = ContextHandle::new();
-        Context::init(stack, &mut child, fptr, args)?;
+        Context::init(stack, &mut child, fptr, args, heap)?;
         Ok(child)
     }
 }
@@ -303,6 +304,7 @@ impl Context {
     ///     &mut child,
     ///     entrypoint as usize,
     ///     &[Val::U64(120), Val::F32(3.14)],
+    ///     std::ptr::null_mut(),
     /// );
     /// assert!(res.is_ok());
     /// ```
@@ -326,6 +328,7 @@ impl Context {
     ///     &mut child,
     ///     entrypoint as usize,
     ///     &[Val::U64(120), Val::F32(3.14)],
+    ///     std::ptr::null_mut(),
     /// );
     /// assert!(res.is_ok());
     /// ```
@@ -367,6 +370,7 @@ impl Context {
         child: &mut Context,
         fptr: usize,
         args: &[Val],
+        heap: *mut core::ffi::c_void,
     ) -> Result<(), Error> {
         Context::init_with_callback(
             stack,
@@ -375,6 +379,7 @@ impl Context {
             ptr::null_mut(),
             fptr,
             args,
+            heap,
         )
     }
 
@@ -393,6 +398,7 @@ impl Context {
         callback_data: *mut Instance,
         fptr: usize,
         args: &[Val],
+        heap: *mut core::ffi::c_void,
     ) -> Result<(), Error> {
         if !stack_is_aligned(stack) {
             return Err(Error::UnalignedStack);
@@ -475,6 +481,10 @@ impl Context {
         // even at the entrypoint of the guest.
         child.gpr.rbp = child as *const Context as u64;
 
+        // Heap pinning: r15 is not used to pass any parameters on Windows/POSIX abis, we simply set this to be the value of the heap always.
+        // This value will be used only when the lucet module loaded is compiled requiring use of the pinned heap register.
+        child.gpr.r15 = heap as u64;
+
         Ok(())
     }
 
@@ -547,6 +557,7 @@ impl Context {
     ///     &mut child,
     ///     entrypoint as usize,
     ///     &[],
+    ///     std::ptr::null_mut(),
     /// ).unwrap();
     ///
     /// unsafe { Context::swap(&mut parent, &mut child); }

lucet-runtime/lucet-runtime-internals/src/context/tests/c_child.rs

@@ -55,6 +55,7 @@ macro_rules! init_and_swap {
                 &mut *$stack,
                 $fn as usize,
                 &[$( $args ),*],
+                std::ptr::null_mut(),
             ).unwrap()));
 
             child_regs = child;

lucet-runtime/lucet-runtime-internals/src/context/tests/mod.rs

@@ -31,7 +31,12 @@ fn init_rejects_unaligned() {
     let mut stack_unaligned = unsafe { slice::from_raw_parts_mut(ptr, len) };
 
     // now we have the unaligned stack, let's make sure it blows up right
-    let res = ContextHandle::create_and_init(&mut stack_unaligned, dummy as usize, &[]);
+    let res = ContextHandle::create_and_init(
+        &mut stack_unaligned,
+        dummy as usize,
+        &[],
+        std::ptr::null_mut(),
+    );
 
     if let Err(Error::UnalignedStack) = res {
         assert!(true);

lucet-runtime/lucet-runtime-internals/src/context/tests/rust_child.rs

@@ -51,6 +51,7 @@ macro_rules! init_and_swap {
                 &mut *$stack,
                 $fn as usize,
                 &[$( $args ),*],
+                std::ptr::null_mut(),
             ).unwrap();
             CHILD = Some(child);
 

lucet-runtime/lucet-runtime-internals/src/instance.rs

@@ -973,7 +973,8 @@ impl Instance {
 
         self.entrypoint = Some(func);
 
-        let mut args_with_vmctx = vec![Val::from(self.alloc.slot().heap)];
+        let heap = self.alloc.slot().heap;
+        let mut args_with_vmctx = vec![Val::from(heap)];
         args_with_vmctx.extend_from_slice(args);
 
         let self_ptr = self as *mut _;
@@ -984,6 +985,7 @@ impl Instance {
             self_ptr,
             func.ptr.as_usize(),
             &args_with_vmctx,
+            heap,
         )?;
 
         self.install_activator();

lucetc/lucetc/main.rs

@@ -100,7 +100,8 @@ pub fn run(opts: &Options) -> Result<(), Error> {
         .with_bindings(bindings)
         .with_opt_level(opts.opt_level)
         .with_cpu_features(opts.cpu_features.clone())
-        .with_target(opts.target.clone());
+        .with_target(opts.target.clone())
+        .with_pinned_heap(opts.pinned_heap);
 
     if let Some(validator) = validator.take() {
         c.validator(validator);

lucetc/lucetc/options.rs

@@ -120,6 +120,7 @@ pub struct Options {
     pub pk_path: Option<PathBuf>,
     pub sk_path: Option<PathBuf>,
     pub count_instructions: bool,
+    pub pinned_heap: bool,
     pub error_style: ErrorStyle,
     pub target: Triple,
 }
@@ -211,6 +212,7 @@ impl Options {
         let sk_path = m.value_of("sk_path").map(PathBuf::from);
         let pk_path = m.value_of("pk_path").map(PathBuf::from);
         let count_instructions = m.is_present("count_instructions");
+        let pinned_heap = m.is_present("pinned_heap");
 
         let error_style = match m.value_of("error_style") {
             None => ErrorStyle::default(),
@@ -239,6 +241,7 @@ impl Options {
             sk_path,
             pk_path,
             count_instructions,
+            pinned_heap,
             error_style,
             target,
         })
@@ -452,6 +455,12 @@ SSE3 but not AVX:
                     .takes_value(false)
                     .help("Instrument the produced binary to count the number of wasm operations the translated program executes")
             )
+            .arg(
+                Arg::with_name("pinned_heap")
+                    .long("--pinned-heap-reg")
+                    .takes_value(false)
+                    .help("This feature is not stable - it may be removed in the future! Pin a register to use as this module's heap base. Typically improves performance.")
+            )
             .arg(
                 Arg::with_name("error_style")
                     .long("error-style")

lucetc/src/compiler.rs

@@ -14,6 +14,7 @@ use crate::traps::{translate_trapcode, trap_sym_for_func};
 use byteorder::{LittleEndian, WriteBytesExt};
 use cranelift_codegen::{
     binemit, ir,
+    isa::RegUnit,
     isa::TargetIsa,
     settings::{self, Configurable},
     Context as ClifContext,
@@ -63,6 +64,7 @@ pub struct CompilerBuilder {
     heap_settings: HeapSettings,
     count_instructions: bool,
     canonicalize_nans: bool,
+    pinned_heap: bool,
     validator: Option<Validator>,
 }
 
@@ -75,6 +77,7 @@ impl CompilerBuilder {
             heap_settings: HeapSettings::default(),
             count_instructions: false,
             canonicalize_nans: false,
+            pinned_heap: false,
             validator: None,
         }
     }
@@ -145,6 +148,15 @@ impl CompilerBuilder {
         self
     }
 
+    pub fn pinned_heap(&mut self, pinned_heap: bool) {
+        self.pinned_heap = pinned_heap;
+    }
+
+    pub fn with_pinned_heap(mut self, pinned_heap: bool) -> Self {
+        self.pinned_heap(pinned_heap);
+        self
+    }
+
     pub fn validator(&mut self, validator: Option<Validator>) {
         self.validator = validator;
     }
@@ -169,6 +181,7 @@ impl CompilerBuilder {
             self.count_instructions,
             &self.validator,
             self.canonicalize_nans,
+            self.pinned_heap,
         )
     }
 }
@@ -182,6 +195,8 @@ pub struct Compiler<'a> {
     count_instructions: bool,
     module_translation_state: ModuleTranslationState,
     canonicalize_nans: bool,
+    pinned_heap: bool,
+    pinned_heap_register: RegUnit,
 }
 
 impl<'a> Compiler<'a> {
@@ -195,8 +210,15 @@ impl<'a> Compiler<'a> {
         count_instructions: bool,
         validator: &Option<Validator>,
         canonicalize_nans: bool,
+        pinned_heap: bool,
     ) -> Result<Self, Error> {
-        let isa = Self::target_isa(target.clone(), opt_level, &cpu_features, canonicalize_nans)?;
+        let isa = Self::target_isa(
+            target.clone(),
+            opt_level,
+            &cpu_features,
+            canonicalize_nans,
+            pinned_heap,
+        )?;
 
         let frontend_config = isa.frontend_config();
         let mut module_info = ModuleInfo::new(frontend_config.clone());
@@ -228,6 +250,8 @@ impl<'a> Compiler<'a> {
             _ => (cranelift_module::default_libcall_names())(libcall),
         });
 
+        let pinned_heap_register = isa.register_info().parse_regunit("r15").unwrap();
+
         let mut builder = ObjectBuilder::new(isa, "lucet_guest".to_owned(), libcalls);
         builder.function_alignment(16);
         let mut clif_module: ClifModule<ObjectBackend> = ClifModule::new(builder);
@@ -250,6 +274,8 @@ impl<'a> Compiler<'a> {
             module_translation_state,
             target,
             canonicalize_nans,
+            pinned_heap,
+            pinned_heap_register,
         })
     }
 
@@ -260,6 +286,8 @@ impl<'a> Compiler<'a> {
     pub fn module_features(&self) -> ModuleFeatures {
         let mut mf: ModuleFeatures = (&self.cpu_features).into();
         mf.instruction_count = self.count_instructions;
+        mf.pinned_heap = self.pinned_heap;
+        mf.pinned_heap_register = self.pinned_heap_register;
         mf
     }
 
@@ -475,6 +503,7 @@ impl<'a> Compiler<'a> {
                 self.opt_level,
                 &self.cpu_features,
                 self.canonicalize_nans,
+                self.pinned_heap,
             )?,
         ))
     }
@@ -484,6 +513,7 @@ impl<'a> Compiler<'a> {
         opt_level: OptLevel,
         cpu_features: &CpuFeatures,
         canonicalize_nans: bool,
+        pinned_heap: bool,
     ) -> Result<Box<dyn TargetIsa>, Error> {
         let mut flags_builder = settings::builder();
         let isa_builder = cpu_features.isa_builder(target)?;
@@ -493,6 +523,10 @@ impl<'a> Compiler<'a> {
         if canonicalize_nans {
             flags_builder.enable("enable_nan_canonicalization").unwrap();
         }
+        if pinned_heap {
+            flags_builder.enable("enable_pinned_reg").unwrap();
+            flags_builder.enable("use_pinned_reg_as_heap_base").unwrap();
+        }
         Ok(isa_builder.finish(settings::Flags::new(flags_builder)))
     }
 }