Add WebGPU Emscripten platform/driver/samples. (#10311)

Rollup of fixes (some simple, a few hacky as we figure out whether we
want to continue supporting native WebGPU platforms), then the actual
Emscripten platform code and samples.

The sample is forked from
[`experimental/web/sample_dynamic/`](https://github.com/iree-org/iree/tree/main/experimental/web/sample_dynamic),
but uses WebGPU instead of WebAssembly with Web Workers.

Both the new sample and the tests at
[`experimental/web/testing/`](https://github.com/iree-org/iree/tree/main/experimental/web/testing)
can successfully create WebGPU devices using Emscripten on my
machine/browser (Chrome Canary on Windows with `#enable-unsafe-webgpu`).
The sample and tests don't get much further than device creation yet
(needs async invocations and readback).

The tests make use of
[Asyncify](https://emscripten.org/docs/porting/asyncify.html) to
implement synchronous versions of `wgpuInstanceRequestAdapter` and
`wgpuAdapterRequestDevice`, while the sample preinitializes a
`WGPUDevice` in JS then passes it in via Emscripten's
`preinitializedWebGPUDevice`/`emscripten_webgpu_get_device()`. There are
a few other ways we can set that up - we'll definitely want to disable
Asyncify for shipping application code.

CMakeLists.txt

@@ -714,10 +714,12 @@ endif()
 if(IREE_TARGET_BACKEND_WEBGPU OR
   (IREE_HAL_DRIVER_WEBGPU AND ("${IREE_HAL_WEBGPU_PLATFORM}" STREQUAL "dawn")))
   # NOTE: these all need to be synchronized with Dawn.
-  set(IREE_SPIRV_TOOLS_TAG  "37d2396cabe56b29d37551ea55d0d745d5748ded")
+  # set(IREE_SPIRV_TOOLS_TAG  "37d2396cabe56b29d37551ea55d0d745d5748ded")
+  set(IREE_SPIRV_TOOLS_TAG  "ff07cfd86fa229525659f6b81058b3171a67bef1")  # 2021-12-10
   # https://dawn.googlesource.com/dawn/+/refs/heads/main/DEPS
   # https://chromium.googlesource.com/vulkan-deps/+/main/DEPS
-  set(IREE_TINT_TAG         "5aca651c524caca8ac85962d1dff378238671eba")
+  # set(IREE_TINT_TAG         "5aca651c524caca8ac85962d1dff378238671eba")
+  set(IREE_TINT_TAG         "188b1fb8f5be52299fb7fbc6db17dbb0c07dbb7e")  # 2021-12-16
   set(IREE_DAWN_ABSEIL_TAG  "789af048b388657987c59d4da406859034fe310f")
   set(IREE_DAWN_GLFW_TAG    "62e175ef9fae75335575964c845a302447c012c7")
   set(IREE_DAWN_TAG         "2f1b0dc47d8316f3db6e5b9a55b873cd3079cea5")

build_tools/cmake/iree_copts.cmake

@@ -355,12 +355,17 @@ iree_select_compiler_opts(IREE_DEFAULT_LINKOPTS
     "-natvis:${IREE_ROOT_DIR}/runtime/iree.natvis"
 )
 
-# Add to LINKOPTS on a binary to configure it for X/Wayland/Windows/etc
-# depending on the target cross-compilation platform.
-if(${CMAKE_SYSTEM_NAME} STREQUAL "Windows")
-  set(IREE_TARGET_GUI_LINKOPTS "-SUBSYSTEM:WINDOWS")
-else()
-  set(IREE_TARGET_GUI_LINKOPTS "")
+if(EMSCRIPTEN AND IREE_HAL_DRIVER_WEBGPU)
+  iree_select_compiler_opts(IREE_DEFAULT_LINKOPTS
+    ALL
+      "-sUSE_WEBGPU"
+      # Hack: Used to create sync versions of requestAdapter and requestDevice
+      # TODO(scotttodd): Only set for test binaries, avoid sync code in apps
+      #   this doesn't _break_ apps that don't use the sync functions, but it
+      #   does bloat their binary size (and each Emscripten flag comes with
+      #   some risk of breaking compatibility with other features)
+      "-sASYNCIFY"
+  )
 endif()
 
 #-------------------------------------------------------------------------------

experimental/web/sample_dynamic/CMakeLists.txt

@@ -37,7 +37,7 @@ target_link_libraries(${_NAME}
 target_link_options(${_NAME} PRIVATE
   # https://emscripten.org/docs/porting/connecting_cpp_and_javascript/Interacting-with-code.html#interacting-with-code-ccall-cwrap
   "-sEXPORTED_FUNCTIONS=['_setup_sample', '_cleanup_sample', '_load_program', '_inspect_program', '_unload_program', '_call_function']"
-  "-sEXPORTED_RUNTIME_METHODS=['ccall','cwrap']"
+  "-sEXPORTED_RUNTIME_METHODS=['ccall','cwrap','UTF8ToString']"
   #
   "-sASSERTIONS=1"
   #

experimental/web/sample_dynamic/build_sample.sh

@@ -76,6 +76,7 @@ emcmake "${CMAKE_BIN?}" -G Ninja .. \
   -DIREE_BUILD_EXPERIMENTAL_WEB_SAMPLES=ON \
   -DIREE_HAL_DRIVER_DEFAULTS=OFF \
   -DIREE_HAL_DRIVER_LOCAL_SYNC=ON \
+  -DIREE_HAL_DRIVER_WEBGPU=OFF \
   -DIREE_BUILD_COMPILER=OFF \
   -DIREE_BUILD_TESTS=OFF
 

experimental/web/sample_dynamic/iree_worker.js

@@ -32,7 +32,7 @@ var Module = {
     wasmInspectProgramFn = Module.cwrap('inspect_program', null, ['number']);
     wasmUnloadProgramFn = Module.cwrap('unload_program', null, ['number']);
     wasmCallFunctionFn = Module.cwrap(
-        'call_function', 'string', ['number', 'string', 'string', 'number']);
+        'call_function', 'number', ['number', 'string', 'string', 'number']);
 
     sampleState = wasmSetupSampleFn();
 
@@ -120,8 +120,11 @@ function callFunction(id, functionParams) {
     return;
   }
 
-  const returnValue =
+  // Receive as a pointer, convert, then free. This avoids a memory leak, see
+  // https://github.com/emscripten-core/emscripten/issues/6484
+  const returnValuePtr =
       wasmCallFunctionFn(programState, functionName, inputsJoined, iterations);
+  const returnValue = Module.UTF8ToString(returnValuePtr);
 
   if (returnValue === '') {
     postMessage({
@@ -130,16 +133,12 @@ function callFunction(id, functionParams) {
       'error': 'Wasm module error, check console for details',
     });
   } else {
+    Module._free(returnValuePtr);
     postMessage({
       'messageType': 'callResult',
       'id': id,
       'payload': JSON.parse(returnValue),
     });
-    // TODO(scotttodd): free char* buffer? Or does Emscripten handle that?
-    // Could refactor to
-    //   1) return void*
-    //   2) convert to String manually using UTF8ToString(pointer)
-    //   3) Module._free(pointer)
   }
 }
 

experimental/web/sample_dynamic/main.c

@@ -6,6 +6,7 @@
 
 #include <stdint.h>
 #include <stdio.h>
+#include <string.h>
 
 #include "iree/base/api.h"
 #include "iree/hal/api.h"
@@ -436,5 +437,7 @@ const char* call_function(iree_program_state_t* program_state,
   }
 
   // Note: this leaks the buffer. It's up to the caller to free it after use.
-  return iree_string_builder_buffer(&outputs_builder);
+  char* outputs = strdup(iree_string_builder_buffer(&outputs_builder));
+  iree_string_builder_deinitialize(&outputs_builder);
+  return outputs;
 }

experimental/web/sample_webgpu/CMakeLists.txt

@@ -0,0 +1,46 @@
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+if(NOT EMSCRIPTEN)
+  return()
+endif()
+
+set(_NAME "iree_experimental_web_sample_webgpu")
+add_executable(${_NAME} "")
+target_sources(${_NAME}
+  PRIVATE
+    main.c
+    device_webgpu.c
+)
+set_target_properties(${_NAME} PROPERTIES OUTPUT_NAME "web-sample-webgpu")
+
+target_compile_options(${_NAME} PRIVATE ${IREE_DEFAULT_COPTS})
+
+# Note: we have to be very careful about dependencies here.
+#
+# The general purpose libraries link in multiple executable loaders and HAL
+# drivers/devices, which include code not compatible with Emscripten.
+target_link_libraries(${_NAME}
+  iree_runtime_runtime
+  iree_hal_drivers_webgpu_webgpu
+  iree_hal_drivers_webgpu_platform_emscripten_emscripten
+)
+
+target_link_options(${_NAME} PRIVATE
+  # https://emscripten.org/docs/porting/connecting_cpp_and_javascript/Interacting-with-code.html#interacting-with-code-ccall-cwrap
+  "-sEXPORTED_FUNCTIONS=['_setup_sample', '_cleanup_sample', '_load_program', '_inspect_program', '_unload_program', '_call_function']"
+  "-sEXPORTED_RUNTIME_METHODS=['ccall','cwrap','UTF8ToString']"
+  #
+  "-sASSERTIONS=1"
+  #
+  # For https://emscripten.org/docs/debugging/Sanitizers.html#address-sanitizer
+  # "-fsanitize=address"
+  # "-sALLOW_MEMORY_GROWTH"
+  #
+  # https://developer.chrome.com/blog/wasm-debugging-2020/
+  "-g"
+  "-gseparate-dwarf"
+)

experimental/web/sample_webgpu/README.md

@@ -0,0 +1,42 @@
+# WebGPU Sample
+
+This experimental sample demonstrates one way to target the web platform with
+IREE, using WebGPU. The output artifact is a web page that loads separately
+provided IREE `.vmfb` (compiled ML model) files and allows for calling
+functions on them.
+
+## Quickstart
+
+**Note**: you will need a WebGPU-compatible browser. Chrome Canary with the
+`#enable-unsafe-webgpu` flag is a good choice (you may need the flag or an
+origin trial token for `localhost`).
+
+1. Install IREE's host tools (e.g. by building the `install` target with CMake)
+2. Install the Emscripten SDK by
+   [following these directions](https://emscripten.org/docs/getting_started/downloads.html)
+3. Initialize your Emscripten environment (e.g. run `emsdk_env.bat`)
+4. From this directory, run `bash ./build_sample.sh [path to install] && bash ./serve_sample.sh`
+5. Open the localhost address linked in the script output
+
+To rebuild most parts of the sample (C runtime, sample HTML, CMake config,
+etc.), just `control + C` to stop the local webserver and rerun the script.
+
+## How it works
+
+[Emscripten](https://emscripten.org/) is used (via the `emcmake` CMake wrapper)
+to compile the runtime into WebAssembly and JavaScript files.
+
+Any supported IREE program, such as
+[simple_abs.mlir](../../../samples/models/simple_abs.mlir), is compiled using
+the WebGPU compiler target. This generates WGSL shader code and IREE VM
+bytecode, which the IREE runtime is able to load and run using the browser's
+WebGPU APIs.
+
+### Asynchronous API
+
+[`iree_api_webgpu.js`](./iree_api_webgpu.js)
+
+* exposes a Promise-based API to the hosting application in
+  [`index.html`](./index.html)
+* preinitializes a WebGPU adapter and device
+* includes Emscripten's JS code and instantiates the WebAssembly module

experimental/web/sample_webgpu/build_sample.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Builds the sample, running host tools, Emscripten, and CMake as needed.
+#
+# Prerequisites:
+#   * Environment must be configured for Emscripten
+#   * Host tools must be built (default at IREE_SOURCE_DIR/build-host/install).
+#     The build_tools/cmake/build_host_tools.sh script can do this for you.
+#
+# Usage:
+#   build_sample.sh (optional install path) && serve_sample.sh
+
+set -e
+
+###############################################################################
+# Setup and checking for dependencies                                         #
+###############################################################################
+
+if ! command -v emcmake &> /dev/null
+then
+  echo "'emcmake' not found, setup environment according to https://emscripten.org/docs/getting_started/downloads.html"
+  exit 1
+fi
+
+CMAKE_BIN=${CMAKE_BIN:-$(which cmake)}
+ROOT_DIR=$(git rev-parse --show-toplevel)
+SOURCE_DIR=${ROOT_DIR}/experimental/web/sample_webgpu
+
+BUILD_DIR=${ROOT_DIR?}/build-emscripten
+mkdir -p ${BUILD_DIR}
+
+BINARY_DIR=${BUILD_DIR}/experimental/web/sample_webgpu
+mkdir -p ${BINARY_DIR}
+
+INSTALL_ROOT="${1:-${ROOT_DIR}/build-host/install}"
+
+###############################################################################
+# Compile from .mlir input to portable .vmfb file using host tools            #
+###############################################################################
+
+COMPILE_TOOL="${INSTALL_ROOT?}/bin/iree-compile"
+
+compile_sample() {
+  echo "  Compiling '$1' sample for WebGPU..."
+  ${COMPILE_TOOL?} $2 \
+    --iree-input-type=mhlo \
+    --iree-hal-target-backends=webgpu \
+    --o ${BINARY_DIR}/$1_webgpu.vmfb
+}
+
+echo "=== Compiling sample MLIR files to VM FlatBuffer outputs (.vmfb) ==="
+compile_sample "simple_abs"     "${ROOT_DIR?}/samples/models/simple_abs.mlir"
+# compile_sample "fullyconnected" "${ROOT_DIR?}/tests/e2e/models/fullyconnected.mlir"
+# compile_sample "collatz"        "${ROOT_DIR?}/tests/e2e/models/collatz.mlir"
+
+###############################################################################
+# Build the web artifacts using Emscripten                                    #
+###############################################################################
+
+echo "=== Building web artifacts using Emscripten ==="
+
+pushd ${BUILD_DIR}
+
+# Configure using Emscripten's CMake wrapper, then build.
+# Note: The sample creates a device directly, so no drivers are required.
+emcmake "${CMAKE_BIN?}" -G Ninja .. \
+  -DCMAKE_BUILD_TYPE=RelWithDebInfo \
+  -DIREE_HOST_BINARY_ROOT=${INSTALL_ROOT} \
+  -DIREE_BUILD_EXPERIMENTAL_WEB_SAMPLES=ON \
+  -DIREE_ENABLE_THREADING=OFF \
+  -DIREE_HAL_DRIVER_DEFAULTS=OFF \
+  -DIREE_HAL_DRIVER_LOCAL_SYNC=OFF \
+  -DIREE_HAL_DRIVER_LOCAL_TASK=OFF \
+  -DIREE_HAL_DRIVER_WEBGPU=ON \
+  -DIREE_HAL_WEBGPU_PLATFORM=emscripten \
+  -DIREE_BUILD_COMPILER=OFF \
+  -DIREE_BUILD_TESTS=OFF
+
+"${CMAKE_BIN?}" --build . --target \
+  iree_experimental_web_sample_webgpu
+
+popd
+
+echo "=== Copying static files (.html, .js) to the build directory ==="
+
+cp ${SOURCE_DIR?}/index.html ${BINARY_DIR}
+cp ${ROOT_DIR?}/docs/website/overrides/ghost.svg ${BINARY_DIR}
+cp ${SOURCE_DIR?}/iree_api_webgpu.js ${BINARY_DIR}

experimental/web/sample_webgpu/device_webgpu.c

@@ -0,0 +1,30 @@
+// Copyright 2022 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <emscripten/html5.h>
+#include <emscripten/html5_webgpu.h>
+
+#include "iree/base/api.h"
+#include "iree/hal/api.h"
+#include "iree/hal/drivers/webgpu/api.h"
+#include "iree/hal/drivers/webgpu/platform/webgpu.h"
+
+iree_status_t create_device(iree_allocator_t host_allocator,
+                            iree_hal_device_t** out_device) {
+  WGPUDevice wgpu_device = emscripten_webgpu_get_device();
+  if (!wgpu_device) {
+    return iree_make_status(
+        IREE_STATUS_UNAVAILABLE,
+        "emscripten_webgpu_get_device() failed to return a WGPUDevice");
+  }
+
+  iree_hal_webgpu_device_options_t default_options;
+  iree_hal_webgpu_device_options_initialize(&default_options);
+
+  return iree_hal_webgpu_wrap_device(IREE_SV("webgpu-emscripten"),
+                                     &default_options, wgpu_device,
+                                     host_allocator, out_device);
+}