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

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;
@@ -201,6 +203,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,
@@ -209,7 +218,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,
@@ -2197,7 +2200,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 {
@@ -2338,6 +2344,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

@@ -1070,7 +1070,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,