Optimize table.init instruction and instantiation (#2847)

* Optimize `table.init` instruction and instantiation

This commit optimizes table initialization as part of instance
instantiation and also applies the same optimization to the `table.init`
instruction. One part of this commit is to remove some preexisting
duplication between instance instantiation and the `table.init`
instruction itself, after this the actual implementation of `table.init`
is optimized to effectively have fewer bounds checks in fewer places and
have a much tighter loop for instantiation.

A big fallout from this change is that memory/table initializer offsets
are now stored as `u32` instead of `usize` to remove a few casts in a
few places. This ended up requiring moving some overflow checks that
happened in parsing to later in code itself because otherwise the wrong
spec test errors are emitted during testing. I've tried to trace where
these can possibly overflow but I think that I managed to get
everything.

In a local synthetic test where an empty module with a single 80,000
element initializer this improves total instantiation time by 4x (562us
=> 141us)

* Review comments

Author: alexcrichton

cranelift/wasm/src/environ/dummy.rs

@@ -747,7 +747,7 @@ impl<'data> ModuleEnvironment<'data> for DummyEnvironment {
         &mut self,
         _table_index: TableIndex,
         _base: Option<GlobalIndex>,
-        _offset: usize,
+        _offset: u32,
         _elements: Box<[FuncIndex]>,
     ) -> WasmResult<()> {
         // We do nothing
@@ -792,7 +792,7 @@ impl<'data> ModuleEnvironment<'data> for DummyEnvironment {
         &mut self,
         _memory_index: MemoryIndex,
         _base: Option<GlobalIndex>,
-        _offset: usize,
+        _offset: u32,
         _data: &'data [u8],
     ) -> WasmResult<()> {
         // We do nothing

cranelift/wasm/src/environ/spec.rs

@@ -937,7 +937,7 @@ pub trait ModuleEnvironment<'data>: TargetEnvironment {
         &mut self,
         table_index: TableIndex,
         base: Option<GlobalIndex>,
-        offset: usize,
+        offset: u32,
         elements: Box<[FuncIndex]>,
     ) -> WasmResult<()>;
 
@@ -984,7 +984,7 @@ pub trait ModuleEnvironment<'data>: TargetEnvironment {
         &mut self,
         memory_index: MemoryIndex,
         base: Option<GlobalIndex>,
-        offset: usize,
+        offset: u32,
         data: &'data [u8],
     ) -> WasmResult<()>;
 

cranelift/wasm/src/sections_translator.rs

@@ -377,7 +377,7 @@ pub fn parse_element_section<'data>(
             } => {
                 let mut init_expr_reader = init_expr.get_binary_reader();
                 let (base, offset) = match init_expr_reader.read_operator()? {
-                    Operator::I32Const { value } => (None, value as u32 as usize),
+                    Operator::I32Const { value } => (None, value as u32),
                     Operator::GlobalGet { global_index } => {
                         (Some(GlobalIndex::from_u32(global_index)), 0)
                     }
@@ -388,12 +388,6 @@ pub fn parse_element_section<'data>(
                         ));
                     }
                 };
-                // Check for offset + len overflow
-                if offset.checked_add(segments.len()).is_none() {
-                    return Err(wasm_unsupported!(
-                        "element segment offset and length overflows"
-                    ));
-                }
                 environ.declare_table_elements(
                     TableIndex::from_u32(table_index),
                     base,
@@ -429,7 +423,7 @@ pub fn parse_data_section<'data>(
             } => {
                 let mut init_expr_reader = init_expr.get_binary_reader();
                 let (base, offset) = match init_expr_reader.read_operator()? {
-                    Operator::I32Const { value } => (None, value as u32 as usize),
+                    Operator::I32Const { value } => (None, value as u32),
                     Operator::GlobalGet { global_index } => {
                         (Some(GlobalIndex::from_u32(global_index)), 0)
                     }
@@ -440,12 +434,6 @@ pub fn parse_data_section<'data>(
                         ))
                     }
                 };
