chore(RUN-1012): infrastructure to increase Wasm64 heap memory size (#…

…3385)

This PR prepares infrastructure to increase Wasm64 heap memory size
allowing Wasm64 canisters to run with different heap memory sizes
alongside the Wasm32 canisters.

Note that this PR is a no-op for now and not changing any functionality,
as the Wasm64 heap memory size is still kept at 4GiB. The intent is that
this only allows that in the future the change can happen seamlessly, by
only changing a constant.

rs/config/src/embedders.rs

@@ -3,7 +3,10 @@ use std::time::Duration;
 use ic_base_types::NumBytes;
 use ic_registry_subnet_type::SubnetType;
 use ic_sys::PAGE_SIZE;
-use ic_types::{NumInstructions, NumOsPages, MAX_STABLE_MEMORY_IN_BYTES, MAX_WASM_MEMORY_IN_BYTES};
+use ic_types::{
+    NumInstructions, NumOsPages, MAX_STABLE_MEMORY_IN_BYTES, MAX_WASM64_MEMORY_IN_BYTES,
+    MAX_WASM_MEMORY_IN_BYTES,
+};
 use serde::{Deserialize, Serialize};
 
 use crate::flag_status::FlagStatus;
@@ -245,6 +248,9 @@ pub struct Config {
     /// The maximum size of the wasm heap memory.
     pub max_wasm_memory_size: NumBytes,
 
+    /// The maximum size of the wasm heap memory for Wasm64 canisters.
+    pub max_wasm64_memory_size: NumBytes,
+
     /// The maximum size of the stable memory.
     pub max_stable_memory_size: NumBytes,
 }
@@ -284,6 +290,7 @@ impl Config {
             dirty_page_copy_overhead: DIRTY_PAGE_COPY_OVERHEAD,
             wasm_max_size: WASM_MAX_SIZE,
             max_wasm_memory_size: NumBytes::new(MAX_WASM_MEMORY_IN_BYTES),
+            max_wasm64_memory_size: NumBytes::new(MAX_WASM64_MEMORY_IN_BYTES),
             max_stable_memory_size: NumBytes::new(MAX_STABLE_MEMORY_IN_BYTES),
             wasm64_dirty_page_overhead_multiplier: WASM64_DIRTY_PAGE_OVERHEAD_MULTIPLIER,
         }

rs/config/src/execution_environment.rs

@@ -2,7 +2,8 @@ use crate::embedders::Config as EmbeddersConfig;
 use crate::flag_status::FlagStatus;
 use ic_base_types::{CanisterId, NumSeconds};
 use ic_types::{
-    Cycles, NumBytes, NumInstructions, MAX_STABLE_MEMORY_IN_BYTES, MAX_WASM_MEMORY_IN_BYTES,
+    Cycles, NumBytes, NumInstructions, MAX_STABLE_MEMORY_IN_BYTES, MAX_WASM64_MEMORY_IN_BYTES,
+    MAX_WASM_MEMORY_IN_BYTES,
 };
 use serde::{Deserialize, Serialize};
 use std::{str::FromStr, time::Duration};
@@ -204,7 +205,12 @@ pub struct Config {
     pub subnet_memory_reservation: NumBytes,
 
     /// The maximum amount of memory that can be utilized by a single canister.
-    pub max_canister_memory_size: NumBytes,
+    /// running in Wasm32 mode.
+    pub max_canister_memory_size_wasm32: NumBytes,
+
+    /// The maximum amount of memory that can be utilized by a single canister.
+    /// running in Wasm64 mode.
+    pub max_canister_memory_size_wasm64: NumBytes,
 
     /// The soft limit on the subnet-wide number of callbacks. Beyond this limit,
     /// canisters are only allowed to make downstream calls up to their individual
@@ -342,9 +348,12 @@ impl Default for Config {
             subnet_wasm_custom_sections_memory_capacity:
                 SUBNET_WASM_CUSTOM_SECTIONS_MEMORY_CAPACITY,
             subnet_memory_reservation: SUBNET_MEMORY_RESERVATION,
-            max_canister_memory_size: NumBytes::new(
+            max_canister_memory_size_wasm32: NumBytes::new(
                 MAX_STABLE_MEMORY_IN_BYTES + MAX_WASM_MEMORY_IN_BYTES,
             ),
+            max_canister_memory_size_wasm64: NumBytes::new(
+                MAX_STABLE_MEMORY_IN_BYTES + MAX_WASM64_MEMORY_IN_BYTES,
+            ),
             subnet_callback_soft_limit: SUBNET_CALLBACK_SOFT_LIMIT,
             canister_guaranteed_callback_quota: CANISTER_GUARANTEED_CALLBACK_QUOTA,
             default_provisional_cycles_balance: Cycles::new(100_000_000_000_000),

rs/drun/src/main.rs

@@ -67,9 +67,9 @@ async fn drun_main() -> Result<(), String> {
             .embedders_config
             .feature_flags
             .best_effort_responses = FlagStatus::Enabled;
-        hypervisor_config.embedders_config.max_wasm_memory_size = MAIN_MEMORY_CAPACITY;
-        hypervisor_config.max_canister_memory_size =
-            hypervisor_config.embedders_config.max_wasm_memory_size
+        hypervisor_config.embedders_config.max_wasm64_memory_size = MAIN_MEMORY_CAPACITY;
+        hypervisor_config.max_canister_memory_size_wasm64 =
+            hypervisor_config.embedders_config.max_wasm64_memory_size
                 + hypervisor_config.embedders_config.max_stable_memory_size;
 
         let cfg = Config::load_with_default(&source, default_config).unwrap_or_else(|err| {

rs/embedders/benches/embedders_bench/src/lib.rs

@@ -57,7 +57,7 @@ fn initialize_execution_test(
     if is_wasm64 {
         test = test.with_wasm64();
         // Set memory size to 8 GiB for Wasm64.
-        test = test.with_max_wasm_memory_size(NumBytes::from(8 * 1024 * 1024 * 1024));
+        test = test.with_max_wasm64_memory_size(NumBytes::from(8 * 1024 * 1024 * 1024));
     }
     let mut test = test.build();
 

rs/embedders/src/wasm_utils.rs

@@ -14,7 +14,9 @@ use ic_types::{methods::WasmMethod, NumInstructions};
 use ic_wasm_types::{BinaryEncodedWasm, WasmInstrumentationError};
 use serde::{Deserialize, Serialize};
 
-use self::{instrumentation::instrument, validation::validate_wasm_binary};
+use self::{
+    instrumentation::instrument, validation::has_wasm64_memory, validation::validate_wasm_binary,
+};
 use crate::wasmtime_embedder::StoreData;
 use crate::{serialized_module::SerializedModule, CompilationResult, WasmtimeEmbedder};
 use wasmtime::InstancePre;
@@ -204,6 +206,13 @@ fn validate_and_instrument(
     config: &EmbeddersConfig,
 ) -> HypervisorResult<(WasmValidationDetails, InstrumentationOutput)> {
     let (wasm_validation_details, module) = validate_wasm_binary(wasm, config)?;
+    // Instrumentation bytemap depends on the Wasm memory size, so for larger heaps we need
+    // to pass in the corresponding Wasm64 heap memory size.
+    let max_wasm_memory_size = if has_wasm64_memory(&module) {
+        config.max_wasm64_memory_size
+    } else {
+        config.max_wasm_memory_size
+    };
     let instrumentation_output = instrument(
         module,
         config.cost_to_compile_wasm_instruction,
@@ -212,7 +221,7 @@ fn validate_and_instrument(
         config.metering_type,
         config.subnet_type,
         config.dirty_page_overhead,
-        config.max_wasm_memory_size,
+        max_wasm_memory_size,
         config.max_stable_memory_size,
     )?;
     Ok((wasm_validation_details, instrumentation_output))

rs/embedders/src/wasm_utils/validation.rs

@@ -1068,6 +1068,11 @@ fn validate_function_section(
     Ok(())
 }
 
+// Checks if the module has a Wasm64 memory.
+pub fn has_wasm64_memory(module: &Module) -> bool {
+    module.memories.first().is_some_and(|m| m.memory64)
+}
+
 // Checks that the initial size of the wasm (heap) memory is not larger than
 // the allowed maximum size. This is only needed for Wasm64, because in Wasm32 this
 // is checked by Wasmtime.
@@ -1552,7 +1557,14 @@ pub(super) fn validate_wasm_binary<'a>(
     validate_data_section(&module)?;
     validate_global_section(&module, config.max_globals)?;
     validate_function_section(&module, config.max_functions)?;
-    validate_initial_wasm_memory_size(&module, config.max_wasm_memory_size)?;
+    // The maximum Wasm memory size is different for Wasm32 and Wasm64 and
+    // each needs to be validated accordingly.
+    let max_wasm_memory_size = if has_wasm64_memory(&module) {
+        config.max_wasm64_memory_size
+    } else {
+        config.max_wasm_memory_size
+    };
+    validate_initial_wasm_memory_size(&module, max_wasm_memory_size)?;
     let (largest_function_instruction_count, max_complexity) = validate_code_section(&module)?;
     let wasm_metadata = validate_custom_section(&module, config)?;
     Ok((

rs/embedders/tests/validation.rs

@@ -1178,7 +1178,7 @@ fn test_wasm64_initial_wasm_memory_size_validation() {
         ..Default::default()
     };
     let allowed_wasm_memory_size_in_pages =
-        embedders_config.max_wasm_memory_size.get() / WASM_PAGE_SIZE as u64;
+        embedders_config.max_wasm64_memory_size.get() / WASM_PAGE_SIZE as u64;
     let declared_wasm_memory_size_in_pages = allowed_wasm_memory_size_in_pages + 10;
     let wasm = wat2wasm(&format!(
         r#"(module

rs/embedders/tests/wasmtime_embedder.rs

@@ -3032,7 +3032,7 @@ fn large_wasm64_stable_read_write_test() {
     config.feature_flags.wasm64 = FlagStatus::Enabled;
     config.feature_flags.wasm_native_stable_memory = FlagStatus::Enabled;
     // Declare a large heap.
-    config.max_wasm_memory_size = NumBytes::from(10 * gb);
+    config.max_wasm64_memory_size = NumBytes::from(10 * gb);
 
     let mut instance = WasmtimeInstanceBuilder::new()
         .with_config(config)

rs/execution_environment/benches/lib/src/common.rs

@@ -281,7 +281,7 @@ where
     };
 
     // Set up larger heap, of 8GB for the Wasm64 feature.
-    embedders_config.max_wasm_memory_size = NumBytes::from(8 * 1024 * 1024 * 1024);
+    embedders_config.max_wasm64_memory_size = NumBytes::from(8 * 1024 * 1024 * 1024);
 
     let config = Config {
         embedders_config,

rs/execution_environment/src/canister_manager.rs

@@ -121,7 +121,8 @@ pub(crate) struct CanisterMgrConfig {
     pub(crate) own_subnet_id: SubnetId,
     pub(crate) own_subnet_type: SubnetType,
     pub(crate) max_controllers: usize,
-    pub(crate) max_canister_memory_size: NumBytes,
+    pub(crate) max_canister_memory_size_wasm32: NumBytes,
+    pub(crate) max_canister_memory_size_wasm64: NumBytes,
     pub(crate) rate_limiting_of_instructions: FlagStatus,
     rate_limiting_of_heap_delta: FlagStatus,
     heap_delta_rate_limit: NumBytes,
@@ -141,7 +142,8 @@ impl CanisterMgrConfig {
         own_subnet_type: SubnetType,
         max_controllers: usize,
         compute_capacity: usize,
-        max_canister_memory_size: NumBytes,
+        max_canister_memory_size_wasm32: NumBytes,
+        max_canister_memory_size_wasm64: NumBytes,
         rate_limiting_of_instructions: FlagStatus,
         allocatable_capacity_in_percent: usize,
         rate_limiting_of_heap_delta: FlagStatus,
@@ -160,7 +162,8 @@ impl CanisterMgrConfig {
             max_controllers,
             compute_capacity: (compute_capacity * allocatable_capacity_in_percent.min(100) / 100)
                 as u64,
-            max_canister_memory_size,
+            max_canister_memory_size_wasm32,
+            max_canister_memory_size_wasm64,
             rate_limiting_of_instructions,
             rate_limiting_of_heap_delta,
             heap_delta_rate_limit,
@@ -2195,7 +2198,10 @@ impl CanisterManager {
         let mut new_canister =
             CanisterState::new(system_state, new_execution_state, scheduler_state);
         let new_memory_usage = new_canister.memory_usage();
-        let memory_allocation_given = canister.memory_limit(self.config.max_canister_memory_size);
+
+        let memory_allocation_given =
+            canister.memory_limit(self.get_max_canister_memory_size(is_wasm64_execution));
+
         if new_memory_usage > memory_allocation_given {
             return (
                 Err(CanisterManagerError::NotEnoughMemoryAllocationGiven {
@@ -2336,6 +2342,16 @@ impl CanisterManager {
         );
         Ok(())
     }
+
+    /// Depending on the canister architecture (Wasm32 or Wasm64), returns the
+    /// maximum memory size that can be allocated by a canister.
+    pub(crate) fn get_max_canister_memory_size(&self, is_wasm64_execution: bool) -> NumBytes {
+        if is_wasm64_execution {
+            self.config.max_canister_memory_size_wasm64
+        } else {
+            self.config.max_canister_memory_size_wasm32
+        }
+    }
 }
 
 #[derive(Eq, PartialEq, Debug)]

rs/execution_environment/src/canister_manager/tests.rs

@@ -306,6 +306,7 @@ fn canister_manager_config(
         // TODO(RUN-319): the capacity should be defined based on actual `scheduler_cores`
         100,
         MAX_CANISTER_MEMORY_SIZE,
+        MAX_CANISTER_MEMORY_SIZE,
         rate_limiting_of_instructions,
         100,
         FlagStatus::Enabled,

rs/execution_environment/src/execution_environment.rs

@@ -402,7 +402,8 @@ impl ExecutionEnvironment {
             own_subnet_type,
             config.max_controllers,
             compute_capacity,
-            config.max_canister_memory_size,
+            config.max_canister_memory_size_wasm32,
+            config.max_canister_memory_size_wasm64,
             config.rate_limiting_of_instructions,
             config.allocatable_compute_capacity_in_percent,
             config.rate_limiting_of_heap_delta,
@@ -1805,8 +1806,12 @@ impl ExecutionEnvironment {
 
     /// Returns the maximum amount of memory that can be utilized by a single
     /// canister.
-    pub fn max_canister_memory_size(&self) -> NumBytes {
-        self.config.max_canister_memory_size
+    pub fn max_canister_memory_size(&self, is_wasm64: bool) -> NumBytes {
+        if is_wasm64 {
+            self.config.max_canister_memory_size_wasm64
+        } else {
+            self.config.max_canister_memory_size_wasm32
+        }
     }
 
     /// Returns the subnet memory capacity.
@@ -1823,9 +1828,17 @@ impl ExecutionEnvironment {
         execution_mode: ExecutionMode,
         subnet_memory_saturation: ResourceSaturation,
     ) -> ExecutionParameters {
+        let is_wasm64_execution = match &canister.execution_state {
+            // The canister is not already installed, so we do not know what kind of canister it is.
+            // Therefore we can assume it is Wasm64 because Wasm64 can have a larger memory limit.
+            None => true,
+            Some(execution_state) => execution_state.is_wasm64,
+        };
+        let max_memory_size = self.max_canister_memory_size(is_wasm64_execution);
+
         ExecutionParameters {
             instruction_limits,
-            canister_memory_limit: canister.memory_limit(self.config.max_canister_memory_size),
+            canister_memory_limit: canister.memory_limit(max_memory_size),
             wasm_memory_limit: canister.wasm_memory_limit(),
             memory_allocation: canister.memory_allocation(),
             canister_guaranteed_callback_quota: self.config.canister_guaranteed_callback_quota

rs/execution_environment/src/query_handler.rs

@@ -258,7 +258,6 @@ impl InternalHttpQueryHandler {
         // instruction limit for the whole composite query tree imposes a much lower
         // implicit bound anyway.
         let subnet_available_callbacks = self.config.subnet_callback_soft_limit as i64;
-        let max_canister_memory_size = self.config.max_canister_memory_size;
 
         let mut context = query_context::QueryContext::new(
             &self.log,
@@ -273,7 +272,8 @@ impl InternalHttpQueryHandler {
             subnet_available_memory,
             subnet_available_callbacks,
             self.config.canister_guaranteed_callback_quota as u64,
-            max_canister_memory_size,
+            self.config.max_canister_memory_size_wasm32,
+            self.config.max_canister_memory_size_wasm64,
             self.max_instructions_per_query,
             self.config.max_query_call_graph_depth,
             self.config.max_query_call_graph_instructions,

rs/execution_environment/src/query_handler/query_context.rs

@@ -95,7 +95,8 @@ pub(super) struct QueryContext<'a> {
     network_topology: Arc<NetworkTopology>,
     // Certificate for certified queries + canister ID of the root query of this context
     data_certificate: (Vec<u8>, CanisterId),
-    max_canister_memory_size: NumBytes,
+    max_canister_memory_size_wasm32: NumBytes,
+    max_canister_memory_size_wasm64: NumBytes,
     max_instructions_per_query: NumInstructions,
     max_query_call_graph_depth: usize,
     instruction_overhead_per_query_call: RoundInstructions,
@@ -130,7 +131,8 @@ impl<'a> QueryContext<'a> {
         subnet_available_memory: SubnetAvailableMemory,
         subnet_available_callbacks: i64,
         canister_guaranteed_callback_quota: u64,
-        max_canister_memory_size: NumBytes,
+        max_canister_memory_size_wasm32: NumBytes,
+        max_canister_memory_size_wasm64: NumBytes,
         max_instructions_per_query: NumInstructions,
         max_query_call_graph_depth: usize,
         max_query_call_graph_instructions: NumInstructions,
@@ -158,7 +160,8 @@ impl<'a> QueryContext<'a> {
             state,
             network_topology,
             data_certificate: (data_certificate, canister_id),
-            max_canister_memory_size,
+            max_canister_memory_size_wasm32,
+            max_canister_memory_size_wasm64,
             max_instructions_per_query,
             max_query_call_graph_depth,
             instruction_overhead_per_query_call: as_round_instructions(
@@ -1083,9 +1086,20 @@ impl<'a> QueryContext<'a> {
         canister: &CanisterState,
         instruction_limits: InstructionLimits,
     ) -> ExecutionParameters {
+        let is_wasm64_execution = canister
+            .execution_state
+            .as_ref()
+            .is_some_and(|es| es.is_wasm64);
+
+        let max_canister_memory_size = if is_wasm64_execution {
+            self.max_canister_memory_size_wasm64
+        } else {
+            self.max_canister_memory_size_wasm32
+        };
+
         ExecutionParameters {
             instruction_limits,
-            canister_memory_limit: canister.memory_limit(self.max_canister_memory_size),
+            canister_memory_limit: canister.memory_limit(max_canister_memory_size),
             wasm_memory_limit: canister.wasm_memory_limit(),
             memory_allocation: canister.memory_allocation(),
             canister_guaranteed_callback_quota: self.canister_guaranteed_callback_quota,

rs/execution_environment/src/scheduler.rs

@@ -1072,7 +1072,15 @@ impl SchedulerImpl {
     ) -> bool {
         for canister_id in canister_ids {
             let canister = state.canister_states.get(canister_id).unwrap();
-            if let Err(err) = canister.check_invariants(self.exec_env.max_canister_memory_size()) {
+
+            let canister_is_wasm64 = canister
+                .execution_state
+                .as_ref()
+                .is_some_and(|es| es.is_wasm64);
+
+            if let Err(err) = canister
+                .check_invariants(self.exec_env.max_canister_memory_size(canister_is_wasm64))
+            {
                 let msg = format!(
                     "{}: At Round {} @ time {}, canister {} has invalid state after execution. Invariant check failed with err: {}",
                     CANISTER_INVARIANT_BROKEN,