diff --git a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp new file mode 100644 index 00000000000..d2f923880bf --- /dev/null +++ b/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 +#include +#include + +#include + +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& 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 diff --git a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl new file mode 100644 index 00000000000..894c32a3dc2 --- /dev/null +++ b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl @@ -0,0 +1,51 @@ +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 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]); +} diff --git a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h b/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h new file mode 100644 index 00000000000..6488e360518 --- /dev/null +++ b/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 + +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 input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 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