Merge branch 'main' into fix-spotless-module-kt

Author: psiddh

.ci/scripts/setup-webgpu-linux-deps.sh

@@ -5,26 +5,93 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# Vendor Dawn (Tint) + SwiftShader for the WebGPU backend CI WITHOUT hosting a
+# private prebuilt:
+#   * Dawn  : Google's official nightly prebuilt, downloaded directly from
+#             github.com/google/dawn/releases (pinned tag+rev+sha256) -- the same
+#             "fetch a pinned upstream prebuilt" pattern used for other CI deps.
+#   * SwiftShader : built from source at a pinned rev compatible with the Dawn
+#             above (the ossci prebuilt is from 2020, too old for current Dawn). No S3.
+# Dawn (Chrome's WebGPU impl; its WGSL compiler Tint is the spec reference) on
+# SwiftShader gives a headless, deterministic, spec-faithful CLI backend.
+#
+# Exports Dawn_DIR / VK_ICD_FILENAMES / LD_LIBRARY_PATH for the cmake build+run.
+# Local/rig override: set DAWN_PREBUILT_DIR=<dir containing lib64/cmake/Dawn> to
+# skip the Dawn download.
 set -ex
 
-# SwiftShader: software Vulkan adapter for GPU-less CI (LunarG SDK not needed).
-install_swiftshader() {
-  _https_amazon_aws=https://ossci-android.s3.amazonaws.com
-  _swiftshader_archive=swiftshader-abe07b943-prebuilt.tar.gz
-  _swiftshader_dir=/tmp/swiftshader
-  mkdir -p $_swiftshader_dir
+# --- pinned versions (bump rev+sha together when upgrading Dawn) --------------
+DAWN_TAG="${DAWN_TAG:-v20260423.175430}"
+DAWN_REV="${DAWN_REV:-31e25af254ab572c77054edec4946d2244e184dd}"
+DAWN_SHA256="${DAWN_SHA256:-ac76fac090162dc1ecea5ed0f28a557bb8f49efc47faab01886105ace82b7b64}"
+# SwiftShader rev verified compatible with DAWN_REV (the old ossci prebuilt is
+# from 2020 and is incompatible with current Dawn -> no adapter / zero compute).
+SWIFTSHADER_REV="${SWIFTSHADER_REV:-9898204d91d6a60b6a08ad74fe4ac52a6913111b}"
 
-  _tmp_archive="/tmp/${_swiftshader_archive}"
+_dawn_dir="${DAWN_PREBUILT_DIR:-/tmp/dawn-ci}"
+_ss_dir=/tmp/swiftshader
 
-  curl --silent --show-error --location --fail --retry 3 --retry-all-errors \
-    --output "${_tmp_archive}" "$_https_amazon_aws/${_swiftshader_archive}"
+# --- toolchain prereqs --------------------------------------------------------
+# Dawn dlopens the system Vulkan loader at runtime (libvulkan1). And the
+# ubuntu-latest prebuilt is built with a bleeding-edge GCC: it references
+# libstdc++ symbols newer than ubuntu-22.04's default (e.g. _M_replace_cold,
+# GCC 13+), so the static .a won't link against the stock runtime. Pull a current
+# libstdc++ from the ubuntu-toolchain-r PPA when the symbol floor isn't met. All
+# of this is scoped to the WebGPU CI job; newer libstdc++ is backward-compatible.
+if command -v apt-get >/dev/null 2>&1; then
+  _SUDO=""; command -v sudo >/dev/null 2>&1 && _SUDO="sudo"
+  ${_SUDO} apt-get update -y || true
+  ${_SUDO} apt-get install -y libvulkan1 software-properties-common || true
+  if ! strings /usr/lib/x86_64-linux-gnu/libstdc++.so.6 2>/dev/null \
+      | grep -q "GLIBCXX_3.4.32"; then
+    ${_SUDO} add-apt-repository -y ppa:ubuntu-toolchain-r/test || true
+    ${_SUDO} apt-get update -y || true
+    ${_SUDO} apt-get install -y libstdc++6 || true  # newest GCC runtime
+  fi
+fi
 
-  tar -C "${_swiftshader_dir}" -xzf "${_tmp_archive}"
+# The native binaries / pybind lib run INSIDE the CI conda env, whose libstdc++
+# predates GLIBCXX_3.4.32 (the Dawn prebuilt's floor) -- the same wall ssjia hit
+# for the vulkan op tests. Upgrade the conda runtime libstdc++ so the loaded
+# libstdc++.so.6 (conda's, not the system one) satisfies Dawn at run time.
+if command -v conda >/dev/null 2>&1; then
+  conda install -y -c conda-forge "libstdcxx-ng>=14" || true
+fi
+
+# --- Dawn: official prebuilt from GitHub (no S3) ------------------------------
+mkdir -p "${_dawn_dir}"
+if [[ ! -d "${_dawn_dir}/lib64/cmake/Dawn" ]]; then
+  _dawn_tar="/tmp/Dawn-${DAWN_REV}-ubuntu-latest-Release.tar.gz"
+  curl --silent --show-error --location --fail --retry 3 --retry-all-errors \
+    --output "${_dawn_tar}" \
+    "https://github.com/google/dawn/releases/download/${DAWN_TAG}/Dawn-${DAWN_REV}-ubuntu-latest-Release.tar.gz"
+  echo "${DAWN_SHA256}  ${_dawn_tar}" | sha256sum -c -
+  # archive top dir is Dawn-<rev>-ubuntu-latest-Release/{lib64,include,bin}
+  tar -C "${_dawn_dir}" --strip-components=1 -xzf "${_dawn_tar}"
+fi
 
-  export VK_ICD_FILENAMES="${_swiftshader_dir}/swiftshader/build/Linux/vk_swiftshader_icd.json"
-  export LD_LIBRARY_PATH="${_swiftshader_dir}/swiftshader/build/Linux/:${LD_LIBRARY_PATH}"
-  export ETVK_USING_SWIFTSHADER=1
-}
+# --- SwiftShader: build from source at a pinned rev (no S3) -------------------
+# The old ossci prebuilt (swiftshader-abe07b943, 2020) is incompatible with the
+# current Dawn; build a matching modern SwiftShader instead. Self-contained
+# cmake build (vendored LLVM); the ICD lands under build/<OS>/.
+if [[ ! -d "${_ss_dir}/build" ]]; then
+  if [[ ! -d "${_ss_dir}/.git" ]]; then
+    git clone https://github.com/google/swiftshader "${_ss_dir}"
+  fi
+  git -C "${_ss_dir}" checkout "${SWIFTSHADER_REV}"
+  # vk_swiftshader's deps are vendored in-tree; tolerate unreachable
+  # disabled-feature submodules (angle, test-only) failing to fetch.
+  git -C "${_ss_dir}" submodule update --init --recursive || true
+  cmake -S "${_ss_dir}" -B "${_ss_dir}/build" -DCMAKE_BUILD_TYPE=Release \
+    -DSWIFTSHADER_BUILD_TESTS=OFF -DSWIFTSHADER_BUILD_PVR=OFF \
+    -DSWIFTSHADER_BUILD_BENCHMARKS=OFF
+  cmake --build "${_ss_dir}/build" --parallel "$(nproc)" --target vk_swiftshader
+fi
+_ss_icd="$(find "${_ss_dir}/build" -name vk_swiftshader_icd.json 2>/dev/null | head -1)"
+[[ -n "${_ss_icd}" ]] || { echo "ERROR: SwiftShader ICD not found after build" >&2; exit 1; }
 
-install_swiftshader
-bash backends/webgpu/scripts/setup-wgpu-native.sh
+_ss_libdir="$(dirname "${_ss_icd}")"
+export Dawn_DIR="${_dawn_dir}/lib64/cmake/Dawn"
+export VK_ICD_FILENAMES="${_ss_icd}"
+export LD_LIBRARY_PATH="${_ss_libdir}:${LD_LIBRARY_PATH:-}"
+export WEBGPU_USING_SWIFTSHADER=1

.ci/scripts/test_backend.sh

@@ -58,11 +58,10 @@ if [[ "$FLOW" == *vulkan* ]]; then
 fi
 
 if [[ "$FLOW" == *webgpu* ]]; then
-    # Setup swiftshader (software Vulkan adapter for GPU-less runners) and wgpu-native,
-    # which are required to build and run the WebGPU delegate.
+    # Dawn (Tint) + SwiftShader, the spec-faithful headless WebGPU backend.
     source .ci/scripts/setup-webgpu-linux-deps.sh
 
-    EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_WEBGPU=ON"
+    EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_WEBGPU=ON -DDawn_DIR=$Dawn_DIR"
 fi
 
 if [[ "$FLOW" == *arm* ]]; then

.github/workflows/test-webgpu-native.yml

@@ -0,0 +1,54 @@
+name: Test WebGPU Native (Dawn)
+
+# The substantive WebGPU op-coverage gate. The shared operators suite only
+# delegates add.Tensor to WebGPU (everything else is CPU fallback), so the real
+# Dawn coverage comes from the native test executables (rms_norm, multi-dispatch
+# ordering, scratch). This runs them on Dawn (Tint) + SwiftShader, headless, on a
+# CPU runner -- separate from _test_backend.yml so that reusable template stays
+# untouched.
+
+# Nightly-only for now: this job builds SwiftShader from source (vendored LLVM),
+# which is too expensive to run on every PR while the workflow's reliability is
+# still being established. Once it has proven stable, re-enable a scoped PR
+# trigger with a paths: filter (backends/webgpu/**, the webgpu CI scripts, and
+# this file). The pull_request-aware ref/concurrency expressions below are kept
+# intentionally so that re-enable is a one-line change.
+on:
+  schedule:
+    - cron: 0 2 * * *
+  push:
+    branches:
+      - main
+      - release/*
+    tags:
+      - ciflow/nightly/*
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}--${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+jobs:
+  test-webgpu-native:
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    with:
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      runner: linux.4xlarge.memory
+      docker-image: ci-image:executorch-ubuntu-22.04-clang12
+      submodules: recursive
+      timeout: 120
+      script: |
+        set -eux
+
+        # The generic Linux job uses the base conda env, not the image's; activate
+        # the image env (it has the pinned from-source torch). Mirrors
+        # test-vulkan-operators-linux in pull.yml.
+        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
+        conda activate "${CONDA_ENV}"
+
+        # Install the python package + runtime deps (the .pte exporters).
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool cmake
+
+        # Vendor Dawn (Tint) + SwiftShader, then build + run the native executables.
+        source .ci/scripts/setup-webgpu-linux-deps.sh
+        bash backends/webgpu/scripts/test_webgpu_native_ci.sh

backends/arm/runtime/VGFSetup.cpp

@@ -486,6 +486,38 @@ static bool is_tensor_like_descriptor_type(VkDescriptorType descriptor_type) {
       descriptor_type == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
 }
 
+static VkResult submit_and_wait_with_fence(
+    VkDevice device,
+    VkQueue queue,
+    const VkSubmitInfo* submit_info) {
+  VkFence fence = VK_NULL_HANDLE;
+
+  const VkFenceCreateInfo fence_info = {
+      .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
+      .pNext = nullptr,
+      .flags = 0,
+  };
+
+  VkResult result = vkCreateFence(device, &fence_info, nullptr, &fence);
+  if (result != VK_SUCCESS) {
+    ET_LOG(Error, "Failed to create Vulkan fence, error %d", result);
+    return result;
+  }
+
+  result = vkQueueSubmit(queue, 1, submit_info, fence);
+  if (result != VK_SUCCESS) {
+    ET_LOG(Error, "Vulkan queue submit failed, error %d", result);
+    vkDestroyFence(device, fence, nullptr);
+    return result;
+  }
+
+  result = vkWaitForFences(
+      device, 1, &fence, VK_TRUE, std::numeric_limits<uint64_t>::max());
+
+  vkDestroyFence(device, fence, nullptr);
+  return result;
+}
+
 static void record_image_layout_transition(
     VkCommandBuffer command_buffer,
     VkImage image,
@@ -1278,10 +1310,11 @@ VkResult transition_image_layout(
       .signalSemaphoreCount = 0,
       .pSignalSemaphores = nullptr,
   };
-  result = vkQueueSubmit(queue, 1, &submit_info, VK_NULL_HANDLE);
-  if (result == VK_SUCCESS) {
-    result = vkQueueWaitIdle(queue);
-  }
+
+  // creates a temporary one-time command buffer, submits it once, waits, and
+  // frees it immediately.
+  result = submit_and_wait_with_fence(device, queue, &submit_info);
+
   vkFreeCommandBuffers(device, command_pool, 1, &command_buffer);
   return result;
 }
@@ -3078,19 +3111,33 @@ bool VgfRepr::process_vgf(
   {
     VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_ALLOCATE_COMMAND_BUFFER");
 
-    // Allocate command buffer
     VkCommandBufferAllocateInfo buffer_allocate_info{
         .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
         .pNext = nullptr,
         .commandPool = vk_command_pool,
         .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
         .commandBufferCount = 1};
+
     result = vkAllocateCommandBuffers(
         vk_device, &buffer_allocate_info, &vk_execute_cmd);
     if (result != VK_SUCCESS) {
       ET_LOG(Error, "Failed to allocate command buffers");
       return false;
     }
+
+    const VkFenceCreateInfo fence_info{
+        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
+        .pNext = nullptr,
+        .flags = VK_FENCE_CREATE_SIGNALED_BIT,
+    };
+
+    result = vkCreateFence(vk_device, &fence_info, nullptr, &vk_execute_fence);
+    if (result != VK_SUCCESS) {
+      ET_LOG(Error, "Failed to create VGF execute fence, error %d", result);
+      vkFreeCommandBuffers(vk_device, vk_command_pool, 1, &vk_execute_cmd);
+      vk_execute_cmd = VK_NULL_HANDLE;
+      return false;
+    }
   }
 
   {
@@ -3392,31 +3439,51 @@ bool VgfRepr::process_vgf(
 bool VgfRepr::execute_vgf(executorch::runtime::EventTracer* event_tracer) {
   ET_LOG(Info, "Executing vgf");
 
-  VkSubmitInfo submit{VK_STRUCTURE_TYPE_SUBMIT_INFO};
-  submit.commandBufferCount = 1;
-  submit.pCommandBuffers = &vk_execute_cmd;
+  VkSubmitInfo submit{
+      .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
+      .pNext = nullptr,
+      .waitSemaphoreCount = 0,
+      .pWaitSemaphores = nullptr,
+      .pWaitDstStageMask = nullptr,
+      .commandBufferCount = 1,
+      .pCommandBuffers = &vk_execute_cmd,
+      .signalSemaphoreCount = 0,
+      .pSignalSemaphores = nullptr,
+  };
 
   VkResult result;
 
   {
-    VGF_PROFILE_SCOPE(event_tracer, "VGF_QUEUE_SUBMIT");
+    VGF_PROFILE_SCOPE(event_tracer, "VGF_QUEUE_SUBMIT_AND_WAIT_FENCE");
 
-    result = vkQueueSubmit(vk_queue, 1, &submit, VK_NULL_HANDLE);
-  }
+    if (vk_execute_fence == VK_NULL_HANDLE) {
+      ET_LOG(Error, "VGF execute fence is not initialized");
+      return false;
+    }
 
-  if (result != VK_SUCCESS) {
-    ET_LOG(Error, "VGF/VkCommandBuffer command submission failed");
-    return false;
-  }
+    result = vkResetFences(vk_device, 1, &vk_execute_fence);
+    if (result != VK_SUCCESS) {
+      ET_LOG(Error, "VGF/VkFence reset failed, error %d", result);
+      return false;
+    }
 
-  {
-    VGF_PROFILE_SCOPE(event_tracer, "VGF_QUEUE_WAIT_IDLE");
+    result = vkQueueSubmit(vk_queue, 1, &submit, vk_execute_fence);
+    if (result != VK_SUCCESS) {
+      ET_LOG(Error, "VGF/VkFence wait failed, error %d", result);
+      return false;
+    }
 
-    result = vkQueueWaitIdle(vk_queue);
+    result = vkWaitForFences(
+        vk_device,
+        1,
+        &vk_execute_fence,
+        VK_TRUE,
+        std::numeric_limits<uint64_t>::max());
   }
 
   if (result != VK_SUCCESS) {
-    ET_LOG(Error, "VGF/VkQueue wait idle failed");
+    ET_LOG(
+        Error, "VGF/VkCommandBuffer command submission or fence wait failed");
     return false;
   }
 
@@ -3431,6 +3498,11 @@ void VgfRepr::free_vgf() {
     vk_timestamp_query_pool = VK_NULL_HANDLE;
   }
 
+  if (vk_execute_fence != VK_NULL_HANDLE) {
+    vkDestroyFence(vk_device, vk_execute_fence, nullptr);
+    vk_execute_fence = VK_NULL_HANDLE;
+  }
+
   vkFreeCommandBuffers(vk_device, vk_command_pool, 1, &vk_execute_cmd);
   vector<VkDeviceMemory> owned_memory;
   auto remember_owned_memory = [&](VkDeviceMemory memory) {

backends/arm/runtime/VGFSetup.h

@@ -163,6 +163,7 @@ class VgfRepr {
   // per-VgfRepr-instance objects allocated in process_vgf, used (can be more
   // than once) in execute_vgf
   VkCommandBuffer vk_execute_cmd = VK_NULL_HANDLE;
+  VkFence vk_execute_fence = VK_NULL_HANDLE;
   // Note: the vector of tensor memory is stored in IOs above
 
   bool init_timestamp_queries();

backends/cuda/tests/test_int4_dispatch.py

@@ -59,7 +59,10 @@ def _make_int4_linear(N, K, group_size=128, symmetric=False, bias=False):
     )
     int4_w = quantize_weight(w_bf16, config)
 
-    module = nn.Linear(K, N, bias=bias, dtype=torch.bfloat16)
+    # device="cuda" so the random init draws from the CUDA RNG to match the
+    # same random weight as regular int4 dispatch and fit the same numerical
+    # error tolerance.
+    module = nn.Linear(K, N, bias=bias, dtype=torch.bfloat16, device="cuda")
     pack_linear_for_cuda(module, {"weight": int4_w})
     module.cuda()
     return module, w_bf16.cuda()

backends/qualcomm/_passes/decompose_cdist.py

@@ -36,14 +36,19 @@ class DecomposeCDist(ExportPass):
     Decompose for math equivalent op.
     """
 
+    cdist_targets = {
+        torch.ops.aten.cdist.default,
+        torch.ops.aten._cdist_forward.default,
+    }
+
     def __init__(self) -> None:
         super().__init__()
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
         for node in graph.nodes:
             model = CDist()
-            if torch.ops.aten.cdist.default == node.target:
+            if node.target in self.cdist_targets:
                 if len(node.args) > 2:
                     assert (
                         node.args[2] == 2