Add initial support for WebGPU

.clang-format-ignore

@@ -12,4 +12,6 @@
 ./tutorial
 # hexagon_remote/bin/src is also special
 ./src/runtime/hexagon_remote/bin/src
+# mini_webgpu.h is copied from upstream with some local mods
+./src/runtime/mini_webgpu.h
 ./dependencies/spirv

Makefile

@@ -124,6 +124,7 @@ WITH_OPENCL ?= not-empty
 WITH_METAL ?= not-empty
 WITH_OPENGLCOMPUTE ?= not-empty
 WITH_D3D12 ?= not-empty
+WITH_WEBGPU ?= not-empty
 WITH_INTROSPECTION ?= not-empty
 WITH_EXCEPTIONS ?=
 WITH_LLVM_INSIDE_SHARED_LIBHALIDE ?= not-empty
@@ -161,6 +162,8 @@ OPENGLCOMPUTE_CXX_FLAGS=$(if $(WITH_OPENGLCOMPUTE), -DWITH_OPENGLCOMPUTE, )
 D3D12_CXX_FLAGS=$(if $(WITH_D3D12), -DWITH_D3D12, )
 D3D12_LLVM_CONFIG_LIB=$(if $(WITH_D3D12), , )
 
+WEBGPU_CXX_FLAGS=$(if $(WITH_WEBGPU), -DWITH_WEBGPU, )
+
 AARCH64_CXX_FLAGS=$(if $(WITH_AARCH64), -DWITH_AARCH64, )
 AARCH64_LLVM_CONFIG_LIB=$(if $(WITH_AARCH64), aarch64, )
 
@@ -206,6 +209,7 @@ CXX_FLAGS += $(OPENCL_CXX_FLAGS)
 CXX_FLAGS += $(METAL_CXX_FLAGS)
 CXX_FLAGS += $(OPENGLCOMPUTE_CXX_FLAGS)
 CXX_FLAGS += $(D3D12_CXX_FLAGS)
+CXX_FLAGS += $(WEBGPU_CXX_FLAGS)
 CXX_FLAGS += $(POWERPC_CXX_FLAGS)
 CXX_FLAGS += $(INTROSPECTION_CXX_FLAGS)
 CXX_FLAGS += $(EXCEPTIONS_CXX_FLAGS)
@@ -436,6 +440,7 @@ SOURCE_FILES = \
   CodeGen_PyTorch.cpp \
   CodeGen_RISCV.cpp \
   CodeGen_WebAssembly.cpp \
+  CodeGen_WebGPU_Dev.cpp \
   CodeGen_X86.cpp \
   CompilerLogger.cpp \
   CPlusPlusMangle.cpp \
@@ -613,6 +618,7 @@ HEADER_FILES = \
   CodeGen_PTX_Dev.h \
   CodeGen_PyTorch.h \
   CodeGen_Targets.h \
+  CodeGen_WebGPU_Dev.h \
   CompilerLogger.h \
   ConciseCasts.h \
   CPlusPlusMangle.h \
@@ -824,6 +830,7 @@ RUNTIME_CPP_COMPONENTS = \
   trace_helper \
   tracing \
   wasm_cpu_features \
+  webgpu \
   windows_clock \
   windows_cuda \
   windows_d3d12compute_arm \
@@ -864,6 +871,7 @@ RUNTIME_EXPORTED_INCLUDES = $(INCLUDE_DIR)/HalideRuntime.h \
                             $(INCLUDE_DIR)/HalideRuntimeOpenGLCompute.h \
                             $(INCLUDE_DIR)/HalideRuntimeMetal.h	\
                             $(INCLUDE_DIR)/HalideRuntimeQurt.h \
+                            $(INCLUDE_DIR)/HalideRuntimeWebGPU.h \
                             $(INCLUDE_DIR)/HalideBuffer.h \
                             $(INCLUDE_DIR)/HalidePyTorchHelpers.h \
                             $(INCLUDE_DIR)/HalidePyTorchCudaHelpers.h

README_webgpu.md

