Update

[ghstack-poisoned]

Author: JCNTH

backends/webgpu/test/native/test_scratch_buffer.cpp

@@ -70,7 +70,9 @@ std::vector<float> readback(
   std::vector<float> out(nbytes / sizeof(float));
   if (cb.status.load(std::memory_order_acquire) == WGPUMapAsyncStatus_Success) {
     const void* m = wgpuBufferGetConstMappedRange(staging, 0, nbytes);
-    std::memcpy(out.data(), m, nbytes);
+    if (m != nullptr) {
+      std::memcpy(out.data(), m, nbytes);
+    }
     wgpuBufferUnmap(staging);
   }
   wgpuBufferRelease(staging);

backends/webgpu/test/ops/dispatch_order/test_dispatch_order.py

@@ -18,6 +18,7 @@
 import torch
 from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
 from executorch.backends.webgpu.test.ops.rms_norm.test_rms_norm import RmsNormModule
+from executorch.backends.webgpu.test.tester import WEBGPU_SUPPORTED_OPS
 from executorch.exir import to_edge_transform_and_lower
 
 
@@ -73,7 +74,8 @@ def _model(kind: str, shape, depth: int) -> torch.nn.Module:
 def _lower(model: torch.nn.Module, x: torch.Tensor):
     ep = torch.export.export(model, (x,))
     return to_edge_transform_and_lower(
-        ep, partitioner=[VulkanPartitioner()]
+        ep,
+        partitioner=[VulkanPartitioner(operator_allowlist=WEBGPU_SUPPORTED_OPS)],
     ).to_executorch()
 
 

backends/webgpu/test/tester.py

@@ -20,6 +20,7 @@
 # Edge ops the WebGPU runtime implements; restricts the Vulkan partitioner.
 WEBGPU_SUPPORTED_OPS = [
     exir_ops.edge.aten.add.Tensor,
+    exir_ops.edge.et_vk.rms_norm.default,
 ]