|
| 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/arange/arange_int_wgsl.h> |
| 13 | +#include <executorch/backends/webgpu/runtime/ops/arange/arange_wgsl.h> |
| 14 | + |
| 15 | +#include <webgpu/webgpu.h> |
| 16 | + |
| 17 | +#include <cstdint> |
| 18 | +#include <stdexcept> |
| 19 | +#include <string> |
| 20 | +#include <vector> |
| 21 | + |
| 22 | +namespace executorch::backends::webgpu { |
| 23 | + |
| 24 | +namespace { |
| 25 | + |
| 26 | +// Uniform layout shared by both shaders (16 bytes). start/step are |
| 27 | +// reinterpreted as i32 by the integer variant, so keep the slots the same |
| 28 | +// width. |
| 29 | +struct ArangeParams { |
| 30 | + uint32_t num_elements; |
| 31 | + union { |
| 32 | + float f; |
| 33 | + int32_t i; |
| 34 | + } start; |
| 35 | + union { |
| 36 | + float f; |
| 37 | + int32_t i; |
| 38 | + } step; |
| 39 | + uint32_t _pad; |
| 40 | +}; |
| 41 | +static_assert( |
| 42 | + sizeof(ArangeParams) % 16u == 0u, |
| 43 | + "WebGPU parameter blocks must have a 16-byte-aligned size"); |
| 44 | + |
| 45 | +// start/end/step arrive as static Int or Double scalars, or as Null when the |
| 46 | +// caller omitted them, in which case aten's defaults apply. A SymInt would make |
| 47 | +// the output length dynamic, which the fixed dispatch below cannot honor. |
| 48 | +// if_omitted == nullptr marks an argument that has no default and must be |
| 49 | +// present. |
| 50 | +double read_scalar( |
| 51 | + WebGPUGraph& graph, |
| 52 | + int id, |
| 53 | + const char* what, |
| 54 | + const double* if_omitted) { |
| 55 | + switch (graph.get_value_type(id)) { |
| 56 | + case WebGPUGraph::ValueType::Int: |
| 57 | + return static_cast<double>(graph.get_int(id)); |
| 58 | + case WebGPUGraph::ValueType::Double: |
| 59 | + return graph.get_double(id); |
| 60 | + case WebGPUGraph::ValueType::Null: |
| 61 | + if (if_omitted == nullptr) { |
| 62 | + throw std::runtime_error( |
| 63 | + std::string("arange: ") + what + " is required"); |
| 64 | + } |
| 65 | + return *if_omitted; |
| 66 | + default: |
| 67 | + throw std::runtime_error( |
| 68 | + std::string("arange: dynamic/unsupported ") + what); |
| 69 | + } |
| 70 | +} |
| 71 | + |
| 72 | +// aten.arange.start_step(start, end, step, *, dtype, layout, device, |
| 73 | +// pin_memory) args: [start, end, step, ..., out]. Writes start + i * step over |
| 74 | +// the pre-allocated output; the element count comes from the output tensor |
| 75 | +// rather than from (end - start) / step so it stays consistent with memory |
| 76 | +// planning. |
| 77 | +void arange_impl(WebGPUGraph& graph, const std::vector<int>& args) { |
| 78 | + if (args.size() < 4) { |
| 79 | + throw std::runtime_error("arange: expected at least 4 args"); |
| 80 | + } |
| 81 | + const int out_id = args.at(args.size() - 1); |
| 82 | + |
| 83 | + WGPUDevice device = graph.device(); |
| 84 | + const auto& out_tensor = graph.get_tensor(out_id); |
| 85 | + // Two 4-byte variants: fp32 and i32. A bool output is 1-byte packed and has |
| 86 | + // no meaningful arange, so reject it rather than mis-stride the buffer. |
| 87 | + if (out_tensor.is_bool || out_tensor.elem_size != 4u) { |
| 88 | + throw std::runtime_error("arange: only 4-byte fp32 or int32 output"); |
| 89 | + } |
| 90 | + const bool is_int = out_tensor.is_int; |
| 91 | + |
| 92 | + constexpr double kStartDefault = 0.0; |
| 93 | + constexpr double kStepDefault = 1.0; |
| 94 | + const double start = read_scalar(graph, args.at(0), "start", &kStartDefault); |
| 95 | + const double step = read_scalar(graph, args.at(2), "step", &kStepDefault); |
| 96 | + // end does not feed the dispatch (the length comes from the output tensor), |
| 97 | + // but it must still be static: a SymInt end changes the live length, and |
| 98 | + // nothing here recomputes the output dims the way Vulkan's resize does, so a |
| 99 | + // dynamic end would silently keep emitting the max-sized output. |
| 100 | + (void)read_scalar(graph, args.at(1), "end", nullptr); |
| 101 | + if (step == 0.0) { |
| 102 | + throw std::runtime_error("arange: step must be non-zero"); |
| 103 | + } |
| 104 | + |
| 105 | + // An empty range (e.g. arange(0, 0)) is valid and has nothing to write. |
| 106 | + const uint64_t numel = out_tensor.nbytes / 4u; |
| 107 | + if (numel == 0) { |
| 108 | + return; |
| 109 | + } |
| 110 | + if (numel > UINT32_MAX) { |
| 111 | + throw std::runtime_error("arange: output numel exceeds u32"); |
| 112 | + } |
| 113 | + if (out_tensor.buffer == nullptr) { |
| 114 | + throw std::runtime_error("arange: null output buffer"); |
| 115 | + } |
| 116 | + const uint32_t num_elements = static_cast<uint32_t>(numel); |
| 117 | + |
| 118 | + uint32_t wg_size = utils::clamp_workgroup_size(device, kArangeWorkgroupSizeX); |
| 119 | + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( |
| 120 | + device, num_elements, wg_size, "arange"); |
| 121 | + |
| 122 | + WGPUConstantEntry wg_size_constant = {}; |
| 123 | + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; |
| 124 | + wg_size_constant.value = static_cast<double>(wg_size); |
| 125 | + |
| 126 | + ArangeParams params = {}; |
| 127 | + params.num_elements = num_elements; |
| 128 | + if (is_int) { |
| 129 | + // Truncating start/step before iterating is torch's integer semantics, not |
| 130 | + // an approximation: arange(0.5, 10.0, 1.5, dtype=int32) is [0, 1, ..., 6] |
| 131 | + // in eager, and a step that truncates to 0 yields all zeros there too. |
| 132 | + // Casting start + i * step per element would give [0, 2, 3, 5, 6, 8, 9] |
| 133 | + // instead. |
| 134 | + params.start.i = static_cast<int32_t>(start); |
| 135 | + params.step.i = static_cast<int32_t>(step); |
| 136 | + } else { |
| 137 | + params.start.f = static_cast<float>(start); |
| 138 | + params.step.f = static_cast<float>(step); |
| 139 | + } |
| 140 | + WGPUBuffer uniform_buffer = |
| 141 | + utils::make_uniform(device, ¶ms, sizeof(ArangeParams)); |
| 142 | + graph.add_uniform_buffer_bytes(sizeof(ArangeParams)); |
| 143 | + |
| 144 | + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( |
| 145 | + device, |
| 146 | + is_int ? kArangeIntWGSL : kArangeWGSL, |
| 147 | + { |
| 148 | + {0, |
| 149 | + WGPUBufferBindingType_Storage, |
| 150 | + out_tensor.buffer, |
| 151 | + out_tensor.nbytes}, |
| 152 | + {1, |
| 153 | + WGPUBufferBindingType_Uniform, |
| 154 | + uniform_buffer, |
| 155 | + sizeof(ArangeParams)}, |
| 156 | + }, |
| 157 | + &wg_size_constant, |
| 158 | + 1); |
| 159 | + |
| 160 | + const size_t dispatch_idx = graph.add_dispatch( |
| 161 | + {bundle.pipeline, |
| 162 | + bundle.bind_group, |
| 163 | + workgroup_count.x, |
| 164 | + "", |
| 165 | + workgroup_count.y}); |
| 166 | + |
| 167 | + // Dynamic shapes: recompute num_elements/dispatch from the live output dims. |
| 168 | + WGPUBuffer params_buf = uniform_buffer; |
| 169 | + graph.add_tensor_resize_hook( |
| 170 | + out_id, |
| 171 | + [out_id, params, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { |
| 172 | + const auto& d = g.cur_dims(out_id); |
| 173 | + const uint64_t n = utils::numel_of(d); |
| 174 | + ArangeParams p = params; |
| 175 | + p.num_elements = static_cast<uint32_t>(n); |
| 176 | + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); |
| 177 | + const utils::WgCount wgc = utils::compute_2d_workgroup_count( |
| 178 | + g.device(), static_cast<uint32_t>(n), wg_size, "arange(resize)"); |
| 179 | + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; |
| 180 | + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; |
| 181 | + }); |
| 182 | + |
| 183 | + graph.own_uniform_buffer(uniform_buffer); |
| 184 | +} |
| 185 | + |
| 186 | +} // namespace |
| 187 | + |
| 188 | +WEBGPU_REGISTER_OPERATORS { |
| 189 | + WEBGPU_REGISTER_OP(aten.arange.start_step, arange_impl); |
| 190 | +} |
| 191 | + |
| 192 | +} // namespace executorch::backends::webgpu |
0 commit comments