[ExecuTorch][WebGPU] Add dequantize_per_tensor op (int8 buffer path)

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 <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+#include <vector>
+
+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<int>& 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<uint64_t>(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<float>(scale);
+  params.zero_point = static_cast<int32_t>(zero_point);
+  params.numel = static_cast<uint32_t>(numel);
+
+  const uint32_t num_words = static_cast<uint32_t>(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<double>(wg_size);
+
+  WGPUBuffer params_buf =
+      utils::make_uniform(device, &params, 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

backends/webgpu/runtime/ops/dequantize/dequantize_per_tensor.wgsl

@@ -0,0 +1,29 @@
+@group(0) @binding(0) var<storage, read> t_in: array<u32>;
+@group(0) @binding(1) var<storage, read_write> t_out: array<f32>;
+
+struct Params {
+  scale: f32,
+  zero_point: i32,
+  numel: u32,
+  pad0: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 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;
+    }
+}

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 <cstdint>
+
+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<storage, read> t_in: array<u32>;
+@group(0) @binding(1) var<storage, read_write> t_out: array<f32>;
+
+struct Params {
+  scale: f32,
+  zero_point: i32,
+  numel: u32,
+  pad0: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 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