GPUDevice.h

#pragma once

#include <memory>
#include <optional>
#include <string>
#include <unordered_set>
#include <utility>
#include <variant>

#include "Unions.h"

#include "RNFHybridObject.h"

#include "rnwgpu/async/AsyncRunner.h"
#include "rnwgpu/async/AsyncTaskHandle.h"

#include "webgpu/webgpu_cpp.h"

#include "GPUBindGroup.h"
#include "GPUBindGroupDescriptor.h"
#include "GPUBindGroupLayout.h"
#include "GPUBindGroupLayoutDescriptor.h"
#include "GPUBuffer.h"
#include "GPUBufferDescriptor.h"
#include "GPUCommandEncoder.h"
#include "GPUCommandEncoderDescriptor.h"
#include "GPUComputePipeline.h"
#include "GPUComputePipelineDescriptor.h"
#include "GPUDeviceLostInfo.h"
#include "GPUError.h"
#include "GPUExternalTexture.h"
#include "GPUExternalTextureDescriptor.h"
#include "GPUPipelineLayout.h"
#include "GPUPipelineLayoutDescriptor.h"
#include "GPUQuerySet.h"
#include "GPUQuerySetDescriptor.h"
#include "GPUQueue.h"
#include "GPURenderBundleEncoder.h"
#include "GPURenderBundleEncoderDescriptor.h"
#include "GPURenderPipeline.h"
#include "GPURenderPipelineDescriptor.h"
#include "GPUSampler.h"
#include "GPUSamplerDescriptor.h"
#include "GPUShaderModule.h"
#include "GPUShaderModuleDescriptor.h"
#include "GPUSupportedLimits.h"
#include "GPUTexture.h"
#include "GPUTextureDescriptor.h"

namespace rnwgpu {

namespace m = margelo;

class GPUDevice : public m::HybridObject {
public:
  explicit GPUDevice(wgpu::Device instance,
                     std::shared_ptr<async::AsyncRunner> async,
                     std::string label)
      : HybridObject("GPUDevice"), _instance(instance), _async(async),
        _label(label) {}

public:
  std::string getBrand() { return _name; }

  void destroy();
  std::shared_ptr<GPUBuffer>
  createBuffer(std::shared_ptr<GPUBufferDescriptor> descriptor);
  std::shared_ptr<GPUTexture>
  createTexture(std::shared_ptr<GPUTextureDescriptor> descriptor);
  std::shared_ptr<GPUSampler> createSampler(
      std::optional<std::shared_ptr<GPUSamplerDescriptor>> descriptor);
  std::shared_ptr<GPUExternalTexture> importExternalTexture(
      std::shared_ptr<GPUExternalTextureDescriptor> descriptor);
  std::shared_ptr<GPUBindGroupLayout> createBindGroupLayout(
      std::shared_ptr<GPUBindGroupLayoutDescriptor> descriptor);
  std::shared_ptr<GPUPipelineLayout>
  createPipelineLayout(std::shared_ptr<GPUPipelineLayoutDescriptor> descriptor);
  std::shared_ptr<GPUBindGroup>
  createBindGroup(std::shared_ptr<GPUBindGroupDescriptor> descriptor);
  std::shared_ptr<GPUShaderModule>
  createShaderModule(std::shared_ptr<GPUShaderModuleDescriptor> descriptor);
  std::shared_ptr<GPUComputePipeline> createComputePipeline(
      std::shared_ptr<GPUComputePipelineDescriptor> descriptor);
  std::shared_ptr<GPURenderPipeline>
  createRenderPipeline(std::shared_ptr<GPURenderPipelineDescriptor> descriptor);
  async::AsyncTaskHandle createComputePipelineAsync(
      std::shared_ptr<GPUComputePipelineDescriptor> descriptor);
  async::AsyncTaskHandle createRenderPipelineAsync(
      std::shared_ptr<GPURenderPipelineDescriptor> descriptor);
  std::shared_ptr<GPUCommandEncoder> createCommandEncoder(
      std::optional<std::shared_ptr<GPUCommandEncoderDescriptor>> descriptor);
  std::shared_ptr<GPURenderBundleEncoder> createRenderBundleEncoder(
      std::shared_ptr<GPURenderBundleEncoderDescriptor> descriptor);
  std::shared_ptr<GPUQuerySet>
  createQuerySet(std::shared_ptr<GPUQuerySetDescriptor> descriptor);
  void pushErrorScope(wgpu::ErrorFilter filter);
  async::AsyncTaskHandle popErrorScope();

  std::unordered_set<std::string> getFeatures();
  std::shared_ptr<GPUSupportedLimits> getLimits();
  std::shared_ptr<GPUQueue> getQueue();
  async::AsyncTaskHandle getLost();
  void notifyDeviceLost(wgpu::DeviceLostReason reason, std::string message);
  void forceLossForTesting();

  std::string getLabel() { return _label; }
  void setLabel(const std::string &label) {
    _label = label;
    _instance.SetLabel(_label.c_str());
  }

  void loadHybridMethods() override {
    registerHybridGetter("__brand", &GPUDevice::getBrand, this);
    registerHybridMethod("destroy", &GPUDevice::destroy, this);
    registerHybridMethod("createBuffer", &GPUDevice::createBuffer, this);
    registerHybridMethod("createTexture", &GPUDevice::createTexture, this);
    registerHybridMethod("createSampler", &GPUDevice::createSampler, this);
    registerHybridMethod("importExternalTexture",
                         &GPUDevice::importExternalTexture, this);
    registerHybridMethod("createBindGroupLayout",
                         &GPUDevice::createBindGroupLayout, this);
    registerHybridMethod("createPipelineLayout",
                         &GPUDevice::createPipelineLayout, this);
    registerHybridMethod("createBindGroup", &GPUDevice::createBindGroup, this);
    registerHybridMethod("createShaderModule", &GPUDevice::createShaderModule,
                         this);
    registerHybridMethod("createComputePipeline",
                         &GPUDevice::createComputePipeline, this);
    registerHybridMethod("createRenderPipeline",
                         &GPUDevice::createRenderPipeline, this);
    registerHybridMethod("createComputePipelineAsync",
                         &GPUDevice::createComputePipelineAsync, this);
    registerHybridMethod("createRenderPipelineAsync",
                         &GPUDevice::createRenderPipelineAsync, this);
    registerHybridMethod("createCommandEncoder",
                         &GPUDevice::createCommandEncoder, this);
    registerHybridMethod("createRenderBundleEncoder",
                         &GPUDevice::createRenderBundleEncoder, this);
    registerHybridMethod("createQuerySet", &GPUDevice::createQuerySet, this);
    registerHybridMethod("pushErrorScope", &GPUDevice::pushErrorScope, this);
    registerHybridMethod("popErrorScope", &GPUDevice::popErrorScope, this);
    registerHybridGetter("features", &GPUDevice::getFeatures, this);
    registerHybridGetter("limits", &GPUDevice::getLimits, this);
    registerHybridGetter("queue", &GPUDevice::getQueue, this);
    registerHybridGetter("lost", &GPUDevice::getLost, this);
    registerHybridGetter("label", &GPUDevice::getLabel, this);
    registerHybridSetter("label", &GPUDevice::setLabel, this);
    registerHybridMethod("forceLossForTesting", &GPUDevice::forceLossForTesting,
                         this);
  }

  inline const wgpu::Device get() { return _instance; }

  size_t getMemoryPressure() override {
    // Devices keep large driver heaps, pipeline caches and residency tracking.
    constexpr size_t kDeviceFloor = 8 * 1024 * 1024; // 8MB baseline
    return kDeviceFloor + _label.capacity();
  }

private:
  friend class GPUAdapter;

  wgpu::Device _instance;
  std::shared_ptr<async::AsyncRunner> _async;
  std::string _label;
  std::optional<async::AsyncTaskHandle> _lostHandle;
  std::shared_ptr<GPUDeviceLostInfo> _lostInfo;
  bool _lostSettled = false;
  std::optional<async::AsyncTaskHandle::ResolveFunction> _lostResolve;
};

} // namespace rnwgpu