From a444cf9097ee629c475093816111f463d52a7cf6 Mon Sep 17 00:00:00 2001 From: Maciej Kot Date: Mon, 22 Apr 2024 06:44:03 +0000 Subject: [PATCH] Feat: RUN-954: Support for wasm64 --- rs/config/src/embedders.rs | 3 + .../src/wasm_utils/instrumentation.rs | 45 +++++++++++++-- rs/embedders/src/wasm_utils/validation.rs | 21 ++++--- rs/embedders/src/wasmtime_embedder.rs | 15 ++++- .../src/wasmtime_embedder/system_api.rs | 48 +++++++++++----- .../wasmtime_embedder_tests.rs | 1 + rs/embedders/tests/spec_tests.rs | 2 +- rs/embedders/tests/wasmtime_embedder.rs | 56 +++++++++++++++++++ .../tests/execution_test.rs | 48 ++++++++++++++++ 9 files changed, 209 insertions(+), 30 deletions(-) diff --git a/rs/config/src/embedders.rs b/rs/config/src/embedders.rs index 923ca641dfc..7d0b39c9656 100644 --- a/rs/config/src/embedders.rs +++ b/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, } } } diff --git a/rs/embedders/src/wasm_utils/instrumentation.rs b/rs/embedders/src/wasm_utils/instrumentation.rs index bd77d8b17e5..02feb020fcb 100644 --- a/rs/embedders/src/wasm_utils/instrumentation.rs +++ b/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 { + 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 = 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 = 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 { diff --git a/rs/embedders/src/wasm_utils/validation.rs b/rs/embedders/src/wasm_utils/validation.rs index 1613ecb431f..778f6b3315a 100644 --- a/rs/embedders/src/wasm_utils/validation.rs +++ b/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)?; diff --git a/rs/embedders/src/wasmtime_embedder.rs b/rs/embedders/src/wasmtime_embedder.rs index 74f839256a4..74168b40822 100644 --- a/rs/embedders/src/wasmtime_embedder.rs +++ b/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> { let mut linker: wasmtime::Linker = 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| { diff --git a/rs/embedders/src/wasmtime_embedder/system_api.rs b/rs/embedders/src/wasmtime_embedder/system_api.rs index 93e3c050ca6..f42aeffd04e 100644 --- a/rs/embedders/src/wasmtime_embedder/system_api.rs +++ b/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( 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", { diff --git a/rs/embedders/src/wasmtime_embedder/wasmtime_embedder_tests.rs b/rs/embedders/src/wasmtime_embedder/wasmtime_embedder_tests.rs index 9d34e8d34c4..6403486596a 100644 --- a/rs/embedders/src/wasmtime_embedder/wasmtime_embedder_tests.rs +++ b/rs/embedders/src/wasmtime_embedder/wasmtime_embedder_tests.rs @@ -130,6 +130,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) diff --git a/rs/embedders/tests/spec_tests.rs b/rs/embedders/tests/spec_tests.rs index 08fd08ffd87..02e12495421 100644 --- a/rs/embedders/tests/spec_tests.rs +++ b/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. diff --git a/rs/embedders/tests/wasmtime_embedder.rs b/rs/embedders/tests/wasmtime_embedder.rs index 3755984fa86..0a2c15aac93 100644 --- a/rs/embedders/tests/wasmtime_embedder.rs +++ b/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)); +} diff --git a/rs/execution_environment/tests/execution_test.rs b/rs/execution_environment/tests/execution_test.rs index c8cfe05870e..01fe959d643 100644 --- a/rs/execution_environment/tests/execution_test.rs +++ b/rs/execution_environment/tests/execution_test.rs @@ -1104,6 +1104,54 @@ fn canister_with_memory_allocation_cannot_grow_wasm_memory_above_allocation() { assert_eq!(err.code(), ErrorCode::CanisterOutOfMemory); } +#[test] +fn canister_with_memory_allocation_cannot_grow_wasm_memory_above_allocation_wasm64() { + let subnet_config = SubnetConfig::new(SubnetType::Application); + let mut embedders_config = ic_config::embedders::Config::default(); + embedders_config.feature_flags.wasm64 = ic_config::flag_status::FlagStatus::Enabled; + let env = StateMachine::new_with_config(StateMachineConfig::new( + subnet_config, + HypervisorConfig { + subnet_memory_capacity: NumBytes::from(100_000_000), + subnet_memory_reservation: NumBytes::from(0), + embedders_config, + ..Default::default() + }, + )); + + let wat = r#" + (module + (import "ic0" "msg_reply" (func $msg_reply)) + (import "ic0" "msg_reply_data_append" + (func $msg_reply_data_append (param i32 i32))) + (func $update + (if (i64.ne (memory.grow (i64.const 400)) (i64.const 1)) + (then (unreachable)) + ) + (call $msg_reply) + ) + (memory $memory i64 1) + (export "canister_update update" (func $update)) + )"#; + + let wasm = wat::parse_str(wat).unwrap(); + + let a_id = create_canister_with_cycles( + &env, + wasm.clone(), + Some( + CanisterSettingsArgsBuilder::new() + .with_memory_allocation(300 * 64 * 1024) + .with_freezing_threshold(0) + .build(), + ), + INITIAL_CYCLES_BALANCE, + ); + + let err = env.execute_ingress(a_id, "update", vec![]).unwrap_err(); + assert_eq!(err.code(), ErrorCode::CanisterOutOfMemory); +} + #[test] fn canister_with_memory_allocation_cannot_grow_stable_memory_above_allocation() { let subnet_config = SubnetConfig::new(SubnetType::Application);