Import matrix-multiply from a separate wasm module (#232)

* Updated marian-dev submodule
* Import wasm gemm from a separate wasm module
 - The fallback implementation of gemm is currently being imported dynamically
   for wasm target
* Updated CI scripts and README to import GEMM from a separate wasm module
* Setting model config to int8shiftAlphaAll in wasm test page

.github/workflows/wasm-custom_marian-mac.yml

@@ -39,6 +39,10 @@ jobs:
         working-directory: build-wasm
         run: bash ../wasm/patch-artifacts-enable-wormhole.sh
 
+      - name: Import GEMM library from a separate wasm module
+        working-directory: build-wasm
+        run: bash ../wasm/patch-artifacts-import-gemm-module.sh
+
       - name: Check artifacts
         working-directory: build-wasm
         run: |

.github/workflows/wasm-custom_marian-ubuntu.yml

@@ -39,6 +39,10 @@ jobs:
         working-directory: build-wasm
         run: bash ../wasm/patch-artifacts-enable-wormhole.sh
 
+      - name: Import GEMM library from a separate wasm module
+        working-directory: build-wasm
+        run: bash ../wasm/patch-artifacts-import-gemm-module.sh
+
       - name: Check artifacts
         working-directory: build-wasm
         run: |

3rd_party/marian-dev

@@ -1 +1 @@
-Subproject commit 62bac858bfd37060beb707d12eb9711649ea4cf6
+Subproject commit a1a82ff64910dc066d64d631cd7a8212df9f88cd

README.md

@@ -46,6 +46,11 @@ To build a version that translates with higher speeds on Firefox Nightly browser
        bash ../wasm/patch-artifacts-enable-wormhole.sh
        ```
 
+   3. Patch generated artifacts to import GEMM library from a separate wasm module
+       ```bash
+       bash ../wasm/patch-artifacts-import-gemm-module.sh
+       ```
+
 To build a version that runs on all browsers (including Firefox Nightly) but translates slowly, follow these instructions:
 
   1. Create a folder where you want to build all the artifacts (`build-wasm` in this case) and compile
@@ -56,6 +61,11 @@ To build a version that runs on all browsers (including Firefox Nightly) but tra
       emmake make -j2
       ```
 
+  2. Patch generated artifacts to import GEMM library from a separate wasm module
+       ```bash
+       bash ../wasm/patch-artifacts-import-gemm-module.sh
+       ```
+
 #### Recompiling
 As long as you don't update any submodule, just follow [Compile](#Compile) steps.\
 If you update a submodule, execute following command in repository root folder before executing

build-wasm.sh

@@ -78,5 +78,8 @@ if [ "$WORMHOLE" = true ]; then
   bash ../wasm/patch-artifacts-enable-wormhole.sh
 fi
 
+#     3. Import GEMM library from a separate wasm module
+bash ../wasm/patch-artifacts-import-gemm-module.sh
+
 # The artifacts (.js and .wasm files) will be available in the build directory
 exit 0

wasm/CMakeLists.txt

@@ -26,6 +26,12 @@ set(LINKER_FLAGS "${LINKER_FLAGS} -s ENVIRONMENT=web,worker")
 # Append version information in the Javascript artifact
 set(LINKER_FLAGS "${LINKER_FLAGS} --extern-pre-js ${CMAKE_CURRENT_BINARY_DIR}/project_version.js")
 
+# Allow importing undefined symbols dynamically
+set(LINKER_FLAGS "${LINKER_FLAGS} -s ERROR_ON_UNDEFINED_SYMBOLS=0 -s DECLARE_ASM_MODULE_EXPORTS=0")
+
+# Export all the functions of fallback implementation of GEMM for wasm target
+set(LINKER_FLAGS "${LINKER_FLAGS} -s EXPORTED_FUNCTIONS=[_int8PrepareAFallback,_int8PrepareBFallback,_int8PrepareBFromTransposedFallback,_int8PrepareBFromQuantizedTransposedFallback,_int8PrepareBiasFallback,_int8MultiplyAndAddBiasFallback,_int8SelectColumnsOfBFallback]")
+
 set_target_properties(bergamot-translator-worker PROPERTIES
                         SUFFIX ".js"
                         LINK_FLAGS ${LINKER_FLAGS}

wasm/patch-artifacts-import-gemm-module.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+usage="Patch wasm artifacts to import fallback implementation of gemm for wasm.
+
+Usage: $(basename "$0") [WASM_ARTIFACTS_FOLDER]
+
+    where:
+    WASM_ARTIFACTS_FOLDER    Folder containing wasm artifacts
+                             (An optional argument, if unspecified the default is: current folder)"
+
+if [ "$#" -gt 1 ]; then
+    echo "Illegal number of parameters passed"
+    echo "$usage"
+    exit
+fi
+
+# Parse wasm artifacts folder if provided via script argument or set it to default
+WASM_ARTIFACTS_FOLDER=$PWD
+if [ "$#" -eq 1 ]; then
+    if [ ! -e "$1" ]; then
+        echo "Error: Folder \""$1"\" doesn't exist"
+        exit
+    fi
+    WASM_ARTIFACTS_FOLDER="$1"
+fi
+
+WASM_ARTIFACTS_JAVASCRIPT_FILE="bergamot-translator-worker.js"
+WASM_ARTIFACTS="$WASM_ARTIFACTS_FOLDER/${WASM_ARTIFACTS_JAVASCRIPT_FILE}"
+if [ ! -e "$WASM_ARTIFACTS" ]; then
+    echo "Error: Artifact \"$WASM_ARTIFACTS\" doesn't exist"
+    exit
+fi
+
+echo "Polyfill the fallback integer (8-bit) gemm implementation from the main module"
+sed -i.bak 's/"env"[[:space:]]*:[[:space:]]*asmLibraryArg,/"env": asmLibraryArg,\
+    "wasm_gemm":{\
+    "int8_prepare_a": (...a) => Module["asm"].int8PrepareAFallback(...a),\
+    "int8_prepare_b": (...a) => Module["asm"].int8PrepareBFallback(...a),\
+    "int8_prepare_b_from_transposed": (...a) => Module["asm"].int8PrepareBFromTransposedFallback(...a),\
+    "int8_prepare_b_from_quantized_transposed": (...a) => Module["asm"].int8PrepareBFromQuantizedTransposedFallback(...a),\
+    "int8_prepare_bias": (...a) => Module["asm"].int8PrepareBiasFallback(...a),\
+    "int8_multiply_and_add_bias": (...a) => Module["asm"].int8MultiplyAndAddBiasFallback(...a),\
+    "int8_select_columns_of_b": (...a) => Module["asm"].int8SelectColumnsOfBFallback(...a),\
+    },/g' ${WASM_ARTIFACTS_JAVASCRIPT_FILE}
+echo "SUCCESS"

wasm/test_page/js/worker.js

@@ -180,7 +180,7 @@ skip-cost: true
 cpu-threads: 0
 quiet: true
 quiet-translation: true
-gemm-precision: int8shiftAll
+gemm-precision: int8shiftAlphaAll
 `;
 
   const modelFile = `${rootURL}/${languagePair}/${modelRegistry[languagePair]["model"].name}`;