perf: inline common memory-zero load offsets

Author: matthargett

crates/parser/src/emitter.rs

@@ -4,7 +4,10 @@ use crate::{
     ParseError, ParserOptions, Result, conversion::FunctionLoweringContext, selection, visit::FunctionDataBuilder,
 };
 use alloc::vec::Vec;
-use tinywasm_types::{BranchTableOperand, ExceptionCatch, ExceptionHandler, Instruction, Operand128, ValueCounts};
+use tinywasm_types::{
+    BranchTableOperand, ExceptionCatch, ExceptionHandler, Instruction, MemoryOperand, Operand128, Operand128Idx,
+    ValueCounts,
+};
 
 const LOOKBEHIND: usize = 3;
 
@@ -308,6 +311,25 @@ impl Emitter {
                 return Err(ParseError::Other("exception handler range out of bounds".into()));
             }
         }
+        // Run after fusion so only surviving memory-0 loads become inline-offset instructions.
+        if self.select {
+            let inline_offset = |index: Operand128Idx<MemoryOperand>| {
+                let arg = data.operand128(index);
+                if arg.memory() == 0 { u32::try_from(arg.offset()).ok() } else { None }
+            };
+            for instruction in &mut self.instructions {
+                use Instruction::*;
+                let replacement = match *instruction {
+                    I32Load(index) => inline_offset(index).map(I32LoadInline),
+                    I32Load8U(index) => inline_offset(index).map(I32Load8UInline),
+                    I32Load16S(index) => inline_offset(index).map(I32Load16SInline),
+                    _ => None,
+                };
+                if let Some(replacement) = replacement {
+                    *instruction = replacement;
+                }
+            }
+        }
         Ok(self.instructions)
     }
 

crates/tinywasm/src/interpreter/executor/instructions.rs

@@ -357,12 +357,15 @@ macro_rules! instruction_handlers {
             I64Store16(idx) => executor.exec_mem_store::<i64, i16, 2>(idx.resolve(&executor.func.data), #[inline(always)] |v| v as i16)?,
             I64Store32(idx) => executor.exec_mem_store::<i64, i32, 4>(idx.resolve(&executor.func.data), #[inline(always)] |v| v as i32)?,
             I32Load(idx) => executor.exec_mem_load::<i32, 4, _>(idx.resolve(&executor.func.data), identity)?,
+            I32LoadInline(offset) => executor.exec_mem_load::<i32, 4, _>(Operand128::<MemoryOperand>::new(u64::from(*offset), 0), identity)?,
             I64Load(idx) => executor.exec_mem_load::<i64, 8, _>(idx.resolve(&executor.func.data), identity)?,
             F32Load(idx) => executor.exec_mem_load::<f32, 4, _>(idx.resolve(&executor.func.data), identity)?,
             F64Load(idx) => executor.exec_mem_load::<f64, 8, _>(idx.resolve(&executor.func.data), identity)?,
             I32Load8S(idx) => executor.exec_mem_load::<i8, 1, _>(idx.resolve(&executor.func.data), i32::from)?,
             I32Load8U(idx) => executor.exec_mem_load::<u8, 1, _>(idx.resolve(&executor.func.data), i32::from)?,
+            I32Load8UInline(offset) => executor.exec_mem_load::<u8, 1, _>(Operand128::<MemoryOperand>::new(u64::from(*offset), 0), i32::from)?,
             I32Load16S(idx) => executor.exec_mem_load::<i16, 2, _>(idx.resolve(&executor.func.data), i32::from)?,
+            I32Load16SInline(offset) => executor.exec_mem_load::<i16, 2, _>(Operand128::<MemoryOperand>::new(u64::from(*offset), 0), i32::from)?,
             I32Load16U(idx) => executor.exec_mem_load::<u16, 2, _>(idx.resolve(&executor.func.data), i32::from)?,
             I64Load8S(idx) => executor.exec_mem_load::<i8, 1, _>(idx.resolve(&executor.func.data), i64::from)?,
             I64Load8U(idx) => executor.exec_mem_load::<u8, 1, _>(idx.resolve(&executor.func.data), i64::from)?,

crates/tinywasm/tests/memory.rs

@@ -200,3 +200,86 @@ fn resource_limiter_allows_guest_memory_grow_by_default() -> TestResult {
     assert_eq!(grow.call(&mut store, ())?, 1);
     Ok(())
 }
+
+#[test]
+fn inline_memory_immediates_preserve_multi_memory_behavior() -> TestResult {
+    use tinywasm::parser::{Parser, ParserOptions};
+    use tinywasm::types::Instruction;
+
+    let wasm = wat::parse_str(
+        r#"
+        (module
+          (memory 1)
+          (memory $other 1)
+          (data (memory $other) (i32.const 1) "\7f")
+          (func (export "inline") (result i32)
+            i32.const 0
+            i32.const 0x12345678
+            i32.store offset=4
+            i32.const 0
+            i32.load offset=4
+            drop
+            i32.const 0
+            i32.const -2
+            i32.store16 offset=8
+            i32.const 0
+            i32.load16_s offset=8
+            drop
+            i32.const 0
+            i32.load8_u offset=4)
+          (func (export "other") (result i32)
+            i32.const 0
+            i32.load8_u $other offset=1))
+        "#,
+    )?;
+    let optimized = Parser::default().parse_module_bytes(&wasm)?;
+    let instructions = &optimized.funcs[0].instructions;
+    assert!(instructions.iter().any(|op| matches!(op, Instruction::I32LoadInline(4))));
+    assert!(instructions.iter().any(|op| matches!(op, Instruction::I32Load16SInline(8))));
+    assert!(instructions.iter().any(|op| matches!(op, Instruction::I32Load8UInline(4))));
+    assert!(optimized.funcs[1].instructions.iter().any(|op| matches!(op, Instruction::I32Load8U(_))));
+
+    for optimize in [true, false] {
+        let module = Parser::new(ParserOptions::new().with_optimize(optimize)).parse_module_bytes(&wasm)?;
+        let mut store = Store::default();
+        let instance = ModuleInstance::instantiate(&mut store, &module, None)?;
+        assert_eq!(instance.func::<(), i32>(&store, "inline")?.call(&mut store, ())?, 0x78);
+        assert_eq!(instance.func::<(), i32>(&store, "other")?.call(&mut store, ())?, 0x7f);
+    }
+    Ok(())
+}
+
+#[test]
+fn inline_memory_loads_preserve_memory64_and_wide_offsets() -> TestResult {
+    use tinywasm::parser::{Parser, ParserOptions};
+    use tinywasm::types::Instruction;
+
+    let wasm = wat::parse_str(
+        r#"
+        (module
+          (memory i64 1)
+          (data (i64.const 0) "\2a")
+          (func (export "small") (result i32)
+            i64.const 0
+            i32.load8_u)
+          (func (export "wide") (result i32)
+            i64.const 0
+            i32.load8_u offset=4294967296))
+        "#,
+    )?;
+    let optimized = Parser::default().parse_module_bytes(&wasm)?;
+    assert!(optimized.funcs[0].instructions.iter().any(|op| matches!(op, Instruction::I32Load8UInline(0))));
+    assert!(optimized.funcs[1].instructions.iter().any(|op| matches!(op, Instruction::I32Load8U(_))));
+
+    for optimize in [true, false] {
+        let module = Parser::new(ParserOptions::new().with_optimize(optimize)).parse_module_bytes(&wasm)?;
+        let mut store = Store::default();
+        let instance = ModuleInstance::instantiate(&mut store, &module, None)?;
+        assert_eq!(instance.func::<(), i32>(&store, "small")?.call(&mut store, ())?, 42);
+        assert!(matches!(
+            instance.func::<(), i32>(&store, "wide")?.call(&mut store, ()),
+            Err(tinywasm::Error::Trap(Trap::MemoryOutOfBounds { .. }))
+        ));
+    }
+    Ok(())
+}

crates/types/src/archive.rs

@@ -7,7 +7,7 @@ use crate::Module;
 #[rustfmt::skip]
 const TWASM_MAGIC: [u8; 16] = [ TWASM_MAGIC_PREFIX[0], TWASM_MAGIC_PREFIX[1], TWASM_MAGIC_PREFIX[2], TWASM_MAGIC_PREFIX[3], TWASM_VERSION[0], TWASM_VERSION[1], 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
 const TWASM_MAGIC_PREFIX: &[u8; 4] = b"TWAS";
-const TWASM_VERSION: &[u8; 2] = b"06";
+const TWASM_VERSION: &[u8; 2] = b"07";
 
 fn validate_magic(wasm: &[u8]) -> Result<usize, TwasmError> {
     if wasm.len() < TWASM_MAGIC.len() || &wasm[..TWASM_MAGIC_PREFIX.len()] != TWASM_MAGIC_PREFIX {
@@ -98,7 +98,7 @@ mod tests {
         let module = Module::from(ModuleInner { funcs: Box::new([Shared::new(function)]), ..ModuleInner::default() });
 
         let archive = module.serialize_twasm().expect("serialize archive");
-        assert_eq!(&archive[..6], b"TWAS06");
+        assert_eq!(&archive[..6], b"TWAS07");
         let decoded = Module::try_from_twasm(&archive).expect("deserialize archive");
         let function = &decoded.funcs[0];
         assert!(function.max_stack == ValueCounts { c32: 2, c64: 3, c128: 4 });

crates/types/src/instructions.rs

@@ -903,6 +903,9 @@ define_instructions! {
     I32x4RelaxedDotI8x16I7x16AddS,
 
     SelectStore32(Operand128Idx<MemoryOperand>), SelectStore64(Operand128Idx<MemoryOperand>),
+
+    // Selected memory-0 operations carry a 32-bit static offset in the instruction.
+    I32LoadInline(u32), I32Load8UInline(u32), I32Load16SInline(u32),
 }
 
 const _: () = assert!(core::mem::size_of::<Instruction>() == 8);