react-native-webgpu/packages/webgpu/cpp/rnwgpu/api/GPUComputePassEncoder.h at e0eed38757770ca11c5bdfd993d797c30967aa20 · wcandillon/react-native-webgpu · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
#pragma once

#include <memory>
#include <optional>
#include <string>
#include <variant>
#include <vector>

#include "Unions.h"

#include "RNFHybridObject.h"

#include "webgpu/webgpu_cpp.h"

#include "GPUBindGroup.h"
#include "GPUBuffer.h"
#include "GPUComputePipeline.h"

namespace rnwgpu {

namespace m = margelo;

class GPUComputePassEncoder : public m::HybridObject {
public:
  explicit GPUComputePassEncoder(wgpu::ComputePassEncoder instance,
                                 std::string label)
      : HybridObject("GPUComputePassEncoder"), _instance(instance),
        _label(label) {}

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

  void setPipeline(std::shared_ptr<GPUComputePipeline> pipeline);
  void dispatchWorkgroups(uint32_t workgroupCountX,
                          std::optional<uint32_t> workgroupCountY,
                          std::optional<uint32_t> workgroupCountZ);
  void dispatchWorkgroupsIndirect(std::shared_ptr<GPUBuffer> indirectBuffer,
                                  uint64_t indirectOffset);
  void end();
  void pushDebugGroup(std::string groupLabel);
  void popDebugGroup();
  void insertDebugMarker(std::string markerLabel);
  void setBindGroup(
      uint32_t index,
      std::variant<std::nullptr_t, std::shared_ptr<GPUBindGroup>> bindGroup,
      std::optional<std::vector<uint32_t>> dynamicOffsets);

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

  void loadHybridMethods() override {
    registerHybridGetter("__brand", &GPUComputePassEncoder::getBrand, this);
    registerHybridMethod("setPipeline", &GPUComputePassEncoder::setPipeline,
                         this);
    registerHybridMethod("dispatchWorkgroups",
                         &GPUComputePassEncoder::dispatchWorkgroups, this);
    registerHybridMethod("dispatchWorkgroupsIndirect",
                         &GPUComputePassEncoder::dispatchWorkgroupsIndirect,
                         this);
    registerHybridMethod("end", &GPUComputePassEncoder::end, this);
    registerHybridMethod("pushDebugGroup",
                         &GPUComputePassEncoder::pushDebugGroup, this);
    registerHybridMethod("popDebugGroup", &GPUComputePassEncoder::popDebugGroup,
                         this);
    registerHybridMethod("insertDebugMarker",
                         &GPUComputePassEncoder::insertDebugMarker, this);
    registerHybridMethod("setBindGroup", &GPUComputePassEncoder::setBindGroup,
                         this);

    registerHybridGetter("label", &GPUComputePassEncoder::getLabel, this);
    registerHybridSetter("label", &GPUComputePassEncoder::setLabel, this);
  }

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

  size_t getMemoryPressure() override { return 1024; }

private:
  wgpu::ComputePassEncoder _instance;
  std::string _label;
};

} // namespace rnwgpu