summaryrefslogtreecommitdiff
path: root/src/gpu/effects/cnn_v2_effect.cc
blob: 04fa74eef909e53019d45374b7dc13fe4dbfcb91 (plain)
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
// CNN v2 Effect Implementation

#include "gpu/effects/cnn_v2_effect.h"

#if defined(USE_TEST_ASSETS)
#include "test_assets.h"
#else
#include "generated/assets.h"
#endif

#include "util/asset_manager.h"
#include "util/fatal_error.h"
#include <cstring>

CNNv2Effect::CNNv2Effect(const GpuContext& ctx)
    : PostProcessEffect(ctx),
      static_pipeline_(nullptr),
      static_bind_group_(nullptr),
      static_features_tex_(nullptr),
      static_features_view_(nullptr),
      input_mip_tex_(nullptr),
      initialized_(false) {
  std::memset(input_mip_view_, 0, sizeof(input_mip_view_));
}

CNNv2Effect::~CNNv2Effect() {
  cleanup();
}

void CNNv2Effect::init(MainSequence* demo) {
  (void)demo;
  if (initialized_) return;

  create_textures();
  create_pipelines();

  initialized_ = true;
}

void CNNv2Effect::resize(int width, int height) {
  PostProcessEffect::resize(width, height);
  cleanup();
  create_textures();
  create_pipelines();
}

void CNNv2Effect::create_textures() {
  const WGPUExtent3D size = {
    static_cast<uint32_t>(width_),
    static_cast<uint32_t>(height_),
    1
  };

  // Static features texture (8×f16 packed as 4×u32)
  WGPUTextureDescriptor static_desc = {};
  static_desc.usage = WGPUTextureUsage_StorageBinding | WGPUTextureUsage_TextureBinding;
  static_desc.dimension = WGPUTextureDimension_2D;
  static_desc.size = size;
  static_desc.format = WGPUTextureFormat_RGBA32Uint;
  static_desc.mipLevelCount = 1;
  static_desc.sampleCount = 1;
  static_features_tex_ = wgpuDeviceCreateTexture(ctx_.device, &static_desc);

  WGPUTextureViewDescriptor view_desc = {};
  view_desc.format = WGPUTextureFormat_RGBA32Uint;
  view_desc.dimension = WGPUTextureViewDimension_2D;
  view_desc.baseMipLevel = 0;
  view_desc.mipLevelCount = 1;
  view_desc.baseArrayLayer = 0;
  view_desc.arrayLayerCount = 1;
  static_features_view_ = wgpuTextureCreateView(static_features_tex_, &view_desc);

  // Input texture with mips (for multi-scale features)
  WGPUTextureDescriptor input_mip_desc = {};
  input_mip_desc.usage = WGPUTextureUsage_TextureBinding | WGPUTextureUsage_CopyDst;
  input_mip_desc.dimension = WGPUTextureDimension_2D;
  input_mip_desc.size = size;
  input_mip_desc.format = WGPUTextureFormat_RGBA8Unorm;
  input_mip_desc.mipLevelCount = 3;  // Levels 0, 1, 2
  input_mip_desc.sampleCount = 1;
  input_mip_tex_ = wgpuDeviceCreateTexture(ctx_.device, &input_mip_desc);

  for (int i = 0; i < 3; ++i) {
    WGPUTextureViewDescriptor mip_view_desc = {};
    mip_view_desc.format = WGPUTextureFormat_RGBA8Unorm;
    mip_view_desc.dimension = WGPUTextureViewDimension_2D;
    mip_view_desc.baseMipLevel = i;
    mip_view_desc.mipLevelCount = 1;
    mip_view_desc.baseArrayLayer = 0;
    mip_view_desc.arrayLayerCount = 1;
    input_mip_view_[i] = wgpuTextureCreateView(input_mip_tex_, &mip_view_desc);
  }

  // Layer textures (placeholder - will be created based on config)
  // TODO: Create layer textures based on layer_configs_
}

void CNNv2Effect::create_pipelines() {
  // Static features compute pipeline
  size_t shader_size = 0;
  const char* static_code = (const char*)GetAsset(AssetId::ASSET_SHADER_CNN_V2_STATIC, &shader_size);

  if (!static_code || shader_size == 0) {
    // Shader not available (e.g., in test mode) - skip pipeline creation
    return;
  }

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(static_code);

  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;

  WGPUShaderModule static_module = wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc);
  if (!static_module) {
    return;
  }

  WGPUComputePipelineDescriptor pipeline_desc = {};
  pipeline_desc.compute.module = static_module;
  pipeline_desc.compute.entryPoint = str_view("main");

  static_pipeline_ = wgpuDeviceCreateComputePipeline(ctx_.device, &pipeline_desc);
  wgpuShaderModuleRelease(static_module);

  // TODO: Create layer pipelines
  // TODO: Create bind groups
}

void CNNv2Effect::update_bind_group(WGPUTextureView input_view) {
  (void)input_view;
  // TODO: Create bind groups for static features and layers
}

void CNNv2Effect::render(WGPURenderPassEncoder pass,
                         const CommonPostProcessUniforms& uniforms) {
  (void)pass;
  (void)uniforms;
  if (!initialized_) return;

  // TODO: Multi-pass execution
  // 1. Compute static features
  // 2. Execute CNN layers
  // 3. Composite to output
}

void CNNv2Effect::cleanup() {
  if (static_features_view_) wgpuTextureViewRelease(static_features_view_);
  if (static_features_tex_) wgpuTextureRelease(static_features_tex_);
  if (static_bind_group_) wgpuBindGroupRelease(static_bind_group_);
  if (static_pipeline_) wgpuComputePipelineRelease(static_pipeline_);

  for (int i = 0; i < 3; ++i) {
    if (input_mip_view_[i]) wgpuTextureViewRelease(input_mip_view_[i]);
  }
  if (input_mip_tex_) wgpuTextureRelease(input_mip_tex_);

  for (auto view : layer_views_) wgpuTextureViewRelease(view);
  for (auto tex : layer_textures_) wgpuTextureRelease(tex);
  for (auto bg : layer_bind_groups_) wgpuBindGroupRelease(bg);
  for (auto pipeline : layer_pipelines_) wgpuComputePipelineRelease(pipeline);

  layer_views_.clear();
  layer_textures_.clear();
  layer_bind_groups_.clear();
  layer_pipelines_.clear();

  initialized_ = false;
}