diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 105bcea89a6..e8f22c8661d 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1746,6 +1746,35 @@ def register_rms_norm(): ) +@update_features( + [ + exir_ops.edge.aten.ne.Scalar, + exir_ops.edge.aten.lt.Scalar, + exir_ops.edge.aten.le.Scalar, + exir_ops.edge.aten.ge.Scalar, + exir_ops.edge.aten.gt.Scalar, + ] +) +def register_compare_scalar_ops(): + return OpFeatures( + inputs_storage=utils.ANY_STORAGE, + inputs_dtypes=utils.FP_INT_T, + outputs_dtypes=utils.BOOL_T, + supports_resize=True, + supports_highdim=True, + ) + + +@update_features(exir_ops.edge.aten.logical_not.default) +def register_logical_not(): + return OpFeatures( + inputs_storage=utils.ANY_STORAGE, + inputs_dtypes=utils.BOOL_T, + supports_resize=True, + supports_highdim=True, + ) + + ####################### ## Utility functions ## ####################### diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 36626b173e0..47c485c2dc6 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -59,6 +59,8 @@ set(WEBGPU_SRCS runtime/ops/reduce/Reduce.cpp runtime/ops/div/BinaryOp.cpp runtime/ops/sub/BinaryOp.cpp + runtime/ops/where/Where.cpp + runtime/ops/boolean_op/BooleanOp.cpp runtime/ops/linear/Linear.cpp ) diff --git a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp new file mode 100644 index 00000000000..27c87a2bbeb --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp @@ -0,0 +1,251 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Shared uniform for the byte-packed-bool op family. `scalar` is unused (0) for +// the scalar-less (unary) logical_not variant. +struct BoolOpParams { + uint32_t num_elements; + float scalar; + uint32_t _pad1; + uint32_t _pad2; +}; + +float read_scalar(WebGPUGraph& graph, int id, const char* op_name) { + if (graph.get_value_type(id) == WebGPUGraph::ValueType::Double) { + return static_cast(graph.get_double(id)); + } + if (graph.get_value_type(id) == WebGPUGraph::ValueType::Int) { + return static_cast(graph.get_int(id)); + } + throw std::runtime_error(std::string(op_name) + ": scalar is not int/double"); +} + +// Dispatch a byte-packed-bool op (scalar compares + logical_not): read an input +// buffer, write a u32 output packing 4 bool bytes per word, one thread per word +// (no inter-thread write race). The caller validates dtypes/shapes and supplies +// `numel` (logical element count = prod(dims)), the scalar (0 when unused), and +// the per-variant shader from boolean_op.yaml. +void dispatch_bool_op( + WebGPUGraph& graph, + int self_id, + int out_id, + uint32_t numel, + float scalar, + const char* wgsl, + uint32_t wg_size_x, + const char* op_name) { + WGPUDevice device = graph.device(); + const auto& self_tensor = graph.get_tensor(self_id); + const auto& out_tensor = graph.get_tensor(out_id); + + // The output (and logical_not's packed-bool input) is a u32 array; round the + // binding up to whole words even when the byte count isn't a multiple of 4. + const size_t self_bind_size = (self_tensor.nbytes + 3) & ~size_t(3); + const size_t out_bind_size = (out_tensor.nbytes + 3) & ~size_t(3); + const uint32_t n_words = (numel + 3u) / 4u; + + uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x); + uint32_t workgroup_count = + utils::compute_1d_workgroup_count(device, n_words, wg_size, op_name); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + BoolOpParams params = {numel, scalar, 0u, 0u}; + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(BoolOpParams)); + graph.add_uniform_buffer_bytes(sizeof(BoolOpParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {wgsl, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[3] = {}; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + for (uint32_t i = 0; i < 3; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + } + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 3; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[3] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = self_tensor.buffer; + bg_entries[0].size = self_bind_size; + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_bind_size; + bg_entries[2].binding = 2; + bg_entries[2].buffer = params_buf; + bg_entries[2].size = sizeof(BoolOpParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = + graph.add_dispatch({pipeline, bind_group, workgroup_count}); + + WGPUBuffer p_buf = params_buf; + auto resize = [self_id, out_id, scalar, wg_size, dispatch_idx, p_buf, op_name]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(self_id); + uint32_t n = 1u; + for (auto x : d) { + n *= static_cast(x); + } + g.set_cur_dims(out_id, d); + BoolOpParams p = {n, scalar, 0u, 0u}; + wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); + const uint32_t nw = (n + 3u) / 4u; + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count(g.device(), nw, wg_size, op_name); + }; + graph.add_tensor_resize_hook(self_id, resize); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +// cmp(self[i], scalar) -> byte-packed bool. args: [self, scalar, out]. +void compare_dispatch( + WebGPUGraph& graph, + const std::vector& args, + const char* wgsl, + uint32_t wg_size_x, + const char* op_name) { + const int self_id = args.at(0); + const int out_id = args.at(args.size() - 1); + const float scalar = read_scalar(graph, args.at(1), op_name); + + const auto& self_tensor = graph.get_tensor(self_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (self_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error(std::string(op_name) + ": null buffer binding"); + } + if (self_tensor.nbytes % sizeof(float) != 0) { + throw std::runtime_error(std::string(op_name) + ": self is not fp32"); + } + const uint32_t numel = + static_cast(self_tensor.nbytes / sizeof(float)); + if (out_tensor.nbytes != static_cast(numel)) { + throw std::runtime_error( + std::string(op_name) + ": out is not a 1-byte (bool) tensor"); + } + dispatch_bool_op( + graph, self_id, out_id, numel, scalar, wgsl, wg_size_x, op_name); +} + +void eq_scalar_impl(WebGPUGraph& g, const std::vector& a) { + compare_dispatch(g, a, kCompareEqWGSL, kCompareEqWorkgroupSizeX, "eq.Scalar"); +} +void ne_scalar_impl(WebGPUGraph& g, const std::vector& a) { + compare_dispatch(g, a, kCompareNeWGSL, kCompareNeWorkgroupSizeX, "ne.Scalar"); +} +void le_scalar_impl(WebGPUGraph& g, const std::vector& a) { + compare_dispatch(g, a, kCompareLeWGSL, kCompareLeWorkgroupSizeX, "le.Scalar"); +} +void ge_scalar_impl(WebGPUGraph& g, const std::vector& a) { + compare_dispatch(g, a, kCompareGeWGSL, kCompareGeWorkgroupSizeX, "ge.Scalar"); +} +void lt_scalar_impl(WebGPUGraph& g, const std::vector& a) { + compare_dispatch(g, a, kCompareLtWGSL, kCompareLtWorkgroupSizeX, "lt.Scalar"); +} +void gt_scalar_impl(WebGPUGraph& g, const std::vector& a) { + compare_dispatch(g, a, kCompareGtWGSL, kCompareGtWorkgroupSizeX, "gt.Scalar"); +} + +// logical_not: byte-packed bool -> byte-packed bool. args: [self, out]. +void logical_not_impl(WebGPUGraph& graph, const std::vector& args) { + const int self_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + const auto& self_tensor = graph.get_tensor(self_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (self_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("logical_not: null buffer binding"); + } + if (out_tensor.nbytes != self_tensor.nbytes) { + throw std::runtime_error("logical_not: self/out byte-size mismatch"); + } + const uint32_t numel = static_cast(self_tensor.nbytes); + dispatch_bool_op( + graph, + self_id, + out_id, + numel, + 0.0f, + kLogicalNotWGSL, + kLogicalNotWorkgroupSizeX, + "logical_not"); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.eq.Scalar, eq_scalar_impl); + WEBGPU_REGISTER_OP(aten.ne.Scalar, ne_scalar_impl); + WEBGPU_REGISTER_OP(aten.le.Scalar, le_scalar_impl); + WEBGPU_REGISTER_OP(aten.ge.Scalar, ge_scalar_impl); + WEBGPU_REGISTER_OP(aten.lt.Scalar, lt_scalar_impl); + WEBGPU_REGISTER_OP(aten.gt.Scalar, gt_scalar_impl); + WEBGPU_REGISTER_OP(aten.logical_not.default, logical_not_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl b/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl new file mode 100644 index 00000000000..4f32792d7dc --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl @@ -0,0 +1,37 @@ +@group(0) @binding(0) var input: array<${IN_TYPE}>; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return ${OP_EXPR}; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} diff --git a/backends/webgpu/runtime/ops/boolean_op/boolean_op.yaml b/backends/webgpu/runtime/ops/boolean_op/boolean_op.yaml new file mode 100644 index 00000000000..4b14860a7e2 --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/boolean_op.yaml @@ -0,0 +1,20 @@ +boolean_op: + parameter_names_with_default_values: + IN_TYPE: f32 + OP_EXPR: input[i] == params.scalar + shader_variants: + - NAME: compare_eq + OP_EXPR: input[i] == params.scalar + - NAME: compare_ne + OP_EXPR: input[i] != params.scalar + - NAME: compare_le + OP_EXPR: input[i] <= params.scalar + - NAME: compare_ge + OP_EXPR: input[i] >= params.scalar + - NAME: compare_lt + OP_EXPR: input[i] < params.scalar + - NAME: compare_gt + OP_EXPR: input[i] > params.scalar + - NAME: logical_not + IN_TYPE: u32 + OP_EXPR: ((input[i >> 2u] >> ((i & 3u) * 8u)) & 0xFFu) == 0u diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h new file mode 100644 index 00000000000..8b7b44e7430 --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h @@ -0,0 +1,61 @@ +/* + * 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 boolean_op.wgsl - DO NOT EDIT. +// wgsl-sha256: 2a3c203d0255086a6e67a9e7cb08538858e8e6cb6a5542f6acebf04b8ccca6b7 +inline constexpr const char* kCompareEqWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return input[i] == params.scalar; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} +)"; + +inline constexpr uint32_t kCompareEqWorkgroupSizeX = 64; +inline constexpr uint32_t kCompareEqWorkgroupSizeY = 1; +inline constexpr uint32_t kCompareEqWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h new file mode 100644 index 00000000000..bcf7239bd7e --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h @@ -0,0 +1,61 @@ +/* + * 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 boolean_op.wgsl - DO NOT EDIT. +// wgsl-sha256: 5fb319a54ed666644f118119639dd8cd33256ea5e5163c4c1392ce6e84b369cc +inline constexpr const char* kCompareGeWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return input[i] >= params.scalar; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} +)"; + +inline constexpr uint32_t kCompareGeWorkgroupSizeX = 64; +inline constexpr uint32_t kCompareGeWorkgroupSizeY = 1; +inline constexpr uint32_t kCompareGeWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h new file mode 100644 index 00000000000..8f911180c8a --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h @@ -0,0 +1,61 @@ +/* + * 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 boolean_op.wgsl - DO NOT EDIT. +// wgsl-sha256: b4a35be8a34774a47be457130125336b5accf8d24615514f3fe1559159644491 +inline constexpr const char* kCompareGtWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return input[i] > params.scalar; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} +)"; + +inline constexpr uint32_t kCompareGtWorkgroupSizeX = 64; +inline constexpr uint32_t kCompareGtWorkgroupSizeY = 1; +inline constexpr uint32_t kCompareGtWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h new file mode 100644 index 00000000000..87b4ff91c2a --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h @@ -0,0 +1,61 @@ +/* + * 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 boolean_op.wgsl - DO NOT EDIT. +// wgsl-sha256: 64b8988db9611c1aaff3ba373d32ecb7b276340df337b38a49e12a92727c00cb +inline constexpr const char* kCompareLeWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return input[i] <= params.scalar; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} +)"; + +inline constexpr uint32_t kCompareLeWorkgroupSizeX = 64; +inline constexpr uint32_t kCompareLeWorkgroupSizeY = 1; +inline constexpr uint32_t kCompareLeWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h new file mode 100644 index 00000000000..930601d229c --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h @@ -0,0 +1,61 @@ +/* + * 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 boolean_op.wgsl - DO NOT EDIT. +// wgsl-sha256: df350a7d0b099d38f9e77a27ffe9e56afa93b2c6a3ff84006227e1c1c0b96521 +inline constexpr const char* kCompareLtWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return input[i] < params.scalar; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} +)"; + +inline constexpr uint32_t kCompareLtWorkgroupSizeX = 64; +inline constexpr uint32_t kCompareLtWorkgroupSizeY = 1; +inline constexpr uint32_t kCompareLtWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h new file mode 100644 index 00000000000..9020fb77671 --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h @@ -0,0 +1,61 @@ +/* + * 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 boolean_op.wgsl - DO NOT EDIT. +// wgsl-sha256: 4f5492c393e494e6a9a734ea797cbfa85f0f1579451fc46f4291f57e1b5daaa2 +inline constexpr const char* kCompareNeWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return input[i] != params.scalar; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} +)"; + +inline constexpr uint32_t kCompareNeWorkgroupSizeX = 64; +inline constexpr uint32_t kCompareNeWorkgroupSizeY = 1; +inline constexpr uint32_t kCompareNeWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h new file mode 100644 index 00000000000..65239060454 --- /dev/null +++ b/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h @@ -0,0 +1,61 @@ +/* + * 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 boolean_op.wgsl - DO NOT EDIT. +// wgsl-sha256: 06dc55c61e55c35eb12abf548fafb495097fa8d22daea930af6a32ed724a5b6a +inline constexpr const char* kLogicalNotWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, + scalar: f32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 64u; + +// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not). +fn elem_bool(i: u32) -> bool { + return ((input[i >> 2u] >> ((i & 3u) * 8u)) & 0xFFu) == 0u; +} + +// One thread per output u32 word packs 4 bool bytes -> no inter-thread race. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + let n_words = (params.num_elements + 3u) / 4u; + if (word_idx >= n_words) { + return; + } + var packed: u32 = 0u; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let i = word_idx * 4u + j; + if (i < params.num_elements) { + if (elem_bool(i)) { + packed = packed | (1u << (j * 8u)); + } + } + } + output[word_idx] = packed; +} +)"; + +inline constexpr uint32_t kLogicalNotWorkgroupSizeX = 64; +inline constexpr uint32_t kLogicalNotWorkgroupSizeY = 1; +inline constexpr uint32_t kLogicalNotWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/where/Where.cpp b/backends/webgpu/runtime/ops/where/Where.cpp new file mode 100644 index 00000000000..bde779498dd --- /dev/null +++ b/backends/webgpu/runtime/ops/where/Where.cpp @@ -0,0 +1,241 @@ +/* + * 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 +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +// where.self selects self/other by cond; broadcasts all three. +void where_impl(WebGPUGraph& graph, const std::vector& args) { + // args: [condition, self, other, out]; out is the last value id. + const int cond_id = args.at(0); + const int a_id = args.at(1); + const int b_id = args.at(2); + const int out_id = args.at(args.size() - 1); + + WGPUDevice device = graph.device(); + + const auto& cond_tensor = graph.get_tensor(cond_id); + const auto& a_tensor = graph.get_tensor(a_id); + const auto& b_tensor = graph.get_tensor(b_id); + const auto& out_tensor = graph.get_tensor(out_id); + + if (cond_tensor.buffer == nullptr || a_tensor.buffer == nullptr || + b_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("where: null buffer binding"); + } + if (out_tensor.dims.size() > kTensorMetaMaxNdim || + cond_tensor.dims.size() > kTensorMetaMaxNdim || + a_tensor.dims.size() > kTensorMetaMaxNdim || + b_tensor.dims.size() > kTensorMetaMaxNdim) { + throw std::runtime_error("where: tensor rank exceeds 4 (MAX_NDIM)"); + } + + const uint32_t out_ndim = static_cast(out_tensor.dims.size()); + + TensorMeta out_meta; + TensorMeta cond_meta; + TensorMeta a_meta; + TensorMeta b_meta; + fill_tensor_meta_broadcast(out_tensor, out_ndim, &out_meta); + fill_tensor_meta_broadcast(cond_tensor, out_ndim, &cond_meta); + fill_tensor_meta_broadcast(a_tensor, out_ndim, &a_meta); + fill_tensor_meta_broadcast(b_tensor, out_ndim, &b_meta); + + // a/b/out are fp32; cond is 1-byte bool (read byte-packed as array). + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float) || + a_tensor.nbytes != static_cast(a_meta.numel) * sizeof(float) || + b_tensor.nbytes != static_cast(b_meta.numel) * sizeof(float)) { + throw std::runtime_error( + "where: non-fp32 self/other (nbytes != numel * 4)"); + } + if (cond_tensor.nbytes != static_cast(cond_meta.numel)) { + throw std::runtime_error("where: condition is not a 1-byte (bool) tensor"); + } + + // Buffer is 4-byte-rounded at alloc; bind the padded span for the u32 read. + const size_t cond_bind_size = (cond_tensor.nbytes + 3) & ~size_t(3); + + uint32_t wg_size = utils::clamp_workgroup_size(device, kWhereWorkgroupSizeX); + uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, out_meta.numel, wg_size, "where"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer cond_meta_buf = + utils::make_uniform(device, &cond_meta, sizeof(TensorMeta)); + WGPUBuffer a_meta_buf = + utils::make_uniform(device, &a_meta, sizeof(TensorMeta)); + WGPUBuffer b_meta_buf = + utils::make_uniform(device, &b_meta, sizeof(TensorMeta)); + graph.add_uniform_buffer_bytes(4 * sizeof(TensorMeta)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kWhereWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[8] = {}; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[3].buffer.type = WGPUBufferBindingType_Storage; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + entries[5].buffer.type = WGPUBufferBindingType_Uniform; + entries[6].buffer.type = WGPUBufferBindingType_Uniform; + entries[7].buffer.type = WGPUBufferBindingType_Uniform; + for (uint32_t i = 0; i < 8; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + } + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 8; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBindGroupEntry bg_entries[8] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = cond_tensor.buffer; + bg_entries[0].size = cond_bind_size; + bg_entries[1].binding = 1; + bg_entries[1].buffer = a_tensor.buffer; + bg_entries[1].size = a_tensor.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = b_tensor.buffer; + bg_entries[2].size = b_tensor.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = out_tensor.buffer; + bg_entries[3].size = out_tensor.nbytes; + bg_entries[4].binding = 4; + bg_entries[4].buffer = out_meta_buf; + bg_entries[4].size = sizeof(TensorMeta); + bg_entries[5].binding = 5; + bg_entries[5].buffer = cond_meta_buf; + bg_entries[5].size = sizeof(TensorMeta); + bg_entries[6].binding = 6; + bg_entries[6].buffer = a_meta_buf; + bg_entries[6].size = sizeof(TensorMeta); + bg_entries[7].binding = 7; + bg_entries[7].buffer = b_meta_buf; + bg_entries[7].size = sizeof(TensorMeta); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 8; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + const size_t dispatch_idx = + graph.add_dispatch({pipeline, bind_group, workgroup_count}); + + // Dynamic shapes: rebuild the 4 broadcast TensorMeta UBOs + dispatch count. + WGPUBuffer o_buf = out_meta_buf, c_buf = cond_meta_buf, a_buf = a_meta_buf, + bb_buf = b_meta_buf; + auto where_resize = [cond_id, + a_id, + b_id, + out_id, + wg_size, + dispatch_idx, + o_buf, + c_buf, + a_buf, + bb_buf](WebGPUGraph& g) { + const auto& c = g.cur_dims(cond_id); + const auto& a = g.cur_dims(a_id); + const auto& b = g.cur_dims(b_id); + const size_t r = std::max({c.size(), a.size(), b.size()}); + auto dim_at = [r](const std::vector& d, size_t i) -> int64_t { + return (i + d.size() < r) ? 1 : d[i - (r - d.size())]; + }; + std::vector out_d(r, 1); + for (size_t i = 0; i < r; i++) { + const int64_t cv = dim_at(c, i), av = dim_at(a, i), bv = dim_at(b, i); + int64_t m = std::max({cv, av, bv}); + if ((cv != m && cv != 1) || (av != m && av != 1) || + (bv != m && bv != 1)) { + throw std::runtime_error( + "where(resize): operands not broadcast-compatible"); + } + out_d[i] = m; + } + g.set_cur_dims(out_id, out_d); + const uint32_t out_ndim = static_cast(r); + WebGPUTensor tc, ta, tb, to; + tc.dims = c; + ta.dims = a; + tb.dims = b; + to.dims = out_d; + TensorMeta om, cm, am, bm; + fill_tensor_meta_broadcast(to, out_ndim, &om); + fill_tensor_meta_broadcast(tc, out_ndim, &cm); + fill_tensor_meta_broadcast(ta, out_ndim, &am); + fill_tensor_meta_broadcast(tb, out_ndim, &bm); + wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om)); + wgpuQueueWriteBuffer(g.queue(), c_buf, 0, &cm, sizeof(cm)); + wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am)); + wgpuQueueWriteBuffer(g.queue(), bb_buf, 0, &bm, sizeof(bm)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), om.numel, wg_size, "where(resize)"); + }; + graph.add_tensor_resize_hook(cond_id, where_resize); + graph.add_tensor_resize_hook(a_id, where_resize); + graph.add_tensor_resize_hook(b_id, where_resize); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(out_meta_buf); + graph.own_uniform_buffer(cond_meta_buf); + graph.own_uniform_buffer(a_meta_buf); + graph.own_uniform_buffer(b_meta_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.where.self, where_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/where/where.wgsl b/backends/webgpu/runtime/ops/where/where.wgsl new file mode 100644 index 00000000000..eb32f314f83 --- /dev/null +++ b/backends/webgpu/runtime/ops/where/where.wgsl @@ -0,0 +1,49 @@ +@group(0) @binding(0) var cond: array; +@group(0) @binding(1) var input_a: array; +@group(0) @binding(2) var input_b: array; +@group(0) @binding(3) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(4) var out_meta: TensorMeta; +@group(0) @binding(5) var cond_meta: TensorMeta; +@group(0) @binding(6) var a_meta: TensorMeta; +@group(0) @binding(7) var b_meta: TensorMeta; + +override wg_size: u32 = 64u; + +// 1-byte bool packed 4-per-u32; extract byte i (mirrors q4gsw weight unpack). +fn cond_is_true(i: u32) -> bool { + let word = cond[i >> 2u]; + return ((word >> ((i & 3u) * 8u)) & 0xFFu) != 0u; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= out_meta.numel) { + return; + } + + var rem = idx; + var lc: u32 = 0u; + var la: u32 = 0u; + var lb: 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]; + lc = lc + min(coord, cond_meta.sizes[d] - 1u) * cond_meta.strides[d]; + la = la + min(coord, a_meta.sizes[d] - 1u) * a_meta.strides[d]; + lb = lb + min(coord, b_meta.sizes[d] - 1u) * b_meta.strides[d]; + } + + if (cond_is_true(lc)) { + output[idx] = input_a[la]; + } else { + output[idx] = input_b[lb]; + } +} diff --git a/backends/webgpu/runtime/ops/where/where_wgsl.h b/backends/webgpu/runtime/ops/where/where_wgsl.h new file mode 100644 index 00000000000..a8ef64fb169 --- /dev/null +++ b/backends/webgpu/runtime/ops/where/where_wgsl.h @@ -0,0 +1,73 @@ +/* + * 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 where.wgsl - DO NOT EDIT. +// wgsl-sha256: e02e6b9e7a073d2aac3d2c110ed5c382ab6663ee9656dc24f4b9c9fdd9e98a38 +inline constexpr const char* kWhereWGSL = R"( +@group(0) @binding(0) var cond: array; +@group(0) @binding(1) var input_a: array; +@group(0) @binding(2) var input_b: array; +@group(0) @binding(3) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: vec4, + strides: vec4, +} +@group(0) @binding(4) var out_meta: TensorMeta; +@group(0) @binding(5) var cond_meta: TensorMeta; +@group(0) @binding(6) var a_meta: TensorMeta; +@group(0) @binding(7) var b_meta: TensorMeta; + +override wg_size: u32 = 64u; + +// 1-byte bool packed 4-per-u32; extract byte i (mirrors q4gsw weight unpack). +fn cond_is_true(i: u32) -> bool { + let word = cond[i >> 2u]; + return ((word >> ((i & 3u) * 8u)) & 0xFFu) != 0u; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= out_meta.numel) { + return; + } + + var rem = idx; + var lc: u32 = 0u; + var la: u32 = 0u; + var lb: 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]; + lc = lc + min(coord, cond_meta.sizes[d] - 1u) * cond_meta.strides[d]; + la = la + min(coord, a_meta.sizes[d] - 1u) * a_meta.strides[d]; + lb = lb + min(coord, b_meta.sizes[d] - 1u) * b_meta.strides[d]; + } + + if (cond_is_true(lc)) { + output[idx] = input_a[la]; + } else { + output[idx] = input_b[lb]; + } +} +)"; + +inline constexpr uint32_t kWhereWorkgroupSizeX = 64; +inline constexpr uint32_t kWhereWorkgroupSizeY = 1; +inline constexpr uint32_t kWhereWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu