diff --git a/backends/webgpu/runtime/ops/amax/Reduce.cpp b/backends/webgpu/runtime/ops/amax/Reduce.cpp new file mode 100644 index 00000000000..702aa21c218 --- /dev/null +++ b/backends/webgpu/runtime/ops/amax/Reduce.cpp @@ -0,0 +1,190 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct AmaxParams { + uint32_t num_rows; + uint32_t reduce_size; + uint32_t _pad[2]; +}; + +// Last-dim reduction; mirrors Vulkan add_reduce_per_row_node. +void amax_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.amax.default args: [in, dim, keepdim, out] + const int in_id = args.at(0); + const int out_id = args.at(3); + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("amax: null buffer binding"); + } + if (in_tensor.is_int || out_tensor.is_int) { + throw std::runtime_error("amax: int dtype unsupported"); + } + + const std::vector& dims = graph.get_int_list(args.at(1)); + const int64_t ndim = static_cast(in_tensor.dims.size()); + if (dims.size() != 1 || (dims[0] != -1 && dims[0] != ndim - 1)) { + throw std::runtime_error("amax: only last-dim reduction is supported"); + } + + const uint32_t reduce_size = static_cast(in_tensor.dims.back()); + const uint32_t num_rows = + static_cast(out_tensor.nbytes / sizeof(float)); + if (reduce_size == 0u || + in_tensor.nbytes / sizeof(float) != + static_cast(num_rows) * reduce_size) { + throw std::runtime_error("amax: shape mismatch (num_rows * reduce_size)"); + } + + uint32_t wg_size = utils::clamp_workgroup_size(device, kAmaxWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_rows, wg_size, "amax"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + AmaxParams params = {}; + params.num_rows = num_rows; + params.reduce_size = reduce_size; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(AmaxParams)); + graph.add_uniform_buffer_bytes(sizeof(AmaxParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kAmaxWGSL, WGPU_STRLEN}; + + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group layout: input (read storage) + output (storage) + params. + WGPUBindGroupLayoutEntry entries[3] = {}; + + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + 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 = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + + bg_entries[2].binding = 2; + bg_entries[2].buffer = uniform_buffer; + bg_entries[2].size = sizeof(AmaxParams); + + 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.x, "amax", workgroup_count.y}); + + // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. + const bool keepdim = graph.get_bool(args.at(2)); + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, keepdim, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const uint32_t rsize = static_cast(d.back()); + if (rsize == 0u) { + throw std::runtime_error("amax(resize): zero reduce dim"); + } + const uint64_t total = utils::numel_of(d); + const uint32_t rows = static_cast(total / rsize); + std::vector od = d; + if (keepdim) { + od.back() = 1; + } else { + od.pop_back(); + } + g.set_cur_dims(out_id, od); + AmaxParams p = {}; + p.num_rows = rows; + p.reduce_size = rsize; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), rows, wg_size, "amax"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); + + // Release intermediates (pipeline + bind_group are kept by dispatch). + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.amax.default, amax_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/amax/amax.wgsl b/backends/webgpu/runtime/ops/amax/amax.wgsl new file mode 100644 index 00000000000..2ff33e0c20f --- /dev/null +++ b/backends/webgpu/runtime/ops/amax/amax.wgsl @@ -0,0 +1,26 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = max(acc, input[base + j]); + } + output[row] = acc; +} diff --git a/backends/webgpu/runtime/ops/amax/amax_wgsl.h b/backends/webgpu/runtime/ops/amax/amax_wgsl.h new file mode 100644 index 00000000000..1f82d7507cd --- /dev/null +++ b/backends/webgpu/runtime/ops/amax/amax_wgsl.h @@ -0,0 +1,50 @@ +/* + * 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 amax.wgsl - DO NOT EDIT. +// wgsl-sha256: aae058ed0c432ac0cb54ea894e03e12a246d0d1c743d9ebb995b9773029e4652 +inline constexpr const char* kAmaxWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = max(acc, input[base + j]); + } + output[row] = acc; +} +)"; + +inline constexpr uint32_t kAmaxWorkgroupSizeX = 256; +inline constexpr uint32_t kAmaxWorkgroupSizeY = 1; +inline constexpr uint32_t kAmaxWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/amin/Reduce.cpp b/backends/webgpu/runtime/ops/amin/Reduce.cpp new file mode 100644 index 00000000000..35c7017436d --- /dev/null +++ b/backends/webgpu/runtime/ops/amin/Reduce.cpp @@ -0,0 +1,190 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the WGSL Params struct; 16-byte aligned. +struct AminParams { + uint32_t num_rows; + uint32_t reduce_size; + uint32_t _pad[2]; +}; + +// Last-dim reduction; mirrors Vulkan add_reduce_per_row_node. +void amin_impl(WebGPUGraph& graph, const std::vector& args) { + // aten.amin.default args: [in, dim, keepdim, out] + const int in_id = args.at(0); + const int out_id = args.at(3); + + WGPUDevice device = graph.device(); + + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("amin: null buffer binding"); + } + if (in_tensor.is_int || out_tensor.is_int) { + throw std::runtime_error("amin: int dtype unsupported"); + } + + const std::vector& dims = graph.get_int_list(args.at(1)); + const int64_t ndim = static_cast(in_tensor.dims.size()); + if (dims.size() != 1 || (dims[0] != -1 && dims[0] != ndim - 1)) { + throw std::runtime_error("amin: only last-dim reduction is supported"); + } + + const uint32_t reduce_size = static_cast(in_tensor.dims.back()); + const uint32_t num_rows = + static_cast(out_tensor.nbytes / sizeof(float)); + if (reduce_size == 0u || + in_tensor.nbytes / sizeof(float) != + static_cast(num_rows) * reduce_size) { + throw std::runtime_error("amin: shape mismatch (num_rows * reduce_size)"); + } + + uint32_t wg_size = utils::clamp_workgroup_size(device, kAminWorkgroupSizeX); + utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_rows, wg_size, "amin"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + AminParams params = {}; + params.num_rows = num_rows; + params.reduce_size = reduce_size; + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(AminParams)); + graph.add_uniform_buffer_bytes(sizeof(AminParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kAminWGSL, WGPU_STRLEN}; + + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + // Bind group layout: input (read storage) + output (storage) + params. + WGPUBindGroupLayoutEntry entries[3] = {}; + + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_Storage; + + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + 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 = in_tensor.buffer; + bg_entries[0].size = in_tensor.nbytes; + + bg_entries[1].binding = 1; + bg_entries[1].buffer = out_tensor.buffer; + bg_entries[1].size = out_tensor.nbytes; + + bg_entries[2].binding = 2; + bg_entries[2].buffer = uniform_buffer; + bg_entries[2].size = sizeof(AminParams); + + 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.x, "amin", workgroup_count.y}); + + // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. + const bool keepdim = graph.get_bool(args.at(2)); + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, keepdim, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + const uint32_t rsize = static_cast(d.back()); + if (rsize == 0u) { + throw std::runtime_error("amin(resize): zero reduce dim"); + } + const uint64_t total = utils::numel_of(d); + const uint32_t rows = static_cast(total / rsize); + std::vector od = d; + if (keepdim) { + od.back() = 1; + } else { + od.pop_back(); + } + g.set_cur_dims(out_id, od); + AminParams p = {}; + p.num_rows = rows; + p.reduce_size = rsize; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), rows, wg_size, "amin"); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); + + // Release intermediates (pipeline + bind_group are kept by dispatch). + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.amin.default, amin_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/amin/amin.wgsl b/backends/webgpu/runtime/ops/amin/amin.wgsl new file mode 100644 index 00000000000..d8f4346d231 --- /dev/null +++ b/backends/webgpu/runtime/ops/amin/amin.wgsl @@ -0,0 +1,26 @@ +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = min(acc, input[base + j]); + } + output[row] = acc; +} diff --git a/backends/webgpu/runtime/ops/amin/amin_wgsl.h b/backends/webgpu/runtime/ops/amin/amin_wgsl.h new file mode 100644 index 00000000000..53f94e2f0e8 --- /dev/null +++ b/backends/webgpu/runtime/ops/amin/amin_wgsl.h @@ -0,0 +1,50 @@ +/* + * 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 amin.wgsl - DO NOT EDIT. +// wgsl-sha256: 974a28fd80f089c8a52bf54d73f3bd03c195b2f9d7904bb4c31e6545813e0459 +inline constexpr const char* kAminWGSL = R"( +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_rows: u32, + reduce_size: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.num_rows) { + return; + } + let base = row * params.reduce_size; + var acc = input[base]; + for (var j = 1u; j < params.reduce_size; j = j + 1u) { + acc = min(acc, input[base + j]); + } + output[row] = acc; +} +)"; + +inline constexpr uint32_t kAminWorkgroupSizeX = 256; +inline constexpr uint32_t kAminWorkgroupSizeY = 1; +inline constexpr uint32_t kAminWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu