From bfdc3e22cf12eac125867c6d31d45f1645433ca6 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 22 Jul 2026 15:32:11 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../ops/dequantize/DequantizePerTensor.cpp | 172 ++++++++++++++++++ .../ops/dequantize/dequantize_per_tensor.wgsl | 29 +++ .../dequantize/dequantize_per_tensor_wgsl.h | 53 ++++++ 3 files changed, 254 insertions(+) create mode 100644 backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp create mode 100644 backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl create mode 100644 backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h diff --git a/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp b/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp new file mode 100644 index 00000000000..83f9b62ce5a --- /dev/null +++ b/backends/webgpu/runtime/ops/dequantize/DequantizePerTensor.cpp @@ -0,0 +1,172 @@ +/* + * 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 { + +struct DequantParams { + float scale; + int32_t zero_point; + uint32_t numel; + uint32_t pad0; +}; +static_assert( + sizeof(DequantParams) == 16, + "DequantParams must match the WGSL Params struct (16 bytes)"); + +// int8->fp32; mirrors Vulkan q8ta_dequantize.glsl ((q-zp)*scale). +void dequantize_per_tensor_impl( + WebGPUGraph& graph, + const std::vector& args) { + // args: [q, scale, zp, ...]; out is always args.back() (skips out_dtype). + const int in_id = args.at(0); + const int out_id = args.at(args.size() - 1); + + if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor || + graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) { + throw std::runtime_error("dequantize_per_tensor: in/out is not a tensor"); + } + + 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("dequantize_per_tensor: null buffer binding"); + } + + const double scale = graph.get_double(args.at(1)); + const int zero_point = graph.get_int(args.at(2)); + + uint64_t numel = 1; + for (int64_t d : out_tensor.dims) { + numel *= static_cast(d); + } + if (numel == 0 || numel % 4 != 0) { + throw std::runtime_error( + "dequantize_per_tensor: numel must be a nonzero " + "multiple of 4"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("dequantize_per_tensor: numel exceeds u32"); + } + // int8 input (not uint8/bool): the kernel sign-extends assuming int8. + if (!in_tensor.is_int8 || in_tensor.nbytes != numel) { + throw std::runtime_error("dequantize_per_tensor: input is not int8"); + } + if (out_tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error("dequantize_per_tensor: output is not fp32"); + } + + DequantParams params = {}; + params.scale = static_cast(scale); + params.zero_point = static_cast(zero_point); + params.numel = static_cast(numel); + + const uint32_t num_words = static_cast(numel / 4); + uint32_t wg_size = + utils::clamp_workgroup_size(device, kDequantizePerTensorWorkgroupSizeX); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, num_words, wg_size, "dequantize_per_tensor"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(DequantParams)); + graph.add_uniform_buffer_bytes(sizeof(DequantParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kDequantizePerTensorWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + 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 = params_buf; + bg_entries[2].size = sizeof(DequantParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 3; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, + bind_group, + workgroup_count.x, + "dequantize_per_tensor", + workgroup_count.y}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(params_buf); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP( + quantized_decomposed.dequantize_per_tensor.default, + dequantize_per_tensor_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl new file mode 100644 index 00000000000..ab5619dfa13 --- /dev/null +++ b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl @@ -0,0 +1,29 @@ +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + scale: f32, + zero_point: i32, + numel: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +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) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let word = t_in[widx]; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let b = (word >> (j * 8u)) & 0xFFu; + let s = i32(b ^ 0x80u) - 128; + t_out[widx * 4u + j] = f32(s - params.zero_point) * params.scale; + } +} diff --git a/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h new file mode 100644 index 00000000000..a5e2c900fde --- /dev/null +++ b/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h @@ -0,0 +1,53 @@ +/* + * 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 dequantize_per_tensor.wgsl - DO NOT EDIT. +// wgsl-sha256: b9e24b5f3b57f6eb842838399985ab2de20513319387fa3341624d310687f90e +inline constexpr const char* kDequantizePerTensorWGSL = R"( +@group(0) @binding(0) var t_in: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + scale: f32, + zero_point: i32, + numel: u32, + pad0: u32, +} +@group(0) @binding(2) var params: Params; + +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) { + // One thread per packed u32 word (4 int8 elems); 2D-fold lifts the 65535 cap. + let widx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (widx >= params.numel / 4u) { + return; + } + let word = t_in[widx]; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let b = (word >> (j * 8u)) & 0xFFu; + let s = i32(b ^ 0x80u) - 128; + t_out[widx * 4u + j] = f32(s - params.zero_point) * params.scale; + } +} +)"; + +inline constexpr uint32_t kDequantizePerTensorWorkgroupSizeX = 64; +inline constexpr uint32_t kDequantizePerTensorWorkgroupSizeY = 1; +inline constexpr uint32_t kDequantizePerTensorWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu