[ExecuTorch][WebGPU] Port pow (aten.pow.Tensor_Tensor)

backends/webgpu/runtime/ops/pow/BinaryOp.cpp

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/binary/BinaryOp.h>
+#include <executorch/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h>
+
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// aten.pow.Tensor_Tensor -> pow(in1, in2), with NumPy broadcasting (mirrors mul
+// + Vulkan). WGSL pow(x,y) is undefined for x<0 (exact Vulkan-GLSL parity).
+void pow_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  add_binary_broadcast_op(
+      graph,
+      args.at(0),
+      args.at(1),
+      args.at(2),
+      kBinaryPowWGSL,
+      kBinaryPowWorkgroupSizeX,
+      "pow");
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.pow.Tensor_Tensor, pow_impl);
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/pow/binary_pow.wgsl

@@ -0,0 +1,51 @@
+@group(0) @binding(0) var<storage, read> input1: array<f32>;
+@group(0) @binding(1) var<storage, read> input2: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
+
+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>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    // Fast path: every input dim matches the output dim -> elementwise.
+    var same = true;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        if (in1_meta.sizes[d] != out_meta.sizes[d] ||
+            in2_meta.sizes[d] != out_meta.sizes[d]) {
+            same = false;
+        }
+    }
+    if (same) {
+        output[idx] = pow(input1[idx], input2[idx]);
+        return;
+    }
+
+    // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize.
+    var rem = idx;
+    var l1: u32 = 0u;
+    var l2: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
+        l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
+    }
+    output[idx] = pow(input1[l1], input2[l2]);
+}

backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h

@@ -0,0 +1,75 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from binary_pow.wgsl - DO NOT EDIT.
+// wgsl-sha256: ffc116e845c861fb2a43f0c3ca86f150f8e1e95f9d3dde73a3886cd2a1ebff93
+inline constexpr const char* kBinaryPowWGSL = R"(
+@group(0) @binding(0) var<storage, read> input1: array<f32>;
+@group(0) @binding(1) var<storage, read> input2: array<f32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
+
+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>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    // Fast path: every input dim matches the output dim -> elementwise.
+    var same = true;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        if (in1_meta.sizes[d] != out_meta.sizes[d] ||
+            in2_meta.sizes[d] != out_meta.sizes[d]) {
+            same = false;
+        }
+    }
+    if (same) {
+        output[idx] = pow(input1[idx], input2[idx]);
+        return;
+    }
+
+    // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize.
+    var rem = idx;
+    var l1: u32 = 0u;
+    var l2: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        l1 = l1 + min(coord, in1_meta.sizes[d] - 1u) * in1_meta.strides[d];
+        l2 = l2 + min(coord, in2_meta.sizes[d] - 1u) * in2_meta.strides[d];
+    }
+    output[idx] = pow(input1[l1], input2[l2]);
+}
+)";
+
+inline constexpr uint32_t kBinaryPowWorkgroupSizeX = 64;
+inline constexpr uint32_t kBinaryPowWorkgroupSizeY = 1;
+inline constexpr uint32_t kBinaryPowWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu