[ExecuTorch][WebGPU] Fuse the SwiGLU elementwise chain

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/CMakeLists.txt

@@ -29,6 +29,7 @@ set(WEBGPU_SRCS
  runtime/WebGPUBackend.cpp
  runtime/WebGPUExecutionOptions.cpp
  runtime/WebGPUGraph.cpp
+ runtime/passes/SwiGLU.cpp
  runtime/WebGPUDelegateHeader.cpp
  runtime/WebGPUDevice.cpp
  runtime/WebGPUQueryPool.cpp

backends/webgpu/runtime/WebGPUGraph.h

@@ -516,6 +516,13 @@ class WebGPUGraph {
  return value_types_[id];
  }
 
+ // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has
+ // none; fusion passes use this to reject candidates aliased with something
+ // the planner may reuse outside the fusion's control.
+ int mem_obj_id(int id) const {
+ return tensor_mem_obj_ids_[id];
+ }
+
  public:
  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
  bool kv_f16() const {
@@ -681,15 +688,6 @@ class WebGPUGraph {
  // resize hook.
  void add_qkv_fused_dispatch(QkvFusionGroup& g);
  void add_qkv_fused_hook(const QkvFusionGroup& g);
-
- // SwiGLU fusion: emit ONE fused elementwise dispatch
- // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
- // muls. `out` is repointed to a private pooled buffer (aliasing guard);
- // `gate` is likewise given a private pooled buffer at its producer op by the
- // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
- // would stomp gate before the fused reads it). Only used during build(); the
- // detection maps are empty (inert) when no SwiGLU triple matches.
- void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
 };
 
 } // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUUtils.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -31,6 +32,18 @@ inline uint64_t numel_of(const std::vector<int64_t>& dims) {
  return numel(dims);
 }
 
+// fp32, non-null-buffer tensor with byte size matching its element count;
+// the dtype/aliasing precondition fusion passes require of every operand.
+inline bool is_fp32_tensor(const WebGPUTensor& tensor) {
+ if (tensor.is_int || tensor.elem_size != sizeof(float) ||
+ tensor.buffer == nullptr) {
+ return false;
+ }
+ const uint64_t elems = numel_of(tensor.dims);
+ return elems <= std::numeric_limits<size_t>::max() / sizeof(float) &&
+ tensor.nbytes == static_cast<size_t>(elems) * sizeof(float);
+}
+
 // Clamp workgroup size to device limit (SwiftShader caps at 128).
 inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
  WGPULimits limits = {};

backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl

@@ -7,14 +7,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }

backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from silu_mul_fused.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
+// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461
 inline constexpr const char* kSiluMulFusedWGSL = R"(
 @group(0) @binding(0) var<storage, read> gate: array<f32>;
 @group(0) @binding(1) var<storage, read> up: array<f32>;
@@ -24,14 +24,13 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
-// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
-// sigmoid + silu are computed in registers (never written to memory), so gate + up
-// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
-// multiply order match the original ops -> bit-exact.
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
- let idx = gid.x;
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+ @builtin(global_invocation_id) gid: vec3<u32>,
+ @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+ let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
  if (idx >= params.num_elements) {
  return;
  }