Merge branch 'mk/wasm64_2' into 'master'

Feat: RUN-954: Support for wasm64

This MR adds support for 64bit main wasm memory and puts it behind a feature flag in EmbeddersConfig.
Support for accessing memory above 4gb with system api calls will be added in a separate MR. 

See merge request dfinity-lab/public/ic!18661

rs/config/src/embedders.rs

@@ -81,6 +81,8 @@ pub struct FeatureFlags {
     // TODO(IC-272): remove this flag once the feature is enabled by default.
     /// Indicates whether canister logging feature is enabled or not.
     pub canister_logging: FlagStatus,
+    /// Indicates whether the support for 64 bit main memory is enabled
+    pub wasm64: FlagStatus,
 }
 
 impl FeatureFlags {
@@ -90,6 +92,7 @@ impl FeatureFlags {
             write_barrier: FlagStatus::Disabled,
             wasm_native_stable_memory: FlagStatus::Enabled,
             canister_logging: FlagStatus::Disabled,
+            wasm64: FlagStatus::Disabled,
         }
     }
 }

rs/embedders/src/wasm_utils/instrumentation.rs

@@ -136,6 +136,12 @@ use wasmparser::{
 use std::collections::BTreeMap;
 use std::convert::TryFrom;
 
+#[derive(Clone, Copy, Debug)]
+pub(crate) enum WasmMemoryType {
+    Wasm32,
+    Wasm64,
+}
+
 // The indices of injected function imports.
 pub(crate) enum InjectedImports {
     OutOfInstructions = 0,
@@ -472,6 +478,7 @@ const BYTEMAP_SIZE_IN_WASM_PAGES: u64 =
     MAX_WASM_MEMORY_IN_BYTES / (PAGE_SIZE as u64) / (WASM_PAGE_SIZE as u64);
 
 const MAX_STABLE_MEMORY_IN_WASM_PAGES: u64 = MAX_STABLE_MEMORY_IN_BYTES / (WASM_PAGE_SIZE as u64);
+const MAX_WASM_MEMORY_IN_WASM_PAGES: u64 = MAX_WASM_MEMORY_IN_BYTES / (WASM_PAGE_SIZE as u64);
 /// There is one byte for each OS page in the stable memory.
 const STABLE_BYTEMAP_SIZE_IN_WASM_PAGES: u64 = MAX_STABLE_MEMORY_IN_WASM_PAGES / (PAGE_SIZE as u64);
 
@@ -558,10 +565,17 @@ fn mutate_function_indices(module: &mut Module, f: impl Fn(u32) -> u32) {
 /// added as the last imports, we'd need to increment only non imported
 /// functions, since imported functions precede all others in the function index
 /// space, but this would be error-prone).
-fn inject_helper_functions(mut module: Module, wasm_native_stable_memory: FlagStatus) -> Module {
+fn inject_helper_functions(
+    mut module: Module,
+    wasm_native_stable_memory: FlagStatus,
+    mem_type: WasmMemoryType,
+) -> Module {
     // insert types
     let ooi_type = FuncType::new([], []);
-    let tgwm_type = FuncType::new([ValType::I32, ValType::I32], [ValType::I32]);
+    let tgwm_type = match mem_type {
+        WasmMemoryType::Wasm32 => FuncType::new([ValType::I32, ValType::I32], [ValType::I32]),
+        WasmMemoryType::Wasm64 => FuncType::new([ValType::I64, ValType::I64], [ValType::I64]),
+    };
 
     let ooi_type_idx = add_func_type(&mut module, ooi_type);
     let tgwm_type_idx = add_func_type(&mut module, tgwm_type);
@@ -670,8 +684,14 @@ pub(super) fn instrument(
     subnet_type: SubnetType,
     dirty_page_overhead: NumInstructions,
 ) -> Result<InstrumentationOutput, WasmInstrumentationError> {
+    let mut main_memory_type = WasmMemoryType::Wasm32;
+    if let Some(mem) = module.memories.first() {
+        if mem.memory64 {
+            main_memory_type = WasmMemoryType::Wasm64;
+        }
+    }
     let stable_memory_index;
-    let mut module = inject_helper_functions(module, wasm_native_stable_memory);
+    let mut module = inject_helper_functions(module, wasm_native_stable_memory, main_memory_type);
     module = export_table(module);
     (module, stable_memory_index) =
         update_memories(module, write_barrier, wasm_native_stable_memory);
@@ -753,7 +773,7 @@ pub(super) fn instrument(
     if !func_types.is_empty() {
         let func_bodies = &mut module.code_sections;
         for (func_ix, func_type) in func_types.into_iter() {
-            inject_try_grow_wasm_memory(&mut func_bodies[func_ix], &func_type);
+            inject_try_grow_wasm_memory(&mut func_bodies[func_ix], &func_type, main_memory_type);
             if write_barrier == FlagStatus::Enabled {
                 inject_mem_barrier(&mut func_bodies[func_ix], &func_type);
             }
@@ -1456,7 +1476,11 @@ fn inject_mem_barrier(func_body: &mut ic_wasm_transform::Body, func_type: &FuncT
 // Scans through the function and adds instrumentation after each `memory.grow`
 // instruction to make sure that there's enough available memory left to support
 // the requested extra memory.
-fn inject_try_grow_wasm_memory(func_body: &mut ic_wasm_transform::Body, func_type: &FuncType) {
+fn inject_try_grow_wasm_memory(
+    func_body: &mut ic_wasm_transform::Body,
+    func_type: &FuncType,
+    mem_type: WasmMemoryType,
+) {
     use Operator::*;
     let mut injection_points: Vec<usize> = Vec::new();
     {
@@ -1474,7 +1498,10 @@ fn inject_try_grow_wasm_memory(func_body: &mut ic_wasm_transform::Body, func_typ
         // over the first field gives the total number of locals.
         let n_locals: u32 = func_body.locals.iter().map(|x| x.0).sum();
         let memory_local_ix = func_type.params().len() as u32 + n_locals;
-        func_body.locals.push((1, ValType::I32));
+        match mem_type {
+            WasmMemoryType::Wasm32 => func_body.locals.push((1, ValType::I32)),
+            WasmMemoryType::Wasm64 => func_body.locals.push((1, ValType::I64)),
+        };
 
         let orig_elems = &func_body.instructions;
         let mut elems: Vec<Operator> = Vec::new();
@@ -1705,6 +1732,12 @@ fn update_memories(
 ) -> (Module, u32) {
     let mut stable_index = 0;
 
+    if let Some(mem) = module.memories.first_mut() {
+        if mem.memory64 && mem.maximum.is_none() {
+            mem.maximum = Some(MAX_WASM_MEMORY_IN_WASM_PAGES);
+        }
+    }
+
     let mut memory_already_exported = false;
     for export in &mut module.exports {
         if let ExternalKind::Memory = export.kind {

rs/embedders/src/wasm_utils/validation.rs

@@ -1324,7 +1324,7 @@ fn validate_code_section(
 }
 
 /// Returns a Wasmtime config that is used for Wasm validation.
-pub fn wasmtime_validation_config() -> wasmtime::Config {
+pub fn wasmtime_validation_config(embedders_config: &EmbeddersConfig) -> wasmtime::Config {
     let mut config = wasmtime::Config::default();
 
     // Keep this in the alphabetical order to simplify comparison with new
@@ -1344,10 +1344,14 @@ pub fn wasmtime_validation_config() -> wasmtime::Config {
     config.wasm_bulk_memory(true);
     config.wasm_function_references(false);
     config.wasm_gc(false);
-    // Wasm memory64 and multi-memory features are disabled during validation,
+    if embedders_config.feature_flags.wasm64 == ic_config::flag_status::FlagStatus::Enabled {
+        config.wasm_memory64(true);
+    } else {
+        config.wasm_memory64(false);
+    }
+    // Wasm multi-memory feature is disabled during validation,
     // but enabled during execution for the Wasm-native stable memory
     // implementation.
-    config.wasm_memory64(false);
     config.wasm_multi_memory(false);
     config.wasm_reference_types(true);
     // The SIMD instructions are disable for determinism.
@@ -1376,12 +1380,15 @@ pub fn wasmtime_validation_config() -> wasmtime::Config {
 
 #[test]
 fn can_create_engine_from_validation_config() {
-    let config = wasmtime_validation_config();
+    let config = wasmtime_validation_config(&EmbeddersConfig::default());
     wasmtime::Engine::new(&config).expect("Cannot create engine from validation config");
 }
 
-fn can_compile(wasm: &BinaryEncodedWasm) -> Result<(), WasmValidationError> {
-    let config = wasmtime_validation_config();
+fn can_compile(
+    wasm: &BinaryEncodedWasm,
+    embedders_config: &EmbeddersConfig,
+) -> Result<(), WasmValidationError> {
+    let config = wasmtime_validation_config(embedders_config);
     let engine = wasmtime::Engine::new(&config).expect("Failed to create wasmtime::Engine");
     wasmtime::Module::validate(&engine, wasm.as_slice()).map_err(|err| {
         WasmValidationError::WasmtimeValidation(format!(
@@ -1435,7 +1442,7 @@ pub(super) fn validate_wasm_binary<'a>(
     config: &EmbeddersConfig,
 ) -> Result<(WasmValidationDetails, Module<'a>), WasmValidationError> {
     check_code_section_size(wasm)?;
-    can_compile(wasm)?;
+    can_compile(wasm, config)?;
     let module = Module::parse(wasm.as_slice(), false)
         .map_err(|err| WasmValidationError::DecodingError(format!("{}", err)))?;
     let imports_details = validate_import_section(&module)?;

rs/embedders/src/wasmtime_embedder.rs

@@ -37,7 +37,7 @@ use memory_tracker::{DirtyPageTracking, PageBitmap, SigsegvMemoryTracker};
 use signal_stack::WasmtimeSignalStack;
 
 use crate::wasm_utils::instrumentation::{
-    ACCESSED_PAGES_COUNTER_GLOBAL_NAME, DIRTY_PAGES_COUNTER_GLOBAL_NAME,
+    WasmMemoryType, ACCESSED_PAGES_COUNTER_GLOBAL_NAME, DIRTY_PAGES_COUNTER_GLOBAL_NAME,
     INSTRUCTIONS_COUNTER_GLOBAL_NAME,
 };
 use crate::{
@@ -189,7 +189,7 @@ impl WasmtimeEmbedder {
     /// canisters __except__ the `host_memory`.
     #[doc(hidden)]
     pub fn wasmtime_execution_config(embedder_config: &EmbeddersConfig) -> wasmtime::Config {
-        let mut config = wasmtime_validation_config();
+        let mut config = wasmtime_validation_config(embedder_config);
 
         // Wasmtime features that differ between Wasm validation and execution.
         // Currently these are multi-memories and the 64-bit memory needed for
@@ -229,12 +229,23 @@ impl WasmtimeEmbedder {
 
     pub fn pre_instantiate(&self, module: &Module) -> HypervisorResult<InstancePre<StoreData>> {
         let mut linker: wasmtime::Linker<StoreData> = Linker::new(module.engine());
+        let mut main_memory_type = WasmMemoryType::Wasm32;
+
+        if let Some(export) = module.get_export(WASM_HEAP_MEMORY_NAME) {
+            if let Some(mem) = export.memory() {
+                if mem.is_64() {
+                    main_memory_type = WasmMemoryType::Wasm64;
+                }
+            }
+        }
+
         system_api::syscalls(
             &mut linker,
             self.config.feature_flags,
             self.config.stable_memory_dirty_page_limit,
             self.config.stable_memory_accessed_page_limit,
             self.config.metering_type,
+            main_memory_type,
         );
 
         let instance_pre = linker.instantiate_pre(module).map_err(|e| {

rs/embedders/src/wasmtime_embedder/system_api.rs

@@ -21,6 +21,7 @@ use wasmtime::{AsContextMut, Caller, Global, Linker, Val};
 use crate::InternalErrorCode;
 use std::convert::TryFrom;
 
+use crate::wasm_utils::instrumentation::WasmMemoryType;
 use crate::wasmtime_embedder::system_api_complexity::system_api;
 use ic_system_api::SystemApiImpl;
 
@@ -308,6 +309,7 @@ pub(crate) fn syscalls(
     stable_memory_dirty_page_limit: NumPages,
     stable_memory_access_page_limit: NumPages,
     metering_type: MeteringType,
+    main_memory_type: WasmMemoryType,
 ) {
     fn with_system_api<T>(
         mut caller: &mut Caller<'_, StoreData>,
@@ -1056,21 +1058,39 @@ pub(crate) fn syscalls(
         })
         .unwrap();
 
-    linker
-        .func_wrap("__", "try_grow_wasm_memory", {
-            move |mut caller: Caller<'_, StoreData>,
-                  native_memory_grow_res: i32,
-                  additional_wasm_pages: u32| {
-                with_system_api(&mut caller, |s| {
-                    s.try_grow_wasm_memory(
-                        native_memory_grow_res as i64,
-                        additional_wasm_pages as u64,
-                    )
+    match main_memory_type {
+        WasmMemoryType::Wasm32 => {
+            linker
+                .func_wrap("__", "try_grow_wasm_memory", {
+                    move |mut caller: Caller<'_, StoreData>,
+                          native_memory_grow_res: i32,
+                          additional_wasm_pages: u32| {
+                        with_system_api(&mut caller, |s| {
+                            s.try_grow_wasm_memory(
+                                native_memory_grow_res as i64,
+                                additional_wasm_pages as u64,
+                            )
+                        })
+                        .map(|()| native_memory_grow_res)
+                    }
                 })
-                .map(|()| native_memory_grow_res)
-            }
-        })
-        .unwrap();
+                .unwrap();
+        }
+        WasmMemoryType::Wasm64 => {
+            linker
+                .func_wrap("__", "try_grow_wasm_memory", {
+                    move |mut caller: Caller<'_, StoreData>,
+                          native_memory_grow_res: i64,
+                          additional_wasm_pages: u64| {
+                        with_system_api(&mut caller, |s| {
+                            s.try_grow_wasm_memory(native_memory_grow_res, additional_wasm_pages)
+                        })
+                        .map(|()| native_memory_grow_res)
+                    }
+                })
+                .unwrap();
+        }
+    }
 
     linker
         .func_wrap("__", "try_grow_stable_memory", {

rs/embedders/src/wasmtime_embedder/wasmtime_embedder_tests.rs

@@ -136,6 +136,7 @@ fn test_wasmtime_system_api() {
         config.stable_memory_dirty_page_limit,
         config.stable_memory_accessed_page_limit,
         config.metering_type,
+        crate::wasmtime_embedder::WasmMemoryType::Wasm32,
     );
     let instance = linker
         .instantiate(&mut store, &module)

rs/embedders/tests/spec_tests.rs

@@ -743,7 +743,7 @@ fn run_testsuite(subdirectory: &str, config: &Config, parsing_multi_memory_enabl
 /// Returns the config that is as close as possible to the actual config used in
 /// production for validation.
 fn default_config() -> Config {
-    let mut config = wasmtime_validation_config();
+    let mut config = wasmtime_validation_config(&ic_config::embedders::Config::default());
     // Some tests require SIMD instructions to run.
     config.wasm_simd(true);
     // This is needed to avoid stack overflows in some tests.

rs/embedders/tests/wasmtime_embedder.rs

@@ -1711,3 +1711,59 @@ fn wasm_logging_new_records_after_exceeding_log_size_limit() {
         }
     }
 }
+
+#[test]
+// Verify that we can create 64 bit memory and write to it
+fn wasm64_basic_test() {
+    let wat = r#"
+    (module
+        (global $g1 (export "g1") (mut i64) (i64.const 0))
+        (func $test (export "canister_update test")
+            (i64.store (i64.const 0) (memory.grow (i64.const 1)))
+            (i64.store (i64.const 20) (i64.const 137))
+            (i64.load (i64.const 20))
+            global.set $g1
+        )
+        (memory (export "memory") i64 10)
+    )"#;
+
+    let mut config = ic_config::embedders::Config::default();
+    config.feature_flags.wasm64 = FlagStatus::Enabled;
+    let mut instance = WasmtimeInstanceBuilder::new()
+        .with_config(config)
+        .with_wat(wat)
+        .build();
+    let res = instance
+        .run(FuncRef::Method(WasmMethod::Update("test".to_string())))
+        .unwrap();
+    assert_eq!(res.exported_globals[0], Global::I64(137));
+}
+
+#[test]
+// Verify behavior of failed memory grow in wasm64 mode
+fn wasm64_handles_memory_grow_failure_test() {
+    let wat = r#"
+    (module
+        (global $g1 (export "g1") (mut i64) (i64.const 0))
+        (global $g2 (export "g2") (mut i64) (i64.const 0))
+        (func $test (export "canister_update test")
+            (memory.grow (i64.const 165536))
+            global.set $g1
+            (i64.const 137)
+            global.set $g2
+        )
+        (memory (export "memory") i64 10)
+    )"#;
+
+    let mut config = ic_config::embedders::Config::default();
+    config.feature_flags.wasm64 = FlagStatus::Enabled;
+    let mut instance = WasmtimeInstanceBuilder::new()
+        .with_config(config)
+        .with_wat(wat)
+        .build();
+    let res = instance
+        .run(FuncRef::Method(WasmMethod::Update("test".to_string())))
+        .unwrap();
+    assert_eq!(res.exported_globals[0], Global::I64(-1));
+    assert_eq!(res.exported_globals[1], Global::I64(137));
+}