-                // Check for offset + len overflow
-                if offset.checked_add(data.len()).is_none() {
-                    return Err(wasm_unsupported!(
-                        "data segment offset and length overflows"
-                    ));
-                }
                 environ.declare_data_initialization(
                     MemoryIndex::from_u32(memory_index),
                     base,

crates/environ/src/module.rs

@@ -7,6 +7,7 @@ use cranelift_wasm::*;
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
+use std::convert::TryFrom;
 use std::sync::Arc;
 
 /// Implemenation styles for WebAssembly linear memory.
@@ -86,7 +87,7 @@ pub struct MemoryInitializer {
     /// Optionally, a global variable giving a base index.
     pub base: Option<GlobalIndex>,
     /// The offset to add to the base.
-    pub offset: usize,
+    pub offset: u32,
     /// The data to write into the linear memory.
     pub data: Box<[u8]>,
 }
@@ -168,7 +169,15 @@ impl MemoryInitialization {
                             // Perform a bounds check on the segment
                             // As this segment is referencing a defined memory without a global base, the last byte
                             // written to by the segment cannot exceed the memory's initial minimum size
-                            if (initializer.offset + initializer.data.len())
+                            let offset = usize::try_from(initializer.offset).unwrap();
+                            let end = match offset.checked_add(initializer.data.len()) {
+                                Some(end) => end,
+                                None => {
+                                    out_of_bounds = true;
+                                    continue;
+                                }
+                            };
+                            if end
                                 > ((module.memory_plans[initializer.memory_index].memory.minimum
                                     as usize)
                                     * WASM_PAGE_SIZE)
@@ -178,8 +187,8 @@ impl MemoryInitialization {
                             }
 
                             let pages = &mut map[index];
-                            let mut page_index = initializer.offset / WASM_PAGE_SIZE;
-                            let mut page_offset = initializer.offset % WASM_PAGE_SIZE;
+                            let mut page_index = offset / WASM_PAGE_SIZE;
+                            let mut page_offset = offset % WASM_PAGE_SIZE;
                             let mut data_offset = 0;
                             let mut data_remaining = initializer.data.len();
 
@@ -268,7 +277,7 @@ pub struct TableInitializer {
     /// Optionally, a global variable giving a base index.
     pub base: Option<GlobalIndex>,
     /// The offset to add to the base.
-    pub offset: usize,
+    pub offset: u32,
     /// The values to write into the table elements.
     pub elements: Box<[FuncIndex]>,
 }

crates/environ/src/module_environ.rs

@@ -705,7 +705,7 @@ impl<'data> cranelift_wasm::ModuleEnvironment<'data> for ModuleEnvironment<'data
         &mut self,
         table_index: TableIndex,
         base: Option<GlobalIndex>,
-        offset: usize,
+        offset: u32,
         elements: Box<[FuncIndex]>,
     ) -> WasmResult<()> {
         for element in elements.iter() {
@@ -794,7 +794,7 @@ impl<'data> cranelift_wasm::ModuleEnvironment<'data> for ModuleEnvironment<'data
         &mut self,
         memory_index: MemoryIndex,
         base: Option<GlobalIndex>,
-        offset: usize,
+        offset: u32,
         data: &'data [u8],
     ) -> WasmResult<()> {
         match &mut self.result.module.memory_initialization {

crates/runtime/src/instance.rs

@@ -573,11 +573,11 @@ impl Instance {
             return None;
         }
 
-        Some(unsafe { &*self.anyfunc_ptr(index) })
+        unsafe { Some(&*self.vmctx_plus_offset(self.offsets.vmctx_anyfunc(index))) }
     }
 
-    unsafe fn anyfunc_ptr(&self, index: FuncIndex) -> *mut VMCallerCheckedAnyfunc {
-        self.vmctx_plus_offset(self.offsets.vmctx_anyfunc(index))
+    unsafe fn anyfunc_base(&self) -> *mut VMCallerCheckedAnyfunc {
+        self.vmctx_plus_offset(self.offsets.vmctx_anyfuncs_begin())
     }
 
     fn find_passive_segment<'a, I, D, T>(
@@ -611,38 +611,56 @@ impl Instance {
         src: u32,
         len: u32,
     ) -> Result<(), Trap> {
-        // https://webassembly.github.io/bulk-memory-operations/core/exec/instructions.html#exec-table-init
-
-        let table = self.get_table(table_index);
-
         let elements = Self::find_passive_segment(
             elem_index,
             &self.module.passive_elements_map,
             &self.module.passive_elements,
             &self.dropped_elements,
         );
+        self.table_init_segment(table_index, elements, dst, src, len)
+    }
 
-        if src
-            .checked_add(len)
-            .map_or(true, |n| n as usize > elements.len())
-            || dst.checked_add(len).map_or(true, |m| m > table.size())
-        {
-            return Err(Trap::wasm(ir::TrapCode::TableOutOfBounds));
-        }
+    pub(crate) fn table_init_segment(
+        &self,
+        table_index: TableIndex,
+        elements: &[FuncIndex],
+        dst: u32,
+        src: u32,
+        len: u32,
+    ) -> Result<(), Trap> {
+        // https://webassembly.github.io/bulk-memory-operations/core/exec/instructions.html#exec-table-init
 
-        // TODO(#983): investigate replacing this get/set loop with a `memcpy`.
-        for (dst, src) in (dst..dst + len).zip(src..src + len) {
-            let elem = self
-                .get_caller_checked_anyfunc(elements[src as usize])
-                .map_or(ptr::null_mut(), |f: &VMCallerCheckedAnyfunc| {
-                    f as *const VMCallerCheckedAnyfunc as *mut _
-                });
-
-            table
-                .set(dst, TableElement::FuncRef(elem))
-                .expect("should never panic because we already did the bounds check above");
-        }
+        let table = self.get_table(table_index);
 
+        let elements = match elements
+            .get(usize::try_from(src).unwrap()..)
+            .and_then(|s| s.get(..usize::try_from(len).unwrap()))
+        {
+            Some(elements) => elements,
+            None => return Err(Trap::wasm(ir::TrapCode::TableOutOfBounds)),
+        };
+
+        match table.element_type() {
+            TableElementType::Func => unsafe {
+                let base = self.anyfunc_base();
+                table.init_funcs(
+                    dst,
+                    elements.iter().map(|idx| {
+                        if *idx == FuncIndex::reserved_value() {
+                            ptr::null_mut()
+                        } else {
+                            debug_assert!(idx.as_u32() < self.offsets.num_defined_functions);
+                            base.add(usize::try_from(idx.as_u32()).unwrap())
+                        }
+                    }),
+                )?;
+            },
+
+            TableElementType::Val(_) => {
+                debug_assert!(elements.iter().all(|e| *e == FuncIndex::reserved_value()));
+                table.fill(dst, TableElement::ExternRef(None), len)?;
+            }
+        }
         Ok(())
     }
 
@@ -773,16 +791,26 @@ impl Instance {
         src: u32,
         len: u32,
     ) -> Result<(), Trap> {
-        // https://webassembly.github.io/bulk-memory-operations/core/exec/instructions.html#exec-memory-init
-
-        let memory = self.get_memory(memory_index);
-
         let data = Self::find_passive_segment(
             data_index,
             &self.module.passive_data_map,
             &self.module.passive_data,
             &self.dropped_data,
         );
+        self.memory_init_segment(memory_index, &data, dst, src, len)
+    }
+
+    pub(crate) fn memory_init_segment(
+        &self,
+        memory_index: MemoryIndex,
+        data: &[u8],
+        dst: u32,
+        src: u32,
+        len: u32,
+    ) -> Result<(), Trap> {
+        // https://webassembly.github.io/bulk-memory-operations/core/exec/instructions.html#exec-memory-init
+
+        let memory = self.get_memory(memory_index);
 
         if src
             .checked_add(len)