diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a39685f..b34f3b2 100644 --- a/.github/workflows/ci.yml +++ b/.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: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..9ae7931 --- /dev/null +++ b/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. diff --git a/benches/benchloop.py b/benches/benchloop.py index 5a125ae..b81553d 100755 --- a/benches/benchloop.py +++ b/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,14 +48,12 @@ 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") @@ -54,11 +61,18 @@ def bench_iteration(sizes, ty, nc, kc, mc, *, layout, threads, file, sleep): 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) @@ -66,8 +80,8 @@ def bench_iteration(sizes, ty, nc, kc, mc, *, layout, threads, file, 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() diff --git a/justfile b/justfile new file mode 100644 index 0000000..b436779 --- /dev/null +++ b/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}} diff --git a/src/lib.rs b/src/lib.rs index 60f09d7..9f1d132 100644 --- a/src/lib.rs +++ b/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` diff --git a/src/sgemm_kernel.rs b/src/sgemm_kernel.rs index 305fd07..afae886 100644 --- a/src/sgemm_kernel.rs +++ b/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