Matrix multiply benches

dfinity · Dec 19, 2024 · 65db4af · 65db4af
1 parent b861f9f
commit 65db4af
Show file tree

Hide file tree

Showing 8 changed files with 184 additions and 2 deletions.
diff --git a/rs/embedders/BUILD.bazel b/rs/embedders/BUILD.bazel
@@ -206,6 +206,7 @@ rust_ic_bench(
     name = "stable_memory_bench",
     testonly = True,
     srcs = ["benches/stable_memory.rs"],
+    compile_data = glob(["benches/matrix-multiply.wasm"]),
     data = DATA + UNIVERSAL_CANISTER_TEST_DEPS,
     env = ENV | UNIVERSAL_CANISTER_TEST_ENV,
     deps = [

diff --git a/rs/embedders/benches/matrix-multiply.wasm b/rs/embedders/benches/matrix-multiply.wasm
diff --git a/rs/embedders/benches/stable_memory.rs b/rs/embedders/benches/stable_memory.rs
@@ -335,6 +335,48 @@ fn vec_u64_sparse_write(c: &mut Criterion) {
     );
 }
 
+const N: u64 = 1024;
+const D: u64 = 1024;
+
+fn multiply_heap(c: &mut Criterion) {
+    embedders_bench::update_bench(
+        c,
+        "multiply_heap",
+        include_bytes!("matrix-multiply.wasm"),
+        &Encode!(&N, &D).unwrap(),
+        "multiply_heap",
+        &Encode!(&()).unwrap(),
+        None,
+        PostSetupAction::PerformCheckpoint,
+    )
+}
+
+fn multiply_stable(c: &mut Criterion) {
+    embedders_bench::update_bench(
+        c,
+        "multiply_stable",
+        include_bytes!("matrix-multiply.wasm"),
+        &Encode!(&N, &D).unwrap(),
+        "multiply_stable",
+        &Encode!(&()).unwrap(),
+        None,
+        PostSetupAction::PerformCheckpoint,
+    )
+}
+
+fn multiply_stable_old(c: &mut Criterion) {
+    embedders_bench::update_bench(
+        c,
+        "multiply_stable_old",
+        include_bytes!("matrix-multiply.wasm"),
+        &Encode!(&N, &D).unwrap(),
+        "multiply_stable_old",
+        &Encode!(&()).unwrap(),
+        None,
+        PostSetupAction::PerformCheckpoint,
+    )
+}
+
 criterion_group!(
     name = benches10;
     config = Criterion::default().sample_size(10);
@@ -349,4 +391,10 @@ criterion_group!(
       vec_u64_single_read, vec_u64_sparse_read, vec_u64_single_write,
 );
 
-criterion_main!(benches10, benches100);
+criterion_group!(
+    name = matrix;
+    config = Criterion::default();
+    targets = multiply_heap, multiply_stable, multiply_stable_old,
+);
+
+criterion_main!(benches10, benches100, matrix);
diff --git a/rs/embedders/src/wasm_utils.rs b/rs/embedders/src/wasm_utils.rs
@@ -156,6 +156,8 @@ enum SystemApiFunc {
     Stable64Read,
     StableWrite,
     Stable64Write,
+    StableLoadV128,
+    StableStoreI32,
 }
 
 impl SystemApiFunc {
@@ -169,6 +171,8 @@ impl SystemApiFunc {
             "stable64_read" => Some(Self::Stable64Read),
             "stable_write" => Some(Self::StableWrite),
             "stable64_write" => Some(Self::Stable64Write),
+            "stable_read_v128" => Some(Self::StableLoadV128),
+            "stable_write_i32" => Some(Self::StableStoreI32),
             _ => None,
         }
     }

diff --git a/rs/embedders/src/wasm_utils/instrumentation.rs b/rs/embedders/src/wasm_utils/instrumentation.rs
@@ -1151,7 +1151,7 @@ fn replace_system_api_functions(
     main_memory_type: WasmMemoryType,
     max_wasm_memory_size: NumBytes,
 ) {
-    let api_indexes = calculate_api_indexes(module);
+    let api_indexes = dbg!(calculate_api_indexes(module));
     let number_of_func_imports = module
         .imports
         .iter()

diff --git a/rs/embedders/src/wasm_utils/system_api_replacements.rs b/rs/embedders/src/wasm_utils/system_api_replacements.rs
@@ -63,6 +63,99 @@ pub(super) fn replacement_functions(
     };
 
     vec![
+        (
+            SystemApiFunc::StableLoadV128,
+            (
+                FuncType::new([ValType::I64], [ValType::V128]),
+                Body {
+                    locals: vec![(1, ValType::I32)],
+                    instructions: vec![
+                        // Calculate bytemap index
+                        LocalGet { local_index: 0 },
+                        I64Const {
+                            value: page_size_shift as i64,
+                        },
+                        I64ShrU,
+                        I32WrapI64,
+                        LocalSet { local_index: 1 },
+                        LocalGet { local_index: 1 }, // address for store later
+                        LocalGet { local_index: 1 },
+                        I32Load8U {
+                            memarg: wasmparser::MemArg {
+                                align: 0,
+                                max_align: 0,
+                                offset: 0,
+                                // We assume the bytemap for stable memory is always
+                                // inserted directly after the stable memory.
+                                memory: special_indices.stable_memory_index + 1,
+                            },
+                        },
+                        I32Const { value: 2 }, // READ_BIT
+                        I32Or,
+                        I32Store8 {
+                            memarg: wasmparser::MemArg {
+                                align: 0,
+                                max_align: 0,
+                                offset: 0,
+                                // We assume the bytemap for stable memory is always
+                                // inserted directly after the stable memory.
+                                memory: special_indices.stable_memory_index + 1,
+                            },
+                        },
+                        LocalGet { local_index: 0 },
+                        V128Load {
+                            memarg: wasmparser::MemArg {
+                                align: 0,
+                                max_align: 0,
+                                offset: 0,
+                                memory: special_indices.stable_memory_index,
+                            },
+                        },
+                        End,
+                    ],
+                },
+            ),
+        ),
+        (
+            SystemApiFunc::StableStoreI32,
+            (
+                FuncType::new([ValType::I64, ValType::I32], []),
+                Body {
+                    locals: vec![],
+                    instructions: vec![
+                        // Calculate bytemap index
+                        LocalGet { local_index: 0 },
+                        I64Const {
+                            value: page_size_shift as i64,
+                        },
+                        I64ShrU,
+                        I32WrapI64,
+                        I32Const { value: 3 }, // Value for accessed and dirty
+                        I32Store8 {
+                            memarg: wasmparser::MemArg {
+                                align: 0,
+                                max_align: 0,
+                                offset: 0,
+                                // We assume the bytemap for stable memory is always
+                                // inserted directly after the stable memory.
+                                memory: special_indices.stable_memory_index + 1,
+                            },
+                        },
+                        LocalGet { local_index: 0 },
+                        LocalGet { local_index: 1 },
+                        I32Store {
+                            memarg: wasmparser::MemArg {
+                                align: 0,
+                                max_align: 0,
+                                offset: 0,
+                                memory: special_indices.stable_memory_index,
+                            },
+                        },
+                        End,
+                    ],
+                },
+            ),
+        ),
         (
             SystemApiFunc::StableSize,
             (

diff --git a/rs/embedders/src/wasm_utils/validation.rs b/rs/embedders/src/wasm_utils/validation.rs
@@ -417,6 +417,26 @@ fn get_valid_system_apis_common(I: ValType) -> HashMap<String, HashMap<String, F
                 },
             )],
         ),
+        (
+            "stable_read_v128",
+            vec![(
+                API_VERSION_IC0,
+                FunctionSignature {
+                    param_types: vec![ValType::I64],
+                    return_type: vec![ValType::V128],
+                },
+            )],
+        ),
+        (
+            "stable_write_i32",
+            vec![(
+                API_VERSION_IC0,
+                FunctionSignature {
+                    param_types: vec![ValType::I64, ValType::I32],
+                    return_type: vec![],
+                },
+            )],
+        ),
         (
             "time",
             vec![(

diff --git a/rs/embedders/src/wasmtime_embedder/system_api.rs b/rs/embedders/src/wasmtime_embedder/system_api.rs
@@ -389,6 +389,22 @@ pub fn syscalls<
         // LINT.ThenChange(logging_charge_bytes_rule)
     }
 
+    linker
+        .func_wrap(
+            "ic0",
+            "stable_read_v128",
+            move |_caller: Caller<'_, StoreData>, _src: i64| -> wasmtime::V128 { unimplemented!() },
+        )
+        .unwrap();
+
+    linker
+        .func_wrap(
+            "ic0",
+            "stable_write_i32",
+            move |_caller: Caller<'_, StoreData>, _dst: i64, _val: i32| unimplemented!(),
+        )
+        .unwrap();
+
     linker
         .func_wrap("ic0", "msg_caller_copy", {
             move |mut caller: Caller<'_, StoreData>, dst: I, offset: I, size: I| {