Wasmtime: Implement the custom-page-sizes proposal

cranelift/codegen/src/isa/aarch64/mod.rs

@@ -179,6 +179,23 @@ impl TargetIsa for AArch64Backend {
         inst::Inst::function_alignment()
     }
 
+    fn page_size_align_log2(&self) -> u8 {
+        use target_lexicon::*;
+        match self.triple().operating_system {
+            OperatingSystem::MacOSX { .. }
+            | OperatingSystem::Darwin
+            | OperatingSystem::Ios
+            | OperatingSystem::Tvos => {
+                debug_assert_eq!(1 << 14, 0x4000);
+                14
+            }
+            _ => {
+                debug_assert_eq!(1 << 16, 0x10000);
+                16
+            }
+        }
+    }
+
     #[cfg(feature = "disas")]
     fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
         use capstone::prelude::*;

cranelift/codegen/src/isa/mod.rs

@@ -238,6 +238,12 @@ pub struct TargetFrontendConfig {
 
     /// The pointer width of the target.
     pub pointer_width: PointerWidth,
+
+    /// The log2 of the target's page size and alignment.
+    ///
+    /// Note that this may be an upper-bound that is larger than necessary for
+    /// some platforms since it may depend on runtime configuration.
+    pub page_size_align_log2: u8,
 }
 
 impl TargetFrontendConfig {
@@ -333,6 +339,12 @@ pub trait TargetIsa: fmt::Display + Send + Sync {
     /// alignment, for performance, required by this ISA.
     fn function_alignment(&self) -> FunctionAlignment;
 
+    /// The log2 of the target's page size and alignment.
+    ///
+    /// Note that this may be an upper-bound that is larger than necessary for
+    /// some platforms since it may depend on runtime configuration.
+    fn page_size_align_log2(&self) -> u8;
+
     /// Create a polymorphic TargetIsa from this specific implementation.
     fn wrapped(self) -> OwnedTargetIsa
     where
@@ -433,6 +445,7 @@ impl<'a> dyn TargetIsa + 'a {
         TargetFrontendConfig {
             default_call_conv: self.default_call_conv(),
             pointer_width: self.pointer_width(),
+            page_size_align_log2: self.page_size_align_log2(),
         }
     }
 }

cranelift/codegen/src/isa/riscv64/mod.rs

@@ -155,6 +155,11 @@ impl TargetIsa for Riscv64Backend {
         inst::Inst::function_alignment()
     }
 
+    fn page_size_align_log2(&self) -> u8 {
+        debug_assert_eq!(1 << 12, 0x1000);
+        12
+    }
+
     #[cfg(feature = "disas")]
     fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
         use capstone::prelude::*;

cranelift/codegen/src/isa/s390x/mod.rs

@@ -156,6 +156,11 @@ impl TargetIsa for S390xBackend {
         inst::Inst::function_alignment()
     }
 
+    fn page_size_align_log2(&self) -> u8 {
+        debug_assert_eq!(1 << 12, 0x1000);
+        12
+    }
+
     #[cfg(feature = "disas")]
     fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
         use capstone::prelude::*;

cranelift/codegen/src/isa/x64/mod.rs

@@ -142,6 +142,11 @@ impl TargetIsa for X64Backend {
         Inst::function_alignment()
     }
 
+    fn page_size_align_log2(&self) -> u8 {
+        debug_assert_eq!(1 << 12, 0x1000);
+        12
+    }
+
     #[cfg(feature = "disas")]
     fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
         use capstone::prelude::*;

cranelift/frontend/src/frontend.rs

@@ -1501,6 +1501,7 @@ mod tests {
         TargetFrontendConfig {
             default_call_conv: CallConv::SystemV,
             pointer_width: PointerWidth::U64,
+            page_size_align_log2: 12,
         }
     }
 

cranelift/wasm/src/code_translator/bounds_checks.rs

@@ -62,6 +62,9 @@ where
     let spectre_mitigations_enabled = env.heap_access_spectre_mitigation();
     let pcc = env.proof_carrying_code();
 
+    let host_page_size_log2 = env.target_config().page_size_align_log2;
+    let can_use_virtual_memory = heap.page_size_log2 >= host_page_size_log2;
+
     let make_compare = |builder: &mut FunctionBuilder,
                         compare_kind: IntCC,
                         lhs: ir::Value,
@@ -188,7 +191,9 @@ where
         //    offset immediates -- which is a common code pattern when accessing
         //    multiple fields in the same struct that is in linear memory --
         //    will all emit the same `index > bound` check, which we can GVN.
-        HeapStyle::Dynamic { bound_gv } if offset_and_size <= heap.offset_guard_size => {
+        HeapStyle::Dynamic { bound_gv }
+            if can_use_virtual_memory && offset_and_size <= heap.offset_guard_size =>
+        {
             let bound = get_dynamic_heap_bound(builder, env, heap);
             let oob = make_compare(
                 builder,
@@ -313,6 +318,10 @@ where
         //    bound`, since we will end up being out-of-bounds regardless of the
         //    given `index`.
         HeapStyle::Static { bound } if offset_and_size > bound.into() => {
+            assert!(
+                can_use_virtual_memory,
+                "static memories require the ability to use virtual memory"
+            );
             env.before_unconditionally_trapping_memory_access(builder)?;
             builder.ins().trap(ir::TrapCode::HeapOutOfBounds);
             Unreachable
@@ -357,10 +366,15 @@ where
         //    within the guard page region, neither of which require emitting an
         //    explicit bounds check.
         HeapStyle::Static { bound }
-            if heap.index_type == ir::types::I32
+            if can_use_virtual_memory
+                && heap.index_type == ir::types::I32
                 && u64::from(u32::MAX)
                     <= u64::from(bound) + u64::from(heap.offset_guard_size) - offset_and_size =>
         {
+            assert!(
+                can_use_virtual_memory,
+                "static memories require the ability to use virtual memory"
+            );
             Reachable(compute_addr(
                 &mut builder.cursor(),
                 heap,
@@ -386,6 +400,10 @@ where
         //    precise, not rely on the virtual memory subsystem at all, and not
         //    factor in the guard pages here.
         HeapStyle::Static { bound } => {
+            assert!(
+                can_use_virtual_memory,
+                "static memories require the ability to use virtual memory"
+            );
             // NB: this subtraction cannot wrap because we didn't hit the first
             // special case.
             let adjusted_bound = u64::from(bound) - offset_and_size;

cranelift/wasm/src/heap.rs

@@ -90,6 +90,9 @@ pub struct HeapData {
 
     /// The memory type for the pointed-to memory, if using proof-carrying code.
     pub memory_type: Option<MemoryType>,
+
+    /// The log2 of this memory's page size.
+    pub page_size_log2: u8,
 }
 
 /// Style of heap including style-specific information.

cranelift/wasm/src/sections_translator.rs

@@ -10,7 +10,7 @@
 use crate::environ::ModuleEnvironment;
 use crate::wasm_unsupported;
 use crate::{
-    DataIndex, ElemIndex, FuncIndex, GlobalIndex, Memory, MemoryIndex, TableIndex, Tag, TagIndex,
+    DataIndex, ElemIndex, FuncIndex, GlobalIndex, MemoryIndex, TableIndex, Tag, TagIndex,
     TypeIndex, WasmError, WasmResult,
 };
 use cranelift_entity::packed_option::ReservedValue;
@@ -20,20 +20,11 @@ use std::vec::Vec;
 use wasmparser::{
     Data, DataKind, DataSectionReader, Element, ElementItems, ElementKind, ElementSectionReader,
     Export, ExportSectionReader, ExternalKind, FunctionSectionReader, GlobalSectionReader,
-    ImportSectionReader, MemorySectionReader, MemoryType, Operator, TableSectionReader,
-    TagSectionReader, TagType, TypeRef, TypeSectionReader,
+    ImportSectionReader, MemorySectionReader, Operator, TableSectionReader, TagSectionReader,
+    TagType, TypeRef, TypeSectionReader,
 };
 use wasmtime_types::ConstExpr;
 
-fn memory(ty: MemoryType) -> Memory {
-    Memory {
-        minimum: ty.initial,
-        maximum: ty.maximum,
-        shared: ty.shared,
-        memory64: ty.memory64,
-    }
-}
-
 fn tag(e: TagType) -> Tag {
     match e.kind {
         wasmparser::TagKind::Exception => Tag {
@@ -75,7 +66,7 @@ pub fn parse_import_section<'data>(
                 )?;
             }
             TypeRef::Memory(ty) => {
-                environ.declare_memory_import(memory(ty), import.module, import.name)?;
+                environ.declare_memory_import(ty.into(), import.module, import.name)?;
             }
             TypeRef::Tag(e) => {
                 environ.declare_tag_import(tag(e), import.module, import.name)?;
@@ -139,8 +130,7 @@ pub fn parse_memory_section(
     environ.reserve_memories(memories.count())?;
 
     for entry in memories {
-        let memory = memory(entry?);
-        environ.declare_memory(memory)?;
+        environ.declare_memory(entry?.into())?;
     }
 
     Ok(())

crates/cranelift/src/func_environ.rs

@@ -20,7 +20,7 @@ use std::mem;
 use wasmparser::Operator;
 use wasmtime_environ::{
     BuiltinFunctionIndex, MemoryPlan, MemoryStyle, Module, ModuleTranslation, ModuleTypesBuilder,
-    PtrSize, TableStyle, Tunables, TypeConvert, VMOffsets, WASM_PAGE_SIZE,
+    PtrSize, TableStyle, Tunables, TypeConvert, VMOffsets,
 };
 use wasmtime_environ::{FUNCREF_INIT_BIT, FUNCREF_MASK};
 
@@ -680,7 +680,13 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         }
     }
 
-    fn cast_pointer_to_memory_index(
+    /// Convert the target pointer-sized integer `val` that is holding a memory
+    /// length (or the `-1` `memory.grow`-failed sentinel) into the memory's
+    /// index type.
+    ///
+    /// This might involve extending or truncating it depending on the memory's
+    /// index type and the target's pointer type.
+    fn convert_memory_length_to_index_type(
         &self,
         mut pos: FuncCursor<'_>,
         val: ir::Value,
@@ -698,18 +704,32 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         } else if pointer_type.bits() > desired_type.bits() {
             pos.ins().ireduce(desired_type, val)
         } else {
-            // Note that we `sextend` instead of the probably expected
-            // `uextend`. This function is only used within the contexts of
-            // `memory.size` and `memory.grow` where we're working with units of
-            // pages instead of actual bytes, so we know that the upper bit is
-            // always cleared for "valid values". The one case we care about
-            // sextend would be when the return value of `memory.grow` is `-1`,
-            // in which case we want to copy the sign bit.
-            //
-            // This should only come up on 32-bit hosts running wasm64 modules,
-            // which at some point also makes you question various assumptions
-            // made along the way...
-            pos.ins().sextend(desired_type, val)
+            // We have a 64-bit memory on a 32-bit host -- this combo doesn't
+            // really make a whole lot of sense to do from a user perspective
+            // but that is neither here nor there. We want to logically do an
+            // unsigned extend *except* when we are given the `-1` sentinel,
+            // which we must preserve as `-1` in the wider type.
+            match self.module.memory_plans[index].memory.page_size_log2 {
+                16 => {
+                    // In the case that we have default page sizes, we can
+                    // always sign extend, since valid memory lengths (in pages)
+                    // never have their sign bit set, and so if the sign bit is
+                    // set then this must be the `-1` sentinel, which we want to
+                    // preserve through the extension.
+                    pos.ins().sextend(desired_type, val)
+                }
+                0 => {
+                    // For single-byte pages, we have to explicitly check for
+                    // `-1` and choose whether to do an unsigned extension or
+                    // return a larger `-1` because there are valid memory
+                    // lengths (in pages) that have the sign bit set.
+                    let extended = pos.ins().uextend(desired_type, val);
+                    let neg_one = pos.ins().iconst(desired_type, -1);
+                    let is_failure = pos.ins().icmp_imm(IntCC::Equal, val, -1);
+                    pos.ins().select(is_failure, neg_one, extended)
+                }
+                _ => unreachable!("only page sizes 2**0 and 2**16 are currently valid"),
+            }
         }
     }
 
@@ -2001,21 +2021,21 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
 
         let min_size = self.module.memory_plans[index]
             .memory
-            .minimum
-            .checked_mul(u64::from(WASM_PAGE_SIZE))
-            .unwrap_or_else(|| {
+            .minimum_byte_size()
+            .unwrap_or_else(|_| {
                 // The only valid Wasm memory size that won't fit in a 64-bit
                 // integer is the maximum memory64 size (2^64) which is one
                 // larger than `u64::MAX` (2^64 - 1). In this case, just say the
                 // minimum heap size is `u64::MAX`.
                 debug_assert_eq!(self.module.memory_plans[index].memory.minimum, 1 << 48);
+                debug_assert_eq!(self.module.memory_plans[index].memory.page_size(), 1 << 16);
                 u64::MAX
             });
 
         let max_size = self.module.memory_plans[index]
             .memory
-            .maximum
-            .and_then(|max| max.checked_mul(u64::from(WASM_PAGE_SIZE)));
+            .maximum_byte_size()
+            .ok();
 
         let (ptr, base_offset, current_length_offset, ptr_memtype) = {
             let vmctx = self.vmctx(func);
@@ -2069,6 +2089,8 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
             }
         };
 
+        let page_size_log2 = self.module.memory_plans[index].memory.page_size_log2;
+
         // If we have a declared maximum, we can make this a "static" heap, which is
         // allocated up front and never moved.
         let (offset_guard_size, heap_style, readonly_base, base_fact, memory_type) =
@@ -2233,6 +2255,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
             style: heap_style,
             index_type: self.memory_index_type(index),
             memory_type,
+            page_size_log2,
         }))
     }
 
@@ -2397,7 +2420,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         let val = self.cast_memory_index_to_i64(&mut pos, val, index);
         let call_inst = pos.ins().call(memory_grow, &[vmctx, val, memory_index]);
         let result = *pos.func.dfg.inst_results(call_inst).first().unwrap();
-        Ok(self.cast_pointer_to_memory_index(pos, result, index))
+        Ok(self.convert_memory_length_to_index_type(pos, result, index))
     }
 
     fn translate_memory_size(
@@ -2469,11 +2492,11 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                 }
             }
         };
-        let current_length_in_pages = pos
-            .ins()
-            .udiv_imm(current_length_in_bytes, i64::from(WASM_PAGE_SIZE));
 
-        Ok(self.cast_pointer_to_memory_index(pos, current_length_in_pages, index))
+        let page_size_log2 = i64::from(self.module.memory_plans[index].memory.page_size_log2);
+        let current_length_in_pages = pos.ins().ushr_imm(current_length_in_bytes, page_size_log2);
+
+        Ok(self.convert_memory_length_to_index_type(pos, current_length_in_pages, index))
     }
 
     fn translate_memory_copy(