Improved wasm benchmark support and add support for wasm +relaxed-simd

.github/workflows/ci.yml

@@ -169,6 +169,10 @@ jobs:
             target: wasm32-wasip1
             rustflags: "-C target-feature=+simd128"
             name: simd128
+          - rust: stable
+            target: wasm32-wasip1
+            rustflags: "-C target-feature=+simd128,+relaxed-simd"
+            name: simd128 relaxed-simd
 
     name: wasm_test/${{ matrix.target }}/${{ matrix.name }}
     steps:
@@ -184,8 +188,7 @@ jobs:
           cargo test -v --tests --lib --release --no-fail-fast --target "${{ matrix.target }}"
         env:
           RUSTFLAGS: ${{ matrix.rustflags }}
-          CARGO_TARGET_WASM32_WASIP1_RUNNER: "wasmtime --dir=."
-          MMTEST_FAST_TEST: 1
+          CARGO_TARGET_WASM32_WASIP1_RUNNER: "wasmtime"
 
 
   cargo-careful:

CONTRIBUTING.md

@@ -0,0 +1,18 @@
+
+## Guidelines
+
+As a current guiding principle, the entrance functions (i.e sgemm, dgemm etc) are non-generic by design, so that the compile time cost of the library is limited.
+
+Threading is supported using thread-tree for its lower dispatch overhead but it's welcome to replace it by rayon - if the low dispatch overhead can be preserved.
+
+## Test tricks
+
+Use MMTEST_FEATURE=fma and so on to restrict target feature detection to the given feature. Note that this currently only supports a single target feature at a time.
+
+## Benchmarks
+
+To run benchmarks, use `./benches/benchloop.py`
+
+## Wasm
+
+To test and benchmark wasm, add the wasm32-wasip1 target using rustup and install wasmtime-cli.

benches/benchloop.py

@@ -5,6 +5,7 @@
 """
 
 import argparse
+import itertools
 import os
 import re
 import subprocess
@@ -22,7 +23,15 @@
 
 _COMPILE = "cargo rustc --example benchmark --release".split()
 _DEFAULT_FEATURES = "constconf cgemm".split()
-_EXEC = "./target/release/examples/benchmark"
+_EXEC = ["./target/release/examples/benchmark"]
+
+_WASM_TARGET = "wasm32-wasip1"
+_WASM_RUSTFLAGS = {
+    "none": "",
+    "simd128": "-C target-feature=+simd128",
+    "relaxed": "-C target-feature=+simd128,+relaxed-simd",
+}
+_WASM_EXEC = ["wasmtime", "run", f"./target/{_WASM_TARGET}/release/examples/benchmark.wasm"]
 
 
 def layout_type(value):
@@ -39,35 +48,40 @@ def bench_loop(args, *, file):
     mcs = [None] if args.mc is None else args.mc
     for threads in args.threads:
         for ty in args.type:
-            for layout in args.layout:
-                for nc in ncs:
-                    for kc in kcs:
-                        for mc in mcs:
-                            bench_iteration(args.size, ty, nc, kc, mc, layout=layout, threads=threads, file=file, sleep=args.sleep)
+            for nc, kc, mc in itertools.product(ncs, kcs, mcs):
+                bench_iteration(args.size, ty, nc, kc, mc,
+                                layouts=args.layout, threads=threads, file=file, sleep=args.sleep, wasm=args.wasm)
 
 
-def bench_iteration(sizes, ty, nc, kc, mc, *, layout, threads, file, sleep):
+def bench_iteration(sizes, ty, nc, kc, mc, *, layouts, threads, file, sleep, wasm):
     features = list(_DEFAULT_FEATURES)
     if threads > 0:
         features.append("threading")
     compile_argv = list(_COMPILE)
     compile_argv.append("--features=" + ",".join(features))
     file.flush()
     env = os.environ.copy()
+    if wasm is not None:
+        compile_argv.extend(["--target", _WASM_TARGET])
+        env["RUSTFLAGS"] = _WASM_RUSTFLAGS[wasm]
     for value, name in zip([nc, kc, mc], ["nc", "kc", "mc"]):
         if value is not None:
             env["MATMUL_" + _GEMMTYPE[ty] + "_" + name.upper()] = str(value)
 
     print("Running", " ".join(compile_argv), file=sys.stderr)
+    flags = env.get("RUSTFLAGS", "")
+    if flags:
+        print("Using RUSTFLAGS='", flags, "'", sep="")
+
     subprocess.run(compile_argv, env=env)
 
     time.sleep(_POST_COMPILE_SLEEP)
 
     exec_env = os.environ.copy()
     exec_env["MATMUL_NUM_THREADS"] = str(threads)
     extra_column = ",".join((str(value) if value is not None else "") for value in [nc, kc, mc, threads])
-    for size in sizes:
-        argv = [_EXEC]
+    for size, layout in itertools.product(sizes, layouts):
+        argv = list(_WASM_EXEC if wasm else _EXEC)
         argv.extend(["--type", ty])
         argv.extend(["--csv", "--layout", layout, "--extra-column", extra_column])
         argv.extend([str(size)] * 3)
@@ -89,6 +103,9 @@ def main():
                         help="Thread use. 0: not enabled; 1: enabled but one thread; n: enabled with n threads.")
     parser.add_argument("--layout", type=layout_type, default=["fcc"], nargs="+",
                         help="Layout (f/c combos, e.g. fcc, fff)")
+    parser.add_argument("--wasm", nargs="?", const="simd128", choices=["none", "simd128", "relaxed"],
+                        help="Build for wasm32-wasip1 and run the benchmark through wasmtime, "
+                        "optionally select which features to enable.")
     parser.add_argument("--sleep", type=int, default=1, help="Time to wait between every run")
     parser.add_argument("--output", type=str, default=None, help="Output file (csv format)")
     args = parser.parse_args()

justfile

@@ -0,0 +1,5 @@
+# This file exists to document some auxiliary commands for development;
+# for most development tasks prefer cargo; i.e. cargo test; which does not need to be here.
+wasmtest *args:
+    CARGO_TARGET_WASM32_WASIP1_RUNNER=wasmtime RUSTFLAGS="-C target-feature=+simd128" cargo test --target wasm32-wasip1 --features=cgemm {{args}}
+    CARGO_TARGET_WASM32_WASIP1_RUNNER=wasmtime RUSTFLAGS="-C target-feature=+simd128,+relaxed-simd" cargo test --target wasm32-wasip1 --features=cgemm {{args}}

src/lib.rs

@@ -64,6 +64,11 @@
 //!
 //!   - `neon`
 //!
+//! - *wasm32* target features need to be enabled at compile-time. The library
+//!   can take advantage of these if enabled:
+//!
+//!   - `simd128` (sgemm) (will also use `relaxed-simd` if enabled)
+//!
 //! ## Features
 //!
 //! ### `std`

src/sgemm_kernel.rs

@@ -680,6 +680,20 @@ unsafe fn kernel_target_wasm_simd(k: usize, alpha: T, a: *const T, b: *const T,
     const MR: usize = KernelWasmSimd::MR;
     const NR: usize = KernelWasmSimd::NR;
 
+    // Use f32x4_relaxed_madd when enabled
+    // by spec relaxed_madd is a fused multiply-add when possible, otherwise it is a multiply then add
+    #[cfg(target_feature = "relaxed-simd")]
+    #[inline(always)]
+    unsafe fn muladd(a: v128, b: v128, c: v128) -> v128 {
+        f32x4_relaxed_madd(a, b, c)
+    }
+
+    #[cfg(not(target_feature = "relaxed-simd"))]
+    #[inline(always)]
+    unsafe fn muladd(a: v128, b: v128, c: v128) -> v128 {
+        f32x4_add(f32x4_mul(a, b), c)
+    }
+
     let (mut a, mut b, rsc, csc) = if rsc == 1 { (b, a, csc, rsc) } else { (a, b, rsc, csc) };
 
     // Kernel 8 x 8 (a x b)
@@ -690,28 +704,38 @@ unsafe fn kernel_target_wasm_simd(k: usize, alpha: T, a: *const T, b: *const T,
     let mut ab21 = [zero; 4];
     let mut ab22 = [zero; 4];
 
+    /// multiple sets of ptr => v1, v2 supported
+    macro_rules! load_vectors {
+        ($($a:ident => $a1:ident, $a2:ident),+) => {
+            $(
+                $a1 = v128_load($a as *const v128);
+                $a2 = v128_load($a.add(4) as *const v128);
+            )+
+        }
+    }
+
+    macro_rules! splat_lane {
+        ($v:expr, $n:tt) =>  { i32x4_shuffle::<$n, $n, $n, $n>($v, $v) };
+    }
+
     // ab_ij = a_i * b_j for all i, j
-    // (wasm SIMD has no lane-FMA; extract+splat into mul+add)
     macro_rules! ab_ij_equals_ai_bj {
         ($dest:ident, $av:expr, $bv:expr) => {
-            $dest[0] = f32x4_add($dest[0], f32x4_mul($bv, f32x4_splat(f32x4_extract_lane::<0>($av))));
-            $dest[1] = f32x4_add($dest[1], f32x4_mul($bv, f32x4_splat(f32x4_extract_lane::<1>($av))));
-            $dest[2] = f32x4_add($dest[2], f32x4_mul($bv, f32x4_splat(f32x4_extract_lane::<2>($av))));
-            $dest[3] = f32x4_add($dest[3], f32x4_mul($bv, f32x4_splat(f32x4_extract_lane::<3>($av))));
+            $dest[0] = muladd($bv, splat_lane!($av, 0), $dest[0]);
+            $dest[1] = muladd($bv, splat_lane!($av, 1), $dest[1]);
+            $dest[2] = muladd($bv, splat_lane!($av, 2), $dest[2]);
+            $dest[3] = muladd($bv, splat_lane!($av, 3), $dest[3]);
         }
     }
 
-    for _ in 0..k {
-        let a1 = v128_load(a as *const v128);
-        let b1 = v128_load(b as *const v128);
-        let a2 = v128_load(a.add(4) as *const v128);
-        let b2 = v128_load(b.add(4) as *const v128);
+    let (mut a1, mut b1, mut a2, mut b2);
 
+    for _ in 0..k {
+        load_vectors!(a => a1, a2, b => b1, b2);
         ab_ij_equals_ai_bj!(ab11, a1, b1);
         ab_ij_equals_ai_bj!(ab12, a1, b2);
         ab_ij_equals_ai_bj!(ab21, a2, b1);
         ab_ij_equals_ai_bj!(ab22, a2, b2);
-
         a = a.add(MR);
         b = b.add(NR);
     }
@@ -760,10 +784,10 @@ unsafe fn kernel_target_wasm_simd(k: usize, alpha: T, a: *const T, b: *const T,
 
         let betav = f32x4_splat(beta);
         // ab += β C
-        loop4!(i, ab11[i] = f32x4_add(ab11[i], f32x4_mul(c11[i], betav)));
-        loop4!(i, ab12[i] = f32x4_add(ab12[i], f32x4_mul(c12[i], betav)));
-        loop4!(i, ab21[i] = f32x4_add(ab21[i], f32x4_mul(c21[i], betav)));
-        loop4!(i, ab22[i] = f32x4_add(ab22[i], f32x4_mul(c22[i], betav)));
+        loop4!(i, ab11[i] = muladd(betav, c11[i], ab11[i]));
+        loop4!(i, ab12[i] = muladd(betav, c12[i], ab12[i]));
+        loop4!(i, ab21[i] = muladd(betav, c21[i], ab21[i]));
+        loop4!(i, ab22[i] = muladd(betav, c22[i], ab22[i]));
     }
 
     // c <- ab