|
| 1 | +/* |
| 2 | + * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | + * All rights reserved. |
| 4 | + * |
| 5 | + * This source code is licensed under the BSD-style license found in the |
| 6 | + * LICENSE file in the root directory of this source tree. |
| 7 | + */ |
| 8 | + |
| 9 | +#include <executorch/backends/webgpu/runtime/WebGPUGraph.h> |
| 10 | +#include <executorch/backends/webgpu/runtime/WebGPUUtils.h> |
| 11 | +#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h> |
| 12 | +#include <executorch/backends/webgpu/runtime/ops/TensorMeta.h> |
| 13 | +#include <executorch/backends/webgpu/runtime/ops/cat/cat_wgsl.h> |
| 14 | + |
| 15 | +#include <webgpu/webgpu.h> |
| 16 | + |
| 17 | +#include <cstdint> |
| 18 | +#include <stdexcept> |
| 19 | +#include <vector> |
| 20 | + |
| 21 | +namespace executorch::backends::webgpu { |
| 22 | + |
| 23 | +namespace { |
| 24 | + |
| 25 | +struct CatParams { |
| 26 | + uint32_t concat_dim; |
| 27 | + uint32_t off_k; |
| 28 | + uint32_t _pad[2]; |
| 29 | +}; |
| 30 | +static_assert( |
| 31 | + sizeof(CatParams) == 16, |
| 32 | + "CatParams must match the WGSL Params uniform (16-byte aligned)"); |
| 33 | + |
| 34 | +// cat: 1 dispatch/input -> disjoint out slab at host off_k (Vulkan concat). |
| 35 | +void cat_impl(WebGPUGraph& graph, const std::vector<int>& args) { |
| 36 | + // args: [tensors (ValueList), dim, out]. |
| 37 | + const int list_id = args.at(0); |
| 38 | + const int out_id = args.at(args.size() - 1); |
| 39 | + |
| 40 | + if (graph.get_value_type(list_id) != WebGPUGraph::ValueType::ValueList) { |
| 41 | + throw std::runtime_error("cat: tensors arg is not a ValueList"); |
| 42 | + } |
| 43 | + if (graph.get_value_type(args.at(1)) != WebGPUGraph::ValueType::Int) { |
| 44 | + throw std::runtime_error("cat: dim arg is not a static Int"); |
| 45 | + } |
| 46 | + if (graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { |
| 47 | + throw std::runtime_error("cat: out arg is not a tensor"); |
| 48 | + } |
| 49 | + |
| 50 | + WGPUDevice device = graph.device(); |
| 51 | + const std::vector<int>& ids = graph.get_value_list(list_id); |
| 52 | + if (ids.empty()) { |
| 53 | + throw std::runtime_error("cat: empty input list"); |
| 54 | + } |
| 55 | + |
| 56 | + const auto& out_tensor = graph.get_tensor(out_id); |
| 57 | + const int ndim = static_cast<int>(out_tensor.dims.size()); |
| 58 | + |
| 59 | + int64_t dim = graph.get_int(args.at(1)); |
| 60 | + if (dim < 0) { |
| 61 | + dim += ndim; |
| 62 | + } |
| 63 | + if (dim < 0 || dim >= ndim) { |
| 64 | + throw std::runtime_error("cat: dim out of range"); |
| 65 | + } |
| 66 | + |
| 67 | + // Workgroup size is invariant across inputs: clamp once, share the constant. |
| 68 | + uint32_t wg_size = utils::clamp_workgroup_size(device, kCatWorkgroupSizeX); |
| 69 | + |
| 70 | + // Validate + cache input meta/wgc BEFORE any GPU alloc (no leak on throw). |
| 71 | + std::vector<TensorMeta> in_metas(ids.size()); |
| 72 | + std::vector<uint32_t> wg_counts(ids.size()); |
| 73 | + int64_t concat_sum = 0; |
| 74 | + for (size_t k = 0; k < ids.size(); k++) { |
| 75 | + const int id = ids[k]; |
| 76 | + if (graph.get_value_type(id) != WebGPUGraph::ValueType::Tensor) { |
| 77 | + throw std::runtime_error("cat: input list element is not a tensor"); |
| 78 | + } |
| 79 | + const auto& in_tensor = graph.get_tensor(id); |
| 80 | + if (static_cast<int>(in_tensor.dims.size()) != ndim) { |
| 81 | + throw std::runtime_error("cat: input rank != output rank"); |
| 82 | + } |
| 83 | + for (int d = 0; d < ndim; d++) { |
| 84 | + if (d != dim && in_tensor.dims[d] != out_tensor.dims[d]) { |
| 85 | + throw std::runtime_error("cat: non-concat dim size mismatch"); |
| 86 | + } |
| 87 | + } |
| 88 | + fill_tensor_meta(in_tensor, &in_metas[k]); |
| 89 | + if (in_tensor.nbytes != |
| 90 | + static_cast<size_t>(in_metas[k].numel) * sizeof(float)) { |
| 91 | + throw std::runtime_error("cat: non-fp32 input (nbytes != numel * 4)"); |
| 92 | + } |
| 93 | + wg_counts[k] = utils::compute_1d_workgroup_count( |
| 94 | + device, in_metas[k].numel, wg_size, "cat"); |
| 95 | + concat_sum += in_tensor.dims[dim]; |
| 96 | + } |
| 97 | + if (concat_sum != out_tensor.dims[dim]) { |
| 98 | + throw std::runtime_error("cat: concat dim sizes do not sum to output"); |
| 99 | + } |
| 100 | + |
| 101 | + TensorMeta out_meta; |
| 102 | + fill_tensor_meta(out_tensor, &out_meta); |
| 103 | + if (out_tensor.nbytes != |
| 104 | + static_cast<size_t>(out_meta.numel) * sizeof(float)) { |
| 105 | + throw std::runtime_error("cat: non-fp32 output (nbytes != numel * 4)"); |
| 106 | + } |
| 107 | + |
| 108 | + WGPUBuffer out_meta_buf = |
| 109 | + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); |
| 110 | + graph.add_uniform_buffer_bytes(sizeof(TensorMeta)); |
| 111 | + |
| 112 | + WGPUConstantEntry wg_size_constant = {}; |
| 113 | + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; |
| 114 | + wg_size_constant.value = static_cast<double>(wg_size); |
| 115 | + |
| 116 | + // Shared shader/layout; fresh pipeline+bind group per input (no double-free). |
| 117 | + WGPUShaderSourceWGSL wgsl_desc = {}; |
| 118 | + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; |
| 119 | + wgsl_desc.code = {kCatWGSL, WGPU_STRLEN}; |
| 120 | + WGPUShaderModuleDescriptor shader_desc = {}; |
| 121 | + shader_desc.nextInChain = &wgsl_desc.chain; |
| 122 | + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); |
| 123 | + |
| 124 | + WGPUBindGroupLayoutEntry entries[5] = {}; |
| 125 | + entries[0].binding = 0; |
| 126 | + entries[0].visibility = WGPUShaderStage_Compute; |
| 127 | + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; |
| 128 | + entries[1].binding = 1; |
| 129 | + entries[1].visibility = WGPUShaderStage_Compute; |
| 130 | + entries[1].buffer.type = WGPUBufferBindingType_Storage; |
| 131 | + entries[2].binding = 2; |
| 132 | + entries[2].visibility = WGPUShaderStage_Compute; |
| 133 | + entries[2].buffer.type = WGPUBufferBindingType_Uniform; |
| 134 | + entries[3].binding = 3; |
| 135 | + entries[3].visibility = WGPUShaderStage_Compute; |
| 136 | + entries[3].buffer.type = WGPUBufferBindingType_Uniform; |
| 137 | + entries[4].binding = 4; |
| 138 | + entries[4].visibility = WGPUShaderStage_Compute; |
| 139 | + entries[4].buffer.type = WGPUBufferBindingType_Uniform; |
| 140 | + |
| 141 | + WGPUBindGroupLayoutDescriptor bgl_desc = {}; |
| 142 | + bgl_desc.entryCount = 5; |
| 143 | + bgl_desc.entries = entries; |
| 144 | + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); |
| 145 | + |
| 146 | + WGPUPipelineLayoutDescriptor pl_desc = {}; |
| 147 | + pl_desc.bindGroupLayoutCount = 1; |
| 148 | + pl_desc.bindGroupLayouts = &bgl; |
| 149 | + WGPUPipelineLayout pipeline_layout = |
| 150 | + wgpuDeviceCreatePipelineLayout(device, &pl_desc); |
| 151 | + |
| 152 | + uint32_t off_k = 0; |
| 153 | + for (size_t k = 0; k < ids.size(); k++) { |
| 154 | + const auto& in_tensor = graph.get_tensor(ids[k]); |
| 155 | + |
| 156 | + CatParams params = {}; |
| 157 | + params.concat_dim = static_cast<uint32_t>(dim); |
| 158 | + params.off_k = off_k; |
| 159 | + |
| 160 | + WGPUBuffer in_meta_buf = |
| 161 | + utils::make_uniform(device, &in_metas[k], sizeof(TensorMeta)); |
| 162 | + WGPUBuffer params_buf = |
| 163 | + utils::make_uniform(device, ¶ms, sizeof(CatParams)); |
| 164 | + graph.add_uniform_buffer_bytes(sizeof(TensorMeta) + sizeof(CatParams)); |
| 165 | + |
| 166 | + WGPUComputePipelineDescriptor pipeline_desc = {}; |
| 167 | + pipeline_desc.layout = pipeline_layout; |
| 168 | + pipeline_desc.compute.module = shader; |
| 169 | + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; |
| 170 | + pipeline_desc.compute.constantCount = 1; |
| 171 | + pipeline_desc.compute.constants = &wg_size_constant; |
| 172 | + WGPUComputePipeline pipeline = |
| 173 | + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); |
| 174 | + |
| 175 | + WGPUBindGroupEntry bg_entries[5] = {}; |
| 176 | + bg_entries[0].binding = 0; |
| 177 | + bg_entries[0].buffer = in_tensor.buffer; |
| 178 | + bg_entries[0].size = in_tensor.nbytes; |
| 179 | + bg_entries[1].binding = 1; |
| 180 | + bg_entries[1].buffer = out_tensor.buffer; |
| 181 | + bg_entries[1].size = out_tensor.nbytes; |
| 182 | + bg_entries[2].binding = 2; |
| 183 | + bg_entries[2].buffer = out_meta_buf; |
| 184 | + bg_entries[2].size = sizeof(TensorMeta); |
| 185 | + bg_entries[3].binding = 3; |
| 186 | + bg_entries[3].buffer = in_meta_buf; |
| 187 | + bg_entries[3].size = sizeof(TensorMeta); |
| 188 | + bg_entries[4].binding = 4; |
| 189 | + bg_entries[4].buffer = params_buf; |
| 190 | + bg_entries[4].size = sizeof(CatParams); |
| 191 | + |
| 192 | + WGPUBindGroupDescriptor bg_desc = {}; |
| 193 | + bg_desc.layout = bgl; |
| 194 | + bg_desc.entryCount = 5; |
| 195 | + bg_desc.entries = bg_entries; |
| 196 | + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); |
| 197 | + |
| 198 | + graph.add_dispatch({pipeline, bind_group, wg_counts[k]}); |
| 199 | + // Drop our refs; this input's bind group keeps its uniforms alive. |
| 200 | + wgpuBufferRelease(in_meta_buf); |
| 201 | + wgpuBufferRelease(params_buf); |
| 202 | + off_k += static_cast<uint32_t>(in_tensor.dims[dim]); |
| 203 | + } |
| 204 | + |
| 205 | + wgpuShaderModuleRelease(shader); |
| 206 | + wgpuBindGroupLayoutRelease(bgl); |
| 207 | + wgpuPipelineLayoutRelease(pipeline_layout); |
| 208 | + // Drop our ref to the shared out_meta; the bind groups keep it alive. |
| 209 | + wgpuBufferRelease(out_meta_buf); |
| 210 | +} |
| 211 | + |
| 212 | +} // namespace |
| 213 | + |
| 214 | +WEBGPU_REGISTER_OPERATORS { |
| 215 | + WEBGPU_REGISTER_OP(aten.cat.default, cat_impl); |
| 216 | +} |
| 217 | + |
| 218 | +} // namespace executorch::backends::webgpu |
0 commit comments