Add an aten.elu kernel and a single-kernel aten.silu (#24800)

Summary:

ELU had no WebGPU kernel and SiLU ran as two kernels (sigmoid, then mul).
This adds single-kernel elu and silu that handle NaN, infinities and exp overflow, kept from decomposition only where the partitioner will delegate them.
This unblocks streaming audio decoders that use ELU and makes SiLU activations one dispatch.

Reviewed By: SS-JIA

Differential Revision: D127010335

Pull Request resolved: #24800

Author: psiddh

backends/vulkan/op_registry.py

@@ -2304,4 +2304,6 @@ def get_op_features(target: Any) -> OpFeatures:
 
 
 def handles_own_prepacking(target: OpKey) -> bool:
-    return get_op_features(target).supports_prepacking
+    # Ops without a Vulkan registration (e.g. ones only a backend's
+    # op_features_overrides delegates) read constants through prepack nodes.
+    return has_impl(target) and get_op_features(target).supports_prepacking

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -4,8 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pyre-strict
-
 import logging
 from typing import (
     Any,
@@ -521,17 +519,17 @@ def filter_fn(node: torch.fx.Node) -> bool:
 
         return (ops_not_to_decompose, filter_fn)
 
-    def partition(self, exported_program: ExportedProgram) -> PartitionResult:
-        # Run the CapabilityBasedPartitioner to return the largest possible
-        # subgraphs containing the nodes with the tags
-        partition_tags = {}
-
-        # Get all fusable subgraphs from fuse_patterns
-        fusable_subgraphs = vk_patterns.get_all_fusable_subgraphs(
-            exported_program.graph_module,
-            self.options.get("skip_fusion_patterns", ()),
-        )
-
+    def make_supported_operators(
+        self,
+        fusable_subgraphs: Optional[List[PatternMatch]] = None,
+        constant_nodes: Optional[Set[torch.fx.Node]] = None,
+        op_features_overrides: Optional[Dict[OpKey, OpFeatures]] = None,
+        operator_allowlist: Optional[Set[OpKey]] = None,
+        operator_blocklist: Optional[Set[OpKey]] = None,
+    ) -> VulkanSupportedOperators:
+        """The node support checker partition() uses, built from this
+        partitioner's options. The optional arguments replace the
+        corresponding partitioner settings."""
         # small_texture_limits opts into the conservative 3D texture limit that is
         # compatible with most desktop/laptop GPUs (the Vulkan spec only guarantees
         # 2048). An explicit texture_limits always takes precedence.
@@ -542,18 +540,46 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         else:
             texture_limits = utils.DEFAULT_TEXTURE_LIMITS
         buffer_limit: int = self.options.get("buffer_limit", utils.DEFAULT_BUFFER_LIMIT)
-        operator_support = VulkanSupportedOperators(
+        return VulkanSupportedOperators(
             texture_limits,
             buffer_limit,
             require_dynamic_shape=self.options.get("require_dynamic_shapes", False),
             skip_bool_tensors=self.options.get("skip_bool_tensors", False),
-            operator_blocklist=self.operator_blocklist,
-            operator_allowlist=self.operator_allowlist,
+            operator_blocklist=(
+                self.operator_blocklist
+                if operator_blocklist is None
+                else operator_blocklist
+            ),
+            operator_allowlist=(
+                self.operator_allowlist
+                if operator_allowlist is None
+                else operator_allowlist
+            ),
             fusable_subgraphs=fusable_subgraphs,
             nn_module_blocklist=self.nn_module_blocklist,
             nn_module_allowlist=self.nn_module_allowlist,
             downcast_64_bit=self.options.get("downcast_64_bit", True),
-            op_features_overrides=self.op_features_overrides,
+            op_features_overrides=(
+                self.op_features_overrides
+                if op_features_overrides is None
+                else op_features_overrides
+            ),
+            constant_nodes=constant_nodes,
+        )
+
+    def partition(self, exported_program: ExportedProgram) -> PartitionResult:
+        # Run the CapabilityBasedPartitioner to return the largest possible
+        # subgraphs containing the nodes with the tags
+        partition_tags = {}
+
+        # Get all fusable subgraphs from fuse_patterns
+        fusable_subgraphs = vk_patterns.get_all_fusable_subgraphs(
+            exported_program.graph_module,
+            self.options.get("skip_fusion_patterns", ()),
+        )
+
+        operator_support = self.make_supported_operators(
+            fusable_subgraphs=fusable_subgraphs,
             constant_nodes=utils.get_constant_nodes(exported_program),
         )
         capability_partitioner = CapabilityBasedPartitioner(

backends/webgpu/partitioner/partitioner.py

@@ -4,8 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pyre-strict
-
 import copy
 from typing import Any, Callable, Dict, final, List, Optional, Tuple
 
@@ -18,7 +16,10 @@
     OpKey,
     vulkan_supported_ops,
 )
-from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import (
+    VulkanPartitioner,
+    VulkanSupportedOperators,
+)
 from executorch.exir.backend.partitioner import Partitioner, PartitionResult
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
@@ -54,8 +55,17 @@ def _layer_norm_check(
     return all(d == 1 for d in normalized_shape[:-1]) and vulkan_check(node)
 
 
-# Ops whose WebGPU kernels cover less than the Vulkan registry advertises, so the
-# unsupported forms stay on CPU instead of failing at delegate init.
+def _fp32_unary_features() -> OpFeatures:
+    return OpFeatures(
+        inputs_storage=vk_utils.ANY_STORAGE,
+        inputs_dtypes={torch.float32},
+        supports_resize=True,
+    )
+
+
+# Ops whose WebGPU kernels cover less (or more) than the Vulkan registry
+# advertises, so unsupported forms stay on CPU instead of failing at delegate
+# init, and WebGPU-only kernels are delegated.
 _WEBGPU_OP_FEATURES: Dict[OpKey, OpFeatures] = {
     # fp32 and int32 kernels only; int64 is downcast to int32 on device.
     exir_ops.edge.aten.arange.start_step: OpFeatures(
@@ -66,12 +76,16 @@ def _layer_norm_check(
     exir_ops.edge.aten.native_layer_norm.default: _webgpu_features(
         exir_ops.edge.aten.native_layer_norm.default, _layer_norm_check
     ),
+    exir_ops.edge.aten.elu.default: _fp32_unary_features(),
+    exir_ops.edge.aten.silu.default: _fp32_unary_features(),
 }
 
 
 # Vulkan fusion patterns whose fused ops the WebGPU runtime does not implement.
 _WEBGPU_SKIP_FUSION_PATTERNS = {"conv2d_with_epilogue"}
 
+_MUL_TARGETS = (torch.ops.aten.mul.Tensor, exir_ops.edge.aten.mul.Tensor)
+
 
 @final
 class WebGPUPartitioner(Partitioner):
@@ -102,7 +116,59 @@ def __init__(
     def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
-        return self._vulkan_partitioner.ops_to_not_decompose(ep)
+        ops, vulkan_filter = self._vulkan_partitioner.ops_to_not_decompose(ep)
+        silu_support = self._silu_support()
+
+        def filter_fn(node: torch.fx.Node) -> bool:
+            if node.target in (
+                torch.ops.aten.silu.default,
+                exir_ops.edge.aten.silu.default,
+            ):
+                return self._keeps_silu(node, ep, silu_support)
+            return vulkan_filter is None or vulkan_filter(node)
+
+        return ops + [torch.ops.aten.silu.default], filter_fn
+
+    def _silu_support(self) -> VulkanSupportedOperators:
+        """The partitioner's own node support check, also accepting the aten
+        silu that ops_to_not_decompose sees before the edge conversion."""
+        vulkan = self._vulkan_partitioner
+        aten_silu = torch.ops.aten.silu.default
+        edge_silu = exir_ops.edge.aten.silu.default
+        allowlist = vulkan.operator_allowlist
+        if allowlist is not None and edge_silu in allowlist:
+            allowlist = allowlist | {aten_silu}
+        blocklist = vulkan.operator_blocklist
+        if edge_silu in blocklist:
+            blocklist = blocklist | {aten_silu}
+        return vulkan.make_supported_operators(
+            op_features_overrides={
+                **_WEBGPU_OP_FEATURES,
+                aten_silu: _WEBGPU_OP_FEATURES[edge_silu],
+            },
+            operator_allowlist=allowlist,
+            operator_blocklist=blocklist,
+        )
+
+    def _keeps_silu(
+        self,
+        node: torch.fx.Node,
+        ep: ExportedProgram,
+        silu_support: VulkanSupportedOperators,
+    ) -> bool:
+        """Whether `node` stays an aten.silu for the WebGPU silu kernel.
+
+        A kept op that is not delegated fails lowering, so keep it only where
+        the partitioner will delegate it: in the top-level graph (control-flow
+        submodules are not partitioned) and passing the partitioner's support
+        check. Where it feeds a mul it decomposes to x * sigmoid(x), as before,
+        so the SwiGLU fusion still matches.
+        """
+        return (
+            node.graph is ep.graph_module.graph
+            and not any(user.target in _MUL_TARGETS for user in node.users)
+            and silu_support.is_node_supported({}, node)
+        )
 
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         return self._vulkan_partitioner.partition(exported_program)

backends/webgpu/runtime/WebGPUShaderRegistry.cpp

@@ -135,12 +135,14 @@
 #include <executorch/backends/webgpu/runtime/ops/unary/abs_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/clamp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/cos_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/elu_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/exp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/neg_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/round_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/silu_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/sin_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/tanh_wgsl.h>
@@ -157,7 +159,7 @@
 namespace executorch::backends::webgpu {
 namespace {
 
-constexpr std::array<WebGPUShaderInfo, 139> kShaderRegistry = {{
+constexpr std::array<WebGPUShaderInfo, 141> kShaderRegistry = {{
     {
         "abs",
         kAbsWGSL,
@@ -466,6 +468,13 @@ constexpr std::array<WebGPUShaderInfo, 139> kShaderRegistry = {{
         kDequantizePerTensorWorkgroupSizeY,
         kDequantizePerTensorWorkgroupSizeZ,
     },
+    {
+        "elu",
+        kEluWGSL,
+        kEluWorkgroupSizeX,
+        kEluWorkgroupSizeY,
+        kEluWorkgroupSizeZ,
+    },
     {
         "embedding",
         kEmbeddingWGSL,
@@ -1012,6 +1021,13 @@ constexpr std::array<WebGPUShaderInfo, 139> kShaderRegistry = {{
         kSigmoidWorkgroupSizeY,
         kSigmoidWorkgroupSizeZ,
     },
+    {
+        "silu",
+        kSiluWGSL,
+        kSiluWorkgroupSizeX,
+        kSiluWorkgroupSizeY,
+        kSiluWorkgroupSizeZ,
+    },
     {
         "silu_mul_fused",
         kSiluMulFusedWGSL,

backends/webgpu/runtime/ops/unary/Activations.cpp

@@ -6,17 +6,20 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/UnaryOp.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/abs_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/clamp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/cos_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/elu_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/exp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/neg_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/pow_scalar_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/round_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/unary/silu_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/sin_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/tanh_wgsl.h>
@@ -100,6 +103,32 @@ void hardswish_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       "hardswish");
 }
 
+void silu_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  add_unary_op(
+      graph, args.at(0), args.at(1), kSiluWGSL, kSiluWorkgroupSizeX, "silu");
+}
+
+void elu_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  // aten.elu.default args: [in, alpha, scale, input_scale, out]; each scalar
+  // defaults to 1.
+  if (args.size() != 5) {
+    throw std::runtime_error("WebGPU elu: expected 5 args");
+  }
+  const float alpha = utils::scalar_or(graph, args.at(1), 1.0f);
+  const float scale = utils::scalar_or(graph, args.at(2), 1.0f);
+  const float input_scale = utils::scalar_or(graph, args.at(3), 1.0f);
+  add_unary_op(
+      graph,
+      args.at(0),
+      args.at(4),
+      kEluWGSL,
+      kEluWorkgroupSizeX,
+      "elu",
+      alpha,
+      scale,
+      input_scale);
+}
+
 void clamp_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   // aten.clamp.default args: [in, min, max, out]; min/max None -> +/-inf.
   const float lo = get_val_or_inf(graph, args.at(1), /*is_max=*/false);
@@ -156,6 +185,8 @@ WEBGPU_REGISTER_OPERATORS {
   WEBGPU_REGISTER_OP(aten.round.default, round_impl);
   WEBGPU_REGISTER_OP(aten.neg.default, neg_impl);
   WEBGPU_REGISTER_OP(aten.hardswish.default, hardswish_impl);
+  WEBGPU_REGISTER_OP(aten.silu.default, silu_impl);
+  WEBGPU_REGISTER_OP(aten.elu.default, elu_impl);
   WEBGPU_REGISTER_OP(aten.clamp.default, clamp_impl);
   WEBGPU_REGISTER_OP(aten.hardtanh.default, hardtanh_impl);
   WEBGPU_REGISTER_OP(aten.pow.Tensor_Scalar, pow_scalar_impl);

backends/webgpu/runtime/ops/unary/UnaryOp.cpp

@@ -20,13 +20,15 @@ namespace executorch::backends::webgpu {
 namespace {
 
 // Canonical unary uniform layout; min/max mirror Vulkan min_max (used by the
-// clamp-family shaders; ignored by no-param activations).
+// clamp-family shaders; ignored by no-param activations). aux is a third
+// scalar (elu's input_scale).
 struct UnaryParams {
   uint32_t num_elements;
   float min;
   float max;
-  uint32_t _pad;
+  float aux;
 };
+static_assert(sizeof(UnaryParams) == 16, "UnaryParams must be 16 bytes");
 
 } // namespace
 
@@ -38,7 +40,8 @@ void add_unary_op(
     uint32_t wg_size_x,
     const char* op_name,
     float min,
-    float max) {
+    float max,
+    float aux) {
   WGPUDevice device = graph.device();
 
   const auto& in_tensor = graph.get_tensor(in_id);
@@ -66,6 +69,7 @@ void add_unary_op(
   params.num_elements = num_elements;
   params.min = min;
   params.max = max;
+  params.aux = aux;
 
   WGPUBuffer uniform_buffer =
       utils::make_uniform(device, &params, sizeof(UnaryParams));
@@ -99,7 +103,7 @@ void add_unary_op(
   WGPUBuffer params_buf = uniform_buffer;
   graph.add_tensor_resize_hook(
       in_id,
-      [in_id, out_id, wg_size, dispatch_idx, params_buf, min, max](
+      [in_id, out_id, wg_size, dispatch_idx, params_buf, min, max, aux](
           WebGPUGraph& g) {
         const auto& d = g.cur_dims(in_id);
         const uint64_t numel = utils::numel_of(d);
@@ -108,6 +112,7 @@ void add_unary_op(
         p.num_elements = static_cast<uint32_t>(numel);
         p.min = min;
         p.max = max;
+        p.aux = aux;
         wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
         const utils::WgCount wgc = utils::compute_2d_workgroup_count(
             g.device(), static_cast<uint32_t>(numel), wg_size, "unary(resize)");