@@ -0,0 +1,102 @@
+# WebGPU support for Halide
+
+Halide has work-in-progress support for generating and running WebGPU shaders.
+This can be used in conjunction with the WebAssembly backend to bring
+GPU-accelerated Halide pipelines to the web.
+
+As the first version of the WebGPU standard is itself still being developed,
+Halide's support has some limitations and may only work with certain browsers
+and versions of Emscripten.
+
+## Known limitations
+
+The following is a non-comprehensive list of known limitations:
+
+-   Only 32-bit integers and floats have efficient support.
+    * 8-bit and 16-bit integers are implemented using emulation. Future
+      extensions to WGSL will allow them to be implemented more efficiently.
+    * 64-bit integers and floats will likely remain unsupported until WGSL gains
+      extensions to support them.
+-   Wrapping native device buffer handles is not yet implemented.
+
+In addition to these functional limitations, the performance of the WebGPU
+backend has not yet been evaluated, and so optimizations in the runtime or
+device codegen may be required before it becomes profitable to use.
+
+## Running with WebAssembly via Emscripten: `HL_TARGET=wasm-32-wasmrt-webgpu`
+
+> _Tested with top-of-tree Emscripten as of 2023-02-23, against Chrome v113._
+
+Halide can generate WebGPU code that can be integrated with WASM code using
+Emscripten.
+Halide must currently be built *without* the `WEBGPU_NATIVE_LIB` flag when
+targeting Emscripten.
+
+When invoking `emcc` to link Halide-generated objects, include these flags:
+`-s USE_WEBGPU=1 -s ASYNCIFY`.
+
+Tests that use AOT compilation can be run using a native WebGPU implementation
+that has Node.js bindings, such as [Dawn](dawn.googlesource.com/dawn/).
+When configuring Halide, use `-DWEBGPU_NODE_BINDINGS=/path/to/dawn.node` to
+enable these tests.
+See [below](#setting-up-dawn) for instructions on building the Dawn Node.js
+bindings.
+
+JIT compilation is not supported when using WebGPU with WASM.
+
+## Running natively: `HL_TARGET=host-webgpu`
+
+> _Tested with top-of-tree Dawn as of 2023-02-28._
+
+For testing purposes, Halide can also target native WebGPU libraries, such as
+[Dawn](dawn.googlesource.com/dawn/) or [wgpu](github.com/gfx-rs/wgpu).
+This is currently the only path that can run the JIT correctness tests.
+See [below](#setting-up-dawn) for instructions on building Dawn.
+
+Due to differences between the APIs implemented by native WebGPU libraries and
+Emscripten, this currently requires a separate build of Halide.
+Pass `-DWEBGPU_NATIVE_LIB=/path/to/native/library.{so,dylib.dll}` to CMake when
+configuring Halide to enable this path, which will automatically use this
+library for the AOT and JIT tests.
+
+## Setting up Dawn
+
+Building Dawn's Node.js bindings currently requires using CMake.
+
+First, [install `depot_tools`](commondatastorage.googleapis.com/chrome-infra-docs/flat/depot_tools/docs/html/depot_tools_tutorial.html#_setting_up) and add it to the
+`PATH` environment variable.
+
+Next, get Dawn and its dependencies:
+
+    # Clone the repo
+    git clone https://dawn.googlesource.com/dawn
+    cd dawn
+
+    # Bootstrap the gclient configuration with Node.js bindings enabled
+    cp scripts/standalone-with-node.gclient .gclient
+
+    # Fetch external dependencies and toolchains with gclient
+    gclient sync
+
+    # Other dependencies that must be installed manually:
+    # - golang
+
+Finally, build Dawn, enabling both the Node.js bindings and shared libraries:
+
+    mkdir -p <build_dir>
+    cd <build_dir>
+
+    cmake <dawn_root_dir> -G Ninja \
+        -DCMAKE_BUILD_TYPE=Release \
+        -DDAWN_BUILD_NODE_BINDINGS=1 \
+        -DDAWN_ENABLE_PIC=1 \
+        -DBUILD_SHARED_LIBS=ON
+
+    ninja dawn.node webgpu_dawn
+
+This will produce the following artifacts:
+- Node.js bindings: `<build_dir>/dawn.node`
+- Native library: `<build_dir>/src/dawn/native/libwebgpu_dawn.{so,dylib,dll}`
+
+These paths can then be used for the `-DWEBGPU_NODE_BINDINGS` and
+`-DWEBGPU_NATIVE_LIB` CMake options when configuring Halide.

cmake/HalideGeneratorHelpers.cmake

@@ -716,6 +716,12 @@ function(_Halide_target_link_gpu_libs TARGET VISIBILITY)
         find_library(METAL_LIBRARY Metal REQUIRED)
         target_link_libraries(${TARGET} ${VISIBILITY} "${FOUNDATION_LIBRARY}" "${METAL_LIBRARY}")
     endif ()
+
+    if ("${ARGN}" MATCHES "webgpu")
+        if (WEBGPU_NATIVE_LIB)
+            target_link_libraries(${TARGET} PRIVATE ${WEBGPU_NATIVE_LIB})
+        endif ()
+    endif ()
 endfunction()
 
 ##

dependencies/wasm/CMakeLists.txt

@@ -70,7 +70,7 @@ endif()
 function(add_wasm_executable TARGET)
     set(options)
     set(oneValueArgs)
-    set(multiValueArgs SRCS DEPS INCLUDES ENABLE_IF)
+    set(multiValueArgs SRCS DEPS INCLUDES OPTIONS ENABLE_IF)
     cmake_parse_arguments(args "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
 
     if (args_ENABLE_IF AND NOT (${args_ENABLE_IF}))
@@ -107,8 +107,23 @@ function(add_wasm_executable TARGET)
         -s ALLOW_MEMORY_GROWTH=1
         -s ENVIRONMENT=node
         -s STACK_SIZE=98304
+        ${args_OPTIONS}
     )
 
+    if ("${Halide_TARGET}" MATCHES "webgpu")
+        set(EMCC_FLAGS
+            ${EMCC_FLAGS}
+            -s USE_WEBGPU=1
+            -s ASYNCIFY
+        )
+    endif ()
+
+    # TODO: Remove this when Emscripten and Dawn implement the same APIs.
+    # See https://github.com/halide/Halide/issues/7248
+    if (WEBGPU_NATIVE_LIB)
+        message(FATAL_ERROR "Cannot generate WASM executables when targeting native WebGPU implementations.")
+    endif ()
+
     set(SRCS)
     foreach (S IN LISTS args_SRCS)
         list(APPEND SRCS "${CMAKE_CURRENT_SOURCE_DIR}/${S}")
@@ -144,9 +159,18 @@ function(add_wasm_halide_test TARGET)
         return()
     endif ()
 
+    set(script "require('./${TARGET}.js')")
+    if (WEBGPU_NODE_BINDINGS)
+        set(script "\
+            const provider = require('${WEBGPU_NODE_BINDINGS}')\n\
+            const gpu = provider.create([])\n\
+            const navigator = { gpu: gpu }\n
+            ${script}")
+    endif ()
+
     add_halide_test("${TARGET}"
                     GROUPS ${args_GROUPS}
-                    COMMAND ${NODE_JS_EXECUTABLE} "${TARGET}.js")
+                    COMMAND ${NODE_JS_EXECUTABLE} "-e" "${script}")
 endfunction()
 
 function(find_node_js)

python_bindings/test/correctness/CMakeLists.txt

@@ -1,3 +1,8 @@
+if (TARGET_WEBASSEMBLY AND Halide_TARGET MATCHES "webgpu")
+    message(WARNING "Python correctness tests are not supported with WebGPU.")
+    return()
+endif ()
+
 add_library(the_sort_function MODULE the_sort_function.c)
 target_link_libraries(the_sort_function PRIVATE Halide::Runtime)
 

python_bindings/tutorial/CMakeLists.txt

@@ -1,3 +1,8 @@
+if (TARGET_WEBASSEMBLY AND Halide_TARGET MATCHES "webgpu")
+    message(WARNING "Python tutorials are not supported with WebGPU.")
+    return()
+endif ()
+
 set(tests
     lesson_01_basics.py
     lesson_02_input_image.py
@@ -58,6 +63,11 @@ else ()
 
     target_link_libraries(lesson_10_halide PRIVATE Halide::Runtime)
 
+    # Undocumented function in HalideGeneratorHelpers. Do not call in external code.
+    # Users of the AOT functions (as opposed to Generators) should link to the relevant
+    # GPU libraries manually.
+    _Halide_target_link_gpu_libs(lesson_10_halide PRIVATE ${Halide_TARGET})
+
     # The fixture "py_lesson_10" orchestrates running the generator part of the lesson first, then the build for the
     # library, and finally runs python_tutorial_lesson_10_aot_compilation_run. The ..._compile test invokes CMake on
     # the current build for the above library.

src/CMakeLists.txt

@@ -39,6 +39,7 @@ set(HEADER_FILES
     CodeGen_PTX_Dev.h
     CodeGen_PyTorch.h
     CodeGen_Targets.h
+    CodeGen_WebGPU_Dev.h
     CompilerLogger.h
     ConciseCasts.h
     CPlusPlusMangle.h
@@ -206,6 +207,7 @@ set(SOURCE_FILES
     CodeGen_PyTorch.cpp
     CodeGen_RISCV.cpp
     CodeGen_WebAssembly.cpp
+    CodeGen_WebGPU_Dev.cpp
     CodeGen_X86.cpp
     CompilerLogger.cpp
     CPlusPlusMangle.cpp
@@ -527,6 +529,18 @@ if (TARGET_OPENGLCOMPUTE)
     target_compile_definitions(Halide PRIVATE WITH_OPENGLCOMPUTE)
 endif ()
 
+option(TARGET_WEBGPU "Include WebGPU target" ON)
+if (TARGET_WEBGPU)
+    target_compile_definitions(Halide PRIVATE WITH_WEBGPU)
+    set(WEBGPU_NATIVE_LIB "" CACHE STRING
+        "WebGPU native library to link against")
+    if (WEBGPU_NATIVE_LIB)
+        set_property(SOURCE JITModule.cpp PROPERTY COMPILE_DEFINITIONS
+                     WEBGPU_NATIVE_LIB=\"${WEBGPU_NATIVE_LIB}\"
+                     APPEND_STRING)
+    endif ()
+endif()
+
 if (TARGET_SPIRV)
     # Our vendored SPIRV headers are only used internally; users do not need
     # them installed.

src/CodeGen_C.cpp

@@ -33,6 +33,7 @@ extern "C" unsigned char halide_internal_runtime_header_HalideRuntimeOpenCL_h[];
 extern "C" unsigned char halide_internal_runtime_header_HalideRuntimeOpenGLCompute_h[];
 extern "C" unsigned char halide_internal_runtime_header_HalideRuntimeQurt_h[];
 extern "C" unsigned char halide_internal_runtime_header_HalideRuntimeD3D12Compute_h[];
+extern "C" unsigned char halide_internal_runtime_header_HalideRuntimeWebGPU_h[];
 
 namespace {
 
@@ -503,6 +504,9 @@ CodeGen_C::~CodeGen_C() {
             if (target.has_feature(Target::D3D12Compute)) {
                 stream << halide_internal_runtime_header_HalideRuntimeD3D12Compute_h << "\n";
             }
+            if (target.has_feature(Target::WebGPU)) {
+                stream << halide_internal_runtime_header_HalideRuntimeWebGPU_h << "\n";
+            }
         }
         stream << "#endif\n";
     }

src/CodeGen_GPU_Dev.cpp

@@ -181,6 +181,8 @@ void CodeGen_GPU_C::visit(const Shuffle *op) {
         std::string storage_name = unique_name('_');
         if (vector_declaration_style == VectorDeclarationStyle::OpenCLSyntax) {
             rhs << "(" << print_type(op->type) << ")(";
+        } else if (vector_declaration_style == VectorDeclarationStyle::WGSLSyntax) {
+            rhs << print_type(op->type) << "(";
         } else {
             rhs << "{";
         }
@@ -192,6 +194,8 @@ void CodeGen_GPU_C::visit(const Shuffle *op) {
         }
         if (vector_declaration_style == VectorDeclarationStyle::OpenCLSyntax) {
             rhs << ")";
+        } else if (vector_declaration_style == VectorDeclarationStyle::WGSLSyntax) {
+            rhs << ")";
         } else {
             rhs << "}";
         }

src/CodeGen_GPU_Dev.h

@@ -87,11 +87,12 @@ struct CodeGen_GPU_Dev {
  * GPU backends derive from and specialize this class. */
 class CodeGen_GPU_C : public CodeGen_C {
 public:
-    /** OpenCL uses a different syntax than C for immediate vectors.  This
+    /** OpenCL and WGSL use different syntax than C for immediate vectors. This
     enum defines which style should be used by the backend. */
     enum class VectorDeclarationStyle {
         CLikeSyntax = 0,
-        OpenCLSyntax = 1
+        OpenCLSyntax = 1,
+        WGSLSyntax = 2,
     };
 
     CodeGen_GPU_C(std::ostream &s, Target t)

src/CodeGen_Internal.cpp

@@ -67,6 +67,7 @@ bool function_takes_user_context(const std::string &name) {
         "halide_openglcompute_run",
         "halide_metal_run",
         "halide_d3d12compute_run",
+        "halide_webgpu_run",
         "halide_msan_annotate_buffer_is_initialized_as_destructor",
         "halide_msan_annotate_buffer_is_initialized",
         "halide_msan_annotate_memory_is_initialized",
@@ -91,6 +92,7 @@ bool function_takes_user_context(const std::string &name) {
         "halide_openglcompute_initialize_kernels",
         "halide_metal_initialize_kernels",
         "halide_d3d12compute_initialize_kernels",
+        "halide_webgpu_initialize_kernels",
         "halide_get_gpu_device",
         "_halide_buffer_crop",
         "_halide_buffer_retire_crop_after_extern_stage",