Rework CMake interface for Dawn/Node bindings (halide#7422)

AOT pipelines that rely on Dawn/WebGPU now depend on a new
Halide_WebGPU find-module. This module honors the make-ish
HL_WEBGPU_NATIVE_LIB variable as a means of initializing the
Halide_WebGPU_NATIVE_LIB cache variable. This is automatically handled
by add_halide_generator and add_halide_runtime and is available to
downstreams.

The JIT tests no longer read the HL_WEBGPU_NODE_BINDINGS environment
variable during the CMake configure or build phase. Instead, a test
launcher reads it at CTest runtime.

Co-authored-by: Alex Reinking <quic_areinkin@quicinc.com>

README_webgpu.md

@@ -18,6 +18,8 @@ The following is a non-comprehensive list of known limitations:
     * 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.
+-   You must use CMake/CTest to build/test Halide for WebGPU; using the Makefile
+    is not supported for WebGPU testing (and probably never will be).
 
 In addition to these functional limitations, the performance of the WebGPU
 backend has not yet been evaluated, and so optimizations in the runtime or
@@ -35,8 +37,9 @@ When invoking `emcc` to link Halide-generated objects, include these flags:
 
 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.
+You must set an environment variable named `HL_WEBGPU_NODE_BINDINGS` that
+has an absolute path to the bindings to run these tests, e.g. `HL_WEBGPU_NODE_BINDINGS=/path/to/dawn.node`.
+
 See [below](#setting-up-dawn) for instructions on building the Dawn Node.js
 bindings.
 
@@ -51,9 +54,14 @@ For testing purposes, Halide can also target native WebGPU libraries, such as
 This is currently the only path that can run the JIT correctness tests.
 See [below](#setting-up-dawn) for instructions on building Dawn.
 
-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.
+When targeting WebGPU with a native target, Halide defaults to looking for a
+build of Dawn (with several common names and suffixes); you can override this
+by setting the `HL_WEBGPU_NATIVE_LIB` environment variable to the absolute path
+to the library you want.
+
+Note that it is explicitly legal to define both `HL_WEBGPU_NATIVE_LIB` and
+`HL_WEBGPU_NODE_BINDINGS` at the same time; the correct executable environment
+will be selected based on the Halide target specified.
 
 Note that it is explicitly legal to specify both WEBGPU_NATIVE_LIB and
 WEBGPU_NODE_BINDINGS for the same build; the correct executable environment
@@ -98,5 +106,5 @@ 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.
+These paths can then be used for the `HL_WEBGPU_NODE_BINDINGS` and
+`HL_WEBGPU_NATIVE_LIB` environment variables when using Halide.

cmake/FindHalide_WebGPU.cmake

@@ -0,0 +1,27 @@
+cmake_minimum_required(VERSION 3.22)
+
+# tip: uncomment this line to get better debugging information if find_library() fails
+# set(CMAKE_FIND_DEBUG_MODE TRUE)
+
+if (EXISTS "$ENV{HL_WEBGPU_NATIVE_LIB}")
+  set(Halide_WebGPU_NATIVE_LIB "$ENV{HL_WEBGPU_NATIVE_LIB}"
+      CACHE FILEPATH "")
+endif ()
+
+find_library(Halide_WebGPU_NATIVE_LIB NAMES webgpu_dawn wgpu)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(
+  Halide_WebGPU
+  REQUIRED_VARS Halide_WebGPU_NATIVE_LIB
+  HANDLE_COMPONENTS
+)
+
+if (Halide_WebGPU_NATIVE_LIB AND NOT TARGET Halide::WebGPU)
+  add_library(Halide::WebGPU UNKNOWN IMPORTED)
+  set_target_properties(
+    Halide::WebGPU
+    PROPERTIES
+    IMPORTED_LOCATION "${Halide_WebGPU_NATIVE_LIB}"
+  )
+endif ()

cmake/HalideGeneratorHelpers.cmake

@@ -723,9 +723,8 @@ function(_Halide_target_link_gpu_libs TARGET VISIBILITY)
     endif ()
 
     if ("${ARGN}" MATCHES "webgpu")
-        if (WEBGPU_NATIVE_LIB)
-            target_link_libraries(${TARGET} INTERFACE ${WEBGPU_NATIVE_LIB})
-        endif ()
+        find_package(Halide_WebGPU REQUIRED)
+        target_link_libraries(${TARGET} ${VISIBILITY} Halide::WebGPU)
     endif ()
 endfunction()
 

dependencies/wasm/CMakeLists.txt

@@ -153,18 +153,9 @@ 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} "-e" "${script}")
+                    COMMAND ${NODE_JS_EXECUTABLE} "${Halide_SOURCE_DIR}/tools/launch_wasm_test.js" "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.js" "${Halide_TARGET}")
 endfunction()
 
 function(find_node_js)

packaging/CMakeLists.txt

@@ -213,6 +213,7 @@ install(FILES
         ${CMAKE_CURRENT_BINARY_DIR}/HalideHelpersConfig.cmake
         ${CMAKE_CURRENT_BINARY_DIR}/HalideHelpersConfigVersion.cmake
         ${Halide_SOURCE_DIR}/cmake/HalideGeneratorHelpers.cmake
+        ${Halide_SOURCE_DIR}/cmake/FindHalide_WebGPU.cmake
         ${Halide_SOURCE_DIR}/cmake/HalideTargetHelpers.cmake
         ${Halide_SOURCE_DIR}/cmake/TargetExportScript.cmake
         DESTINATION ${Halide_INSTALL_HELPERSDIR}

src/CMakeLists.txt

@@ -594,13 +594,6 @@ 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()
 
 

src/CodeGen_WebGPU_Dev.cpp

@@ -345,7 +345,7 @@ void CodeGen_WebGPU_Dev::CodeGen_WGSL::add_kernel(
                 user_warning
                     << "buffers of small integer types are currently emulated "
                     << "using atomics in the WebGPU backend, and accesses to "
-                    << "them will be slow";
+                    << "them will be slow.";
                 buffers_with_emulated_accesses.insert(arg.name);
                 type_decl = "atomic<u32>";
             } else {

src/JITModule.cpp

@@ -138,35 +138,43 @@ void load_vulkan() {
 
 void load_webgpu() {
     debug(1) << "Looking for a native WebGPU implementation...\n";
-    const char *libnames[] = {
-#ifdef WEBGPU_NATIVE_LIB
-        // Specified via CMake.
-        WEBGPU_NATIVE_LIB,
-#endif
-        // Dawn (Chromium).
-        "libwebgpu_dawn.so",
-        "libwebgpu_dawn.dylib",
-        "webgpu_dawn.dll",
-
-        // wgpu (Firefox).
-        "libwgpu.so",
-        "libwgpu.dylib",
-        "wgpu.dll",
-    };
-    string error;
-    for (const char *libname : libnames) {
+
+    const auto try_load = [](const char *libname) -> string {
         debug(1) << "Trying " << libname << "... ";
-        error.clear();
+        string error;
         llvm::sys::DynamicLibrary::LoadLibraryPermanently(libname, &error);
-        if (error.empty()) {
-            debug(1) << "found!\n";
-            break;
-        } else {
-            debug(1) << "not found.\n";
+        debug(1) << (error.empty() ? "found!\n" : "not found.\n");
+        return error;
+    };
+
+    string error;
+
+    auto env_libname = get_env_variable("HL_WEBGPU_NATIVE_LIB");
+    if (!env_libname.empty()) {
+        error = try_load(env_libname.c_str());
+    }
+    if (!error.empty()) {
+        const char *libnames[] = {
+            // Dawn (Chromium).
+            "libwebgpu_dawn.so",
+            "libwebgpu_dawn.dylib",
+            "webgpu_dawn.dll",
+
+            // wgpu (Firefox).
+            "libwgpu.so",
+            "libwgpu.dylib",
+            "wgpu.dll",
+        };
+
+        for (const char *libname : libnames) {
+            error = try_load(libname);
+            if (error.empty()) {
+                break;
+            }
         }
     }
-    user_assert(error.empty()) << "Could not find a native WebGPU library: "
-                               << error << "\n";
+    user_assert(error.empty()) << "Could not find a native WebGPU library: " << error << "\n"
+                               << "(Try setting the env var HL_WEBGPU_NATIVE_LIB to an explicit path to fix this.)\n";
 }
 
 }  // namespace