fix(benchmarks): update reinstrument bench (#4000)

pallets/gear/src/benchmarking/code.rs

@@ -104,11 +104,21 @@ pub struct ModuleDefinition {
     pub stack_end: Option<WasmPage>,
 }
 
+#[derive(Default)]
+pub enum InitElements {
+    NoInit,
+    Number(u32),
+    #[default]
+    All,
+}
+
 pub struct TableSegment {
     /// How many elements should be created inside the table.
     pub num_elements: u32,
     /// The function index with which all table elements should be initialized.
     pub function_index: u32,
+    /// Generate element segment which initialize the table.
+    pub init_elements: InitElements,
 }
 
 pub struct TypeSegment {
@@ -312,12 +322,21 @@ where
 
         // Add function pointer table
         if let Some(table) = def.table {
-            program = program
+            let mut table_builder = program
                 .table()
                 .with_min(table.num_elements)
-                .with_max(Some(table.num_elements))
-                .with_element(0, vec![table.function_index; table.num_elements as usize])
-                .build();
+                .with_max(Some(table.num_elements));
+
+            table_builder = match table.init_elements {
+                InitElements::NoInit => table_builder,
+                InitElements::Number(num) => {
+                    table_builder.with_element(0, vec![table.function_index; num as usize])
+                }
+                InitElements::All => table_builder
+                    .with_element(0, vec![table.function_index; table.num_elements as usize]),
+            };
+
+            program = table_builder.build();
         }
 
         // Add the dummy section
@@ -458,22 +477,24 @@ where
         module.into()
     }
 
-    /// Creates a WebAssembly module with a table size of `target_bytes` bytes.
+    /// Creates a WebAssembly module with a table size of `num_elements` bytes.
     /// Each element in the table points to function index `0` and occupies 1 byte.
-    pub fn sized_table_section(target_bytes: u32) -> Self {
+    pub fn sized_table_section(num_elements: u32, init_elements: Option<u32>) -> Self {
         let mut module = ModuleDefinition {
             memory: Some(ImportedMemory::max::<T>()),
             ..Default::default()
         };
 
         module.init_body = Some(body::empty());
 
-        // 1 element with function index value `0` takes 1 byte to encode.
-        let num_elements = target_bytes;
-
+        // 1 element with function index value `0` occupies 1 byte
         module.table = Some(TableSegment {
             num_elements,
             function_index: 0,
+            init_elements: match init_elements {
+                Some(num) => InitElements::Number(num),
+                None => InitElements::NoInit,
+            },
         });
 
         module.into()

pallets/gear/src/benchmarking/mod.rs

@@ -116,6 +116,7 @@ const MAX_PAYLOAD_LEN: u32 = 32 * 64 * 1024;
 const MAX_PAYLOAD_LEN_KB: u32 = MAX_PAYLOAD_LEN / 1024;
 const MAX_SALT_SIZE_BYTES: u32 = 4 * 1024 * 1024;
 const MAX_NUMBER_OF_DATA_SEGMENTS: u32 = 1024;
+const MAX_TABLE_ENTRIES: u32 = 10_000_000;
 
 /// How many batches we do per API benchmark.
 const API_BENCHMARK_BATCHES: u32 = 20;
@@ -405,11 +406,11 @@ benchmarks! {
         Environment::new(ext, &code, DispatchKind::Init, Default::default(), max_pages::<T>().into()).unwrap();
     }
 
-    // `t`: Size of the table section in kilobytes.
+    // `t`: Size of the memory allocated for the table after instantiation, in kilobytes.
     instantiate_module_table_section_per_kb {
-        let t in 0 .. T::Schedule::get().limits.code_len / 1024;
+        let t in 0 .. MAX_TABLE_ENTRIES / 1024;
 
-        let WasmModule { code, .. } = WasmModule::<T>::sized_table_section(t * 1024);
+        let WasmModule { code, .. } = WasmModule::<T>::sized_table_section(t * 1024, None);
         let ext = Ext::new(ProcessorContext::new_mock());
     }: {
         Environment::new(ext, &code, DispatchKind::Init, Default::default(), max_pages::<T>().into()).unwrap();
@@ -581,8 +582,11 @@ benchmarks! {
     // first time after a new schedule was deployed: For every new schedule a program needs
     // to re-run the instrumentation once.
     reinstrument_per_kb {
-        let c in 0 .. T::Schedule::get().limits.code_len / 1_024;
-        let WasmModule { code, hash, .. } = WasmModule::<T>::sized(c * 1_024, Location::Handle);
+        let e in 0 .. T::Schedule::get().limits.code_len / 1_024;
+
+        let max_table_size = T::Schedule::get().limits.code_len;
+        // NOTE: We use a program filled with table/element sections here because it is the heaviest weight-wise.
+        let WasmModule { code, hash, .. } = WasmModule::<T>::sized_table_section(max_table_size, Some(e * 1024));
         let code = Code::try_new_mock_const_or_no_rules(code, false, Default::default()).unwrap();
         let code_and_id = CodeAndId::new(code);
         let code_id = code_and_id.code_id();
@@ -1732,6 +1736,7 @@ benchmarks! {
             table: Some(TableSegment {
                 num_elements,
                 function_index: OFFSET_AUX,
+                init_elements: Default::default(),
             }),
             .. Default::default()
         }));
@@ -1756,6 +1761,7 @@ benchmarks! {
             table: Some(TableSegment {
                 num_elements,
                 function_index: OFFSET_AUX,
+                init_elements: Default::default(),
             }),
             .. Default::default()
         }));