Allow use of pinned heap registers

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,7 @@ pub struct ModuleFeatures {
     pub lzcnt: bool,
     pub popcnt: bool,
     pub instruction_count: bool,
+    pub pinned_heap: bool,
     _hidden: (),
 }
 
@@ -75,6 +76,7 @@ impl ModuleFeatures {
             lzcnt: false,
             popcnt: false,
             instruction_count: false,
+            pinned_heap: false,
             _hidden: (),
         }
     }

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

@@ -655,6 +655,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);
@@ -704,6 +705,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/context_asm.S

@@ -118,6 +118,7 @@ _lucet_context_swap:
     mov %r14, (6*8)(%rdi)
     mov %r15, (7*8)(%rdi)
     mov %rsi, (8*8)(%rdi)
+    mov %r8,  (9*8)(%rdi)
 
     movdqu %xmm0, (10*8 + 0*16)(%rdi)
     movdqu %xmm1, (10*8 + 1*16)(%rdi)
@@ -137,6 +138,7 @@ _lucet_context_swap:
     mov (5*8)(%rsi), %r13
     mov (6*8)(%rsi), %r14
     mov (7*8)(%rsi), %r15
+    mov (9*8)(%rsi), %r8
 
     movdqu (10*8 + 0*16)(%rsi), %xmm0
     movdqu (10*8 + 1*16)(%rsi), %xmm1

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

@@ -34,6 +34,7 @@ pub(crate) struct GpRegs {
     r14: u64,
     r15: u64,
     pub(crate) rsi: u64,
+    r8: u64,
 }
 
 impl GpRegs {
@@ -48,6 +49,7 @@ impl GpRegs {
             r14: 0,
             r15: 0,
             rsi: 0,
+            r8: 0,
         }
     }
 }
@@ -208,9 +210,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 +306,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 +330,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 +372,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 +381,7 @@ impl Context {
             ptr::null_mut(),
             fptr,
             args,
+            heap,
         )
     }
 
@@ -393,6 +400,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 +483,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 +559,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

@@ -972,7 +972,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 _;
@@ -983,6 +984,7 @@ impl Instance {
             self_ptr,
             func.ptr.as_usize(),
             &args_with_vmctx,
+            heap,
         )?;
 
         self.install_activator();

lucet-runtime/lucet-runtime-tests/src/stack.rs

@@ -162,7 +162,9 @@ macro_rules! stack_tests {
             expect_stack_overflow(
                 // Same note as `expect_ok_locals64_481`
                 stack_testcase(64 - 4).expect("generate stack_testcase 64"),
-                481,
+                // TODO: pick some high enough number that this test overflows still
+                // cranelift changes made 480 the wrong magic number
+                591,
                 true,
             );
         }

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")