webgpu_backend.h

#pragma once
 
// Optional WebGPU C backend. The consumer supplies and includes the stable
// <webgpu/webgpu.h> API before this header and links Dawn/Emdawnwebgpu (or a
// compatible implementation). CPU-only builds never parse or link this code.
 
#if defined(TESS_ENABLE_WEBGPU)
 
#ifndef WEBGPU_H_
#error "tess/gpu/webgpu_backend.h requires <webgpu/webgpu.h> first"
#endif
 
#include <tess/gpu/backend.h>
 
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <new>
#include <optional>
#include <type_traits>
#include <vector>
 
namespace tess::gpu {

struct WebGpuBackendConfig {
  std::uint64_t max_buffer_bytes = std::uint64_t{256} * 1024u * 1024u;
  std::uint64_t max_dispatch_chunks = 65535;
  // WebGPU guarantees at least 65,535 workgroups in one dispatch dimension.
  // Supply a lower adapter/device limit here when one is imposed externally.
  std::uint32_t max_dispatch_workgroups_x = 65535;
  std::uint64_t max_inflight_readback_bytes = std::uint64_t{4} * 1024u * 1024u;
  std::size_t field_capacity = 16;
  std::size_t product_capacity = 64;
  bool allow_full_field_readback = false;
};

enum class WebGpuReadbackStatus : std::uint8_t {
  Complete,
  Failed,
};

using WebGpuReadbackCallback = void (*)(GpuProductHandle, WebGpuReadbackStatus,
                                        const void*, std::size_t,
                                        void*) noexcept;

struct WebGpuProductDesc {
  std::uint64_t product_key = 0;
  std::uint32_t input_field_index = 0;
  WGPUComputePipeline pipeline = nullptr;
  WGPUBindGroup bind_group = nullptr;
  WGPUBuffer readback_source = nullptr;
  std::uint64_t readback_source_offset = 0;
  std::uint64_t readback_byte_size = 0;
  WebGpuReadbackCallback readback_callback = nullptr;
  void* readback_userdata = nullptr;
};
 
namespace detail {
 
struct WebGpuSharedState {
  std::atomic<bool> available{true};
  std::atomic<std::uint64_t> inflight_readback_bytes{0};
  std::uint64_t max_inflight_readback_bytes = 0;
};
 
struct WebGpuReadbackOperation {
  WGPUBuffer staging = nullptr;
  GpuProductHandle handle{};
  std::size_t byte_size = 0;
  WebGpuReadbackCallback callback = nullptr;
  void* userdata = nullptr;
  std::shared_ptr<WebGpuSharedState> shared;
};
 
template <typename Size>
[[nodiscard]] constexpr bool fits_size(std::uint64_t value) noexcept {
  static_assert(std::is_integral_v<Size> && std::is_unsigned_v<Size>);
  if constexpr (std::numeric_limits<Size>::digits >=
                std::numeric_limits<std::uint64_t>::digits) {
    return true;
  } else {
    return value <=
           static_cast<std::uint64_t>(std::numeric_limits<Size>::max());
  }
}
 
inline void webgpu_readback_complete(WGPUMapAsyncStatus status, WGPUStringView,
                                     void* userdata1, void*) noexcept {
  auto* operation = static_cast<WebGpuReadbackOperation*>(userdata1);
  const void* data = nullptr;
  auto result = WebGpuReadbackStatus::Failed;
  // This acquire load is the linearization point against device notifications:
  // whichever atomic operation wins decides whether this already-submitted
  // result is still safe to expose to the application callback.
  const auto device_available =
      operation->shared->available.load(std::memory_order_acquire);
  if (status == WGPUMapAsyncStatus_Success && device_available) {
    data = wgpuBufferGetConstMappedRange(operation->staging, 0,
                                         operation->byte_size);
    if (data != nullptr) {
      result = WebGpuReadbackStatus::Complete;
    }
  }
  operation->callback(operation->handle, result, data,
                      data == nullptr ? 0 : operation->byte_size,
                      operation->userdata);
  if (status == WGPUMapAsyncStatus_Success) {
    wgpuBufferUnmap(operation->staging);
  }
  wgpuBufferRelease(operation->staging);
  operation->shared->inflight_readback_bytes.fetch_sub(
      operation->byte_size, std::memory_order_relaxed);
  delete operation;
}
 
[[nodiscard]] inline bool reserve_readback_bytes(
    const std::shared_ptr<WebGpuSharedState>& shared,
    std::uint64_t bytes) noexcept {
  auto current =
      shared->inflight_readback_bytes.load(std::memory_order_relaxed);
  for (;;) {
    if (bytes > shared->max_inflight_readback_bytes - current) {
      return false;
    }
    if (shared->inflight_readback_bytes.compare_exchange_weak(
            current, current + bytes, std::memory_order_relaxed)) {
      return true;
    }
  }
}
 
}  // namespace detail

class WebGpuBackend {
 public:
  explicit WebGpuBackend(WGPUDevice device, WebGpuBackendConfig config = {})
      : device_(device),
        config_(config),
        shared_(std::make_shared<detail::WebGpuSharedState>()) {
    shared_->max_inflight_readback_bytes = config_.max_inflight_readback_bytes;
    fields_.reserve(config_.field_capacity);
    products_.reserve(config_.product_capacity);
    if (device_ == nullptr || config_.max_buffer_bytes == 0 ||
        config_.max_dispatch_chunks == 0 ||
        config_.max_dispatch_workgroups_x == 0) {
      shared_->available.store(false, std::memory_order_relaxed);
      // No reference was retained on this disabled construction path.
      device_ = nullptr;
      return;
    }
    wgpuDeviceAddRef(device_);
    queue_ = wgpuDeviceGetQueue(device_);
    if (queue_ == nullptr) {
      shared_->available.store(false, std::memory_order_relaxed);
    }
  }
 
  WebGpuBackend(const WebGpuBackend&) = delete;
  auto operator=(const WebGpuBackend&) -> WebGpuBackend& = delete;
  WebGpuBackend(WebGpuBackend&&) = delete;
  auto operator=(WebGpuBackend&&) -> WebGpuBackend& = delete;

  ~WebGpuBackend() {
    for (auto& product : products_) {
      release_product(product);
    }
    for (auto& field : fields_) {
      if (field.buffer != nullptr) {
        wgpuBufferRelease(field.buffer);
      }
    }
    if (queue_ != nullptr) {
      wgpuQueueRelease(queue_);
    }
    if (device_ != nullptr) {
      wgpuDeviceRelease(device_);
    }
  }

  [[nodiscard]] auto capabilities() const noexcept -> GpuCapabilities {
    if (!available()) {
      return {};
    }
    return GpuCapabilities{
        .compute = true,
        .async_dispatch = true,
        .async_readback = true,
        .max_buffer_bytes = config_.max_buffer_bytes,
        .max_dispatch_chunks = config_.max_dispatch_chunks,
        .buffer_alignment = 4,
    };
  }

  [[nodiscard]] bool register_field(FieldMirrorDesc desc) {
    if (!available() || !desc.total_bytes_fits() || desc.total_bytes() == 0 ||
        desc.total_bytes() > config_.max_buffer_bytes ||
        fields_.size() >= config_.field_capacity ||
        find_field(desc.field_index) != nullptr) {
      return false;
    }
    const auto bytes = aligned_size(desc.total_bytes());
    if (!bytes.has_value() || *bytes > config_.max_buffer_bytes) {
      return false;
    }
    auto buffer_desc = WGPU_BUFFER_DESCRIPTOR_INIT;
    buffer_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst;
    buffer_desc.size = *bytes;
    auto buffer = wgpuDeviceCreateBuffer(device_, &buffer_desc);
    if (buffer == nullptr) {
      return false;
    }
    fields_.push_back(FieldSlot{desc, buffer, *bytes});
    return true;
  }

  [[nodiscard]] auto field_buffer(std::uint32_t field_index) const noexcept
      -> WGPUBuffer {
    const auto* field = find_field(field_index);
    return field == nullptr ? nullptr : field->buffer;
  }

  [[nodiscard]] auto register_product(WebGpuProductDesc desc)
      -> std::optional<GpuProductHandle> {
    if (!available() || desc.product_key == 0 || desc.pipeline == nullptr ||
        desc.bind_group == nullptr ||
        find_product(desc.product_key) != nullptr) {
      return std::nullopt;
    }
    if (desc.readback_source == nullptr) {
      // Source-less products do not perform readback. Reject orphaned callback
      // metadata instead of accepting a configuration whose callback can
      // never fire.
      if (desc.readback_byte_size != 0 || desc.readback_source_offset != 0 ||
          desc.readback_callback != nullptr ||
          desc.readback_userdata != nullptr) {
        return std::nullopt;
      }
    } else {
      if (desc.readback_byte_size == 0 || desc.readback_callback == nullptr ||
          (wgpuBufferGetUsage(desc.readback_source) &
           WGPUBufferUsage_CopySrc) == 0 ||
          (desc.readback_source_offset & 3u) != 0 ||
          (desc.readback_byte_size & 3u) != 0 ||
          desc.readback_source_offset >
              wgpuBufferGetSize(desc.readback_source) ||
          desc.readback_byte_size > wgpuBufferGetSize(desc.readback_source) -
                                        desc.readback_source_offset) {
        return std::nullopt;
      }
    }
    ProductSlot* slot = nullptr;
    for (auto& candidate : products_) {
      if (!candidate.active) {
        slot = &candidate;
        break;
      }
    }
    if (slot == nullptr) {
      if (products_.size() >= config_.product_capacity) {
        return std::nullopt;
      }
      products_.emplace_back();
      slot = &products_.back();
    }
    wgpuComputePipelineAddRef(desc.pipeline);
    wgpuBindGroupAddRef(desc.bind_group);
    if (desc.readback_source != nullptr) {
      wgpuBufferAddRef(desc.readback_source);
    }
    slot->active = true;
    slot->generation = ++generation_clock_;
    if (slot->generation == 0) {
      slot->generation = ++generation_clock_;
    }
    slot->desc = desc;
    return GpuProductHandle{desc.product_key, slot->generation};
  }

  bool unregister_product(GpuProductHandle handle) noexcept {
    auto* product = find_product(handle.key);
    if (product == nullptr || product->generation != handle.generation) {
      return false;
    }
    release_product(*product);
    return true;
  }

  [[nodiscard]] bool valid(GpuProductHandle handle) const noexcept {
    const auto* product = find_product(handle.key);
    return product != nullptr && product->generation == handle.generation;
  }

  [[nodiscard]] bool upload(const UploadDesc& upload) {
    const auto* field = find_field(upload.field_index);
    if (!available() || field == nullptr || upload.data == nullptr ||
        upload.byte_size == 0 || (upload.buffer_offset & 3u) != 0 ||
        (upload.byte_size & 3u) != 0 ||
        !detail::fits_size<std::size_t>(upload.byte_size) ||
        upload.buffer_offset > field->allocated_bytes ||
        upload.byte_size > field->allocated_bytes - upload.buffer_offset) {
      return false;
    }
    wgpuQueueWriteBuffer(queue_, field->buffer, upload.buffer_offset,
                         upload.data,
                         static_cast<std::size_t>(upload.byte_size));
    return true;
  }

  [[nodiscard]] bool dispatch(const DispatchDesc& dispatch) {
    const auto* product = find_product(dispatch.handle.key);
    if (!available() || product == nullptr ||
        find_field(dispatch.input_field_index) == nullptr ||
        product->generation != dispatch.handle.generation ||
        product->desc.input_field_index != dispatch.input_field_index ||
        dispatch.chunk_count == 0 || dispatch.workgroups_per_chunk == 0 ||
        dispatch.chunk_count > config_.max_dispatch_chunks ||
        dispatch.chunk_count >
            config_.max_dispatch_workgroups_x / dispatch.workgroups_per_chunk) {
      return false;
    }
    auto encoder_desc = WGPU_COMMAND_ENCODER_DESCRIPTOR_INIT;
    auto encoder = wgpuDeviceCreateCommandEncoder(device_, &encoder_desc);
    if (encoder == nullptr) {
      return false;
    }
    auto pass_desc = WGPU_COMPUTE_PASS_DESCRIPTOR_INIT;
    auto pass = wgpuCommandEncoderBeginComputePass(encoder, &pass_desc);
    if (pass == nullptr) {
      wgpuCommandEncoderRelease(encoder);
      return false;
    }
    wgpuComputePassEncoderSetPipeline(pass, product->desc.pipeline);
    wgpuComputePassEncoderSetBindGroup(pass, 0, product->desc.bind_group, 0,
                                       nullptr);
    wgpuComputePassEncoderDispatchWorkgroups(
        pass,
        static_cast<std::uint32_t>(dispatch.chunk_count) *
            dispatch.workgroups_per_chunk,
        1, 1);
    wgpuComputePassEncoderEnd(pass);
    wgpuComputePassEncoderRelease(pass);
    const auto submitted = submit(encoder);
    wgpuCommandEncoderRelease(encoder);
    return submitted;
  }

  [[nodiscard]] bool readback(const ReadbackDesc& readback) {
    const auto* product = find_product(readback.handle.key);
    if (!available() || product == nullptr ||
        product->generation != readback.handle.generation ||
        readback.policy == ReadbackPolicy::None ||
        (readback.policy == ReadbackPolicy::FullField &&
         !config_.allow_full_field_readback) ||
        readback.byte_size == 0 || (readback.byte_size & 3u) != 0 ||
        product->desc.readback_source == nullptr ||
        product->desc.readback_callback == nullptr ||
        readback.byte_size > product->desc.readback_byte_size ||
        !detail::fits_size<std::size_t>(readback.byte_size) ||
        !detail::reserve_readback_bytes(shared_, readback.byte_size)) {
      return false;
    }
    const auto size = static_cast<std::size_t>(readback.byte_size);
    auto buffer_desc = WGPU_BUFFER_DESCRIPTOR_INIT;
    buffer_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst;
    buffer_desc.size = readback.byte_size;
    auto staging = wgpuDeviceCreateBuffer(device_, &buffer_desc);
    if (staging == nullptr) {
      release_readback_bytes(size);
      return false;
    }
    auto* operation = new (std::nothrow) detail::WebGpuReadbackOperation{
        staging,
        readback.handle,
        size,
        product->desc.readback_callback,
        product->desc.readback_userdata,
        shared_,
    };
    if (operation == nullptr) {
      wgpuBufferRelease(staging);
      release_readback_bytes(size);
      return false;
    }
 
    auto encoder_desc = WGPU_COMMAND_ENCODER_DESCRIPTOR_INIT;
    auto encoder = wgpuDeviceCreateCommandEncoder(device_, &encoder_desc);
    if (encoder == nullptr) {
      abandon_readback(operation);
      return false;
    }
    wgpuCommandEncoderCopyBufferToBuffer(encoder, product->desc.readback_source,
                                         product->desc.readback_source_offset,
                                         staging, 0, readback.byte_size);
    const auto submitted = submit(encoder);
    wgpuCommandEncoderRelease(encoder);
    if (!submitted) {
      abandon_readback(operation);
      return false;
    }
    auto callback = WGPU_BUFFER_MAP_CALLBACK_INFO_INIT;
    // Spontaneous delivery avoids making this transport own or pump the
    // application's WGPUInstance. It can run inline or on another thread.
    // Backend cleanup below uses only buffer operations that the stable
    // WebGPU API explicitly permits in spontaneous callbacks; application
    // callbacks inherit the stricter public contract documented above.
    callback.mode = WGPUCallbackMode_AllowSpontaneous;
    callback.callback = detail::webgpu_readback_complete;
    callback.userdata1 = operation;
    const auto future =
        wgpuBufferMapAsync(staging, WGPUMapMode_Read, 0, size, callback);
    // A null future never delivers its callback, so this call still owns the
    // operation. A non-null future may already have completed inline and
    // deleted it; do not inspect operation after this branch.
    if (future.id == 0) {
      abandon_readback(operation);
      return false;
    }
    return true;
  }

  void notify_device_lost() noexcept {
    shared_->available.store(false, std::memory_order_release);
  }

  void notify_device_error() noexcept {
    shared_->available.store(false, std::memory_order_release);
  }
 
 private:
  struct FieldSlot {
    FieldMirrorDesc desc{};
    WGPUBuffer buffer = nullptr;
    std::uint64_t allocated_bytes = 0;
  };
 
  struct ProductSlot {
    bool active = false;
    std::uint64_t generation = 0;
    WebGpuProductDesc desc{};
  };
 
  [[nodiscard]] bool available() const noexcept {
    return device_ != nullptr && queue_ != nullptr &&
           shared_->available.load(std::memory_order_acquire);
  }
 
  [[nodiscard]] static auto aligned_size(std::uint64_t size) noexcept
      -> std::optional<std::uint64_t> {
    if (size > std::numeric_limits<std::uint64_t>::max() - 3) {
      return std::nullopt;
    }
    return (size + 3) & ~std::uint64_t{3};
  }
 
  [[nodiscard]] auto find_field(std::uint32_t field_index) noexcept
      -> FieldSlot* {
    for (auto& field : fields_) {
      if (field.desc.field_index == field_index) {
        return &field;
      }
    }
    return nullptr;
  }
 
  [[nodiscard]] auto find_field(std::uint32_t field_index) const noexcept
      -> const FieldSlot* {
    for (const auto& field : fields_) {
      if (field.desc.field_index == field_index) {
        return &field;
      }
    }
    return nullptr;
  }
 
  [[nodiscard]] auto find_product(std::uint64_t key) noexcept -> ProductSlot* {
    for (auto& product : products_) {
      if (product.active && product.desc.product_key == key) {
        return &product;
      }
    }
    return nullptr;
  }
 
  [[nodiscard]] auto find_product(std::uint64_t key) const noexcept
      -> const ProductSlot* {
    for (const auto& product : products_) {
      if (product.active && product.desc.product_key == key) {
        return &product;
      }
    }
    return nullptr;
  }
 
  static void release_product(ProductSlot& product) noexcept {
    if (!product.active) {
      return;
    }
    wgpuComputePipelineRelease(product.desc.pipeline);
    wgpuBindGroupRelease(product.desc.bind_group);
    if (product.desc.readback_source != nullptr) {
      wgpuBufferRelease(product.desc.readback_source);
    }
    product.active = false;
    product.desc = {};
  }
 
  [[nodiscard]] bool submit(WGPUCommandEncoder encoder) {
    auto command_desc = WGPU_COMMAND_BUFFER_DESCRIPTOR_INIT;
    auto command = wgpuCommandEncoderFinish(encoder, &command_desc);
    if (command == nullptr) {
      return false;
    }
    wgpuQueueSubmit(queue_, 1, &command);
    wgpuCommandBufferRelease(command);
    return true;
  }
 
  void release_readback_bytes(std::size_t size) noexcept {
    shared_->inflight_readback_bytes.fetch_sub(size, std::memory_order_relaxed);
  }
 
  void abandon_readback(detail::WebGpuReadbackOperation* operation) noexcept {
    wgpuBufferRelease(operation->staging);
    release_readback_bytes(operation->byte_size);
    delete operation;
  }
 
  WGPUDevice device_ = nullptr;
  WGPUQueue queue_ = nullptr;
  WebGpuBackendConfig config_{};
  std::shared_ptr<detail::WebGpuSharedState> shared_;
  std::vector<FieldSlot> fields_;
  std::vector<ProductSlot> products_;
  std::uint64_t generation_clock_ = 0;
};
 
static_assert(GpuBackend<WebGpuBackend>);
 
}  // namespace tess::gpu
 
#endif  // defined(TESS_ENABLE_WEBGPU)