diff options
Diffstat (limited to 'tools')
| -rw-r--r-- | tools/cnn_test.cc | 666 |
1 files changed, 666 insertions, 0 deletions
diff --git a/tools/cnn_test.cc b/tools/cnn_test.cc index c2983a9..5823110 100644 --- a/tools/cnn_test.cc +++ b/tools/cnn_test.cc @@ -28,6 +28,7 @@ #include <cstdlib> #include <cstring> #include <vector> +#include <cmath> // Helper to get asset string or empty string static const char* SafeGetAsset(AssetId id) { @@ -44,6 +45,7 @@ struct Args { const char* save_intermediates = nullptr; int num_layers = 3; // Default to 3 layers bool debug_hex = false; // Print first 8 pixels as hex + int cnn_version = 1; // 1=CNNEffect, 2=CNNv2Effect }; // Parse command-line arguments @@ -83,6 +85,12 @@ static bool parse_args(int argc, char** argv, Args* args) { } } else if (strcmp(argv[i], "--debug-hex") == 0) { args->debug_hex = true; + } else if (strcmp(argv[i], "--cnn-version") == 0 && i + 1 < argc) { + args->cnn_version = atoi(argv[++i]); + if (args->cnn_version < 1 || args->cnn_version > 2) { + fprintf(stderr, "Error: cnn-version must be 1 or 2\n"); + return false; + } } else if (strcmp(argv[i], "--help") == 0) { return false; } else { @@ -103,6 +111,7 @@ static void print_usage(const char* prog) { fprintf(stderr, " --layers N Number of CNN layers (1-10, default: 3)\n"); fprintf(stderr, " --save-intermediates DIR Save intermediate layers to directory\n"); fprintf(stderr, " --debug-hex Print first 8 pixels as hex (debug)\n"); + fprintf(stderr, " --cnn-version N CNN version: 1 (default) or 2\n"); fprintf(stderr, " --help Show this help\n"); } @@ -257,6 +266,650 @@ static bool save_ppm(const char* path, const std::vector<uint8_t>& pixels, return true; } +// CNN v2 structures (matching CNNv2Effect) +struct CNNv2LayerInfo { + uint32_t kernel_size; + uint32_t in_channels; + uint32_t out_channels; + uint32_t weight_offset; + uint32_t weight_count; +}; + +struct CNNv2LayerParams { + uint32_t kernel_size; + uint32_t in_channels; + uint32_t out_channels; + uint32_t weight_offset; + uint32_t is_output_layer; + float blend_amount; +}; + +struct CNNv2StaticFeatureParams { + uint32_t mip_level; + uint32_t padding[3]; +}; + +// Convert RGBA32Uint (packed f16) texture to BGRA8Unorm +static std::vector<uint8_t> readback_rgba32uint_to_bgra8( + WGPUDevice device, WGPUQueue queue, WGPUTexture texture, + int width, int height) { + // Create staging buffer + const uint32_t bytes_per_row = width * 16; // 4×u32 per pixel + const uint32_t padded_bytes_per_row = (bytes_per_row + 255) & ~255; + const size_t buffer_size = padded_bytes_per_row * height; + + WGPUBufferDescriptor buffer_desc = {}; + buffer_desc.size = buffer_size; + buffer_desc.usage = WGPUBufferUsage_CopyDst | WGPUBufferUsage_MapRead; + buffer_desc.mappedAtCreation = false; + + WGPUBuffer staging = wgpuDeviceCreateBuffer(device, &buffer_desc); + + // Copy texture to buffer + WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device, nullptr); + + WGPUTexelCopyTextureInfo src = {}; + src.texture = texture; + src.mipLevel = 0; + + WGPUTexelCopyBufferInfo dst = {}; + dst.buffer = staging; + dst.layout.bytesPerRow = padded_bytes_per_row; + dst.layout.rowsPerImage = height; + + WGPUExtent3D copy_size = { + static_cast<uint32_t>(width), + static_cast<uint32_t>(height), + 1}; + + wgpuCommandEncoderCopyTextureToBuffer(encoder, &src, &dst, ©_size); + + WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, nullptr); + wgpuQueueSubmit(queue, 1, &commands); + wgpuCommandBufferRelease(commands); + wgpuCommandEncoderRelease(encoder); + + // Wait for copy to complete + wgpuDevicePoll(device, true, nullptr); + + // Map and read buffer + struct MapState { + bool done = false; + }; + MapState map_state; + + auto map_cb = [](WGPUMapAsyncStatus status, WGPUStringView message, + void* userdata1, void* userdata2) { + (void)message; + (void)userdata2; + MapState* state = (MapState*)userdata1; + state->done = (status == WGPUMapAsyncStatus_Success); + }; + + WGPUBufferMapCallbackInfo map_info = {}; + map_info.mode = WGPUCallbackMode_AllowProcessEvents; + map_info.callback = map_cb; + map_info.userdata1 = &map_state; + + wgpuBufferMapAsync(staging, WGPUMapMode_Read, 0, buffer_size, map_info); + + // Wait for mapping to complete + for (int i = 0; i < 100 && !map_state.done; ++i) { + wgpuDevicePoll(device, true, nullptr); + } + + if (!map_state.done) { + fprintf(stderr, "Error: Buffer mapping timed out\n"); + wgpuBufferRelease(staging); + return std::vector<uint8_t>(); + } + + const uint32_t* mapped = + (const uint32_t*)wgpuBufferGetConstMappedRange(staging, 0, buffer_size); + + std::vector<uint8_t> result(width * height * 4); + + // Unpack f16 to u8 (BGRA) + for (int y = 0; y < height; ++y) { + const uint32_t* row = + (const uint32_t*)((const uint8_t*)mapped + y * padded_bytes_per_row); + for (int x = 0; x < width; ++x) { + // Read 4×u32 (8×f16) + uint32_t data[4]; + data[0] = row[x * 4 + 0]; + data[1] = row[x * 4 + 1]; + data[2] = row[x * 4 + 2]; + data[3] = row[x * 4 + 3]; + + // Extract RGBA channels (first 4 f16 values) + uint16_t r16 = data[0] & 0xFFFF; + uint16_t g16 = (data[0] >> 16) & 0xFFFF; + uint16_t b16 = data[1] & 0xFFFF; + uint16_t a16 = (data[1] >> 16) & 0xFFFF; + + // Convert f16 to f32 (simple decode) + auto f16_to_f32 = [](uint16_t h) -> float { + uint32_t sign = (h >> 15) & 1; + uint32_t exp = (h >> 10) & 0x1F; + uint32_t frac = h & 0x3FF; + + if (exp == 0) { + if (frac == 0) return sign ? -0.0f : 0.0f; + // Denormal + float val = frac / 1024.0f / 16384.0f; + return sign ? -val : val; + } + if (exp == 31) { + return frac ? NAN : (sign ? -INFINITY : INFINITY); + } + + int32_t e = exp - 15; + float val = (1.0f + frac / 1024.0f) * powf(2.0f, e); + return sign ? -val : val; + }; + + float r = f16_to_f32(r16); + float g = f16_to_f32(g16); + float b = f16_to_f32(b16); + float a = f16_to_f32(a16); + + // Clamp to [0,1] and convert to u8 + auto clamp_u8 = [](float v) -> uint8_t { + if (v <= 0.0f) return 0; + if (v >= 1.0f) return 255; + return static_cast<uint8_t>(v * 255.0f + 0.5f); + }; + + result[(y * width + x) * 4 + 0] = clamp_u8(b); + result[(y * width + x) * 4 + 1] = clamp_u8(g); + result[(y * width + x) * 4 + 2] = clamp_u8(r); + result[(y * width + x) * 4 + 3] = clamp_u8(a); + } + } + + wgpuBufferUnmap(staging); + wgpuBufferRelease(staging); + + return result; +} + +// Process image with CNN v2 +static bool process_cnn_v2(WGPUDevice device, WGPUQueue queue, + WGPUInstance instance, WGPUTexture input_texture, + int width, int height, const Args& args) { + printf("Using CNN v2 (storage buffer architecture)\n"); + + // Load weights + size_t weights_size = 0; + const uint8_t* weights_data = + (const uint8_t*)GetAsset(AssetId::ASSET_WEIGHTS_CNN_V2, &weights_size); + + if (!weights_data || weights_size < 20) { + fprintf(stderr, "Error: CNN v2 weights not available\n"); + return false; + } + + // Parse header + const uint32_t* header = (const uint32_t*)weights_data; + uint32_t magic = header[0]; + uint32_t version = header[1]; + uint32_t num_layers = header[2]; + uint32_t total_weights = header[3]; + + if (magic != 0x324e4e43) { // 'CNN2' + fprintf(stderr, "Error: Invalid CNN v2 weights magic\n"); + return false; + } + + uint32_t mip_level = 0; + if (version == 2) { + mip_level = header[4]; + } + + printf("Loaded CNN v2 weights: %u layers, %u weights, version %u\n", + num_layers, total_weights, version); + + // Parse layer info + const uint32_t header_u32_count = (version == 1) ? 4 : 5; + const uint32_t* layer_data = header + header_u32_count; + std::vector<CNNv2LayerInfo> layer_info; + + for (uint32_t i = 0; i < num_layers; ++i) { + CNNv2LayerInfo info; + info.kernel_size = layer_data[i * 5 + 0]; + info.in_channels = layer_data[i * 5 + 1]; + info.out_channels = layer_data[i * 5 + 2]; + info.weight_offset = layer_data[i * 5 + 3]; + info.weight_count = layer_data[i * 5 + 4]; + layer_info.push_back(info); + + printf(" Layer %u: %ux%u conv, %u→%u channels, %u weights\n", i, + info.kernel_size, info.kernel_size, info.in_channels, + info.out_channels, info.weight_count); + } + + // Create weights storage buffer + WGPUBufferDescriptor weights_buffer_desc = {}; + weights_buffer_desc.size = weights_size; + weights_buffer_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + weights_buffer_desc.mappedAtCreation = false; + + WGPUBuffer weights_buffer = + wgpuDeviceCreateBuffer(device, &weights_buffer_desc); + wgpuQueueWriteBuffer(queue, weights_buffer, 0, weights_data, weights_size); + + // Create input view + const WGPUTextureViewDescriptor view_desc = { + .format = WGPUTextureFormat_BGRA8Unorm, + .dimension = WGPUTextureViewDimension_2D, + .baseMipLevel = 0, + .mipLevelCount = 1, + .baseArrayLayer = 0, + .arrayLayerCount = 1, + }; + WGPUTextureView input_view = wgpuTextureCreateView(input_texture, &view_desc); + + // Create static features texture (RGBA32Uint) + const WGPUTextureDescriptor static_desc = { + .usage = WGPUTextureUsage_StorageBinding | WGPUTextureUsage_TextureBinding | WGPUTextureUsage_CopySrc, + .dimension = WGPUTextureDimension_2D, + .size = {static_cast<uint32_t>(width), static_cast<uint32_t>(height), 1}, + .format = WGPUTextureFormat_RGBA32Uint, + .mipLevelCount = 1, + .sampleCount = 1, + }; + WGPUTexture static_features_tex = + wgpuDeviceCreateTexture(device, &static_desc); + WGPUTextureView static_features_view = + wgpuTextureCreateView(static_features_tex, nullptr); + + // Create layer textures (ping-pong) + WGPUTexture layer_textures[2] = { + wgpuDeviceCreateTexture(device, &static_desc), + wgpuDeviceCreateTexture(device, &static_desc), + }; + WGPUTextureView layer_views[2] = { + wgpuTextureCreateView(layer_textures[0], nullptr), + wgpuTextureCreateView(layer_textures[1], nullptr), + }; + + // Load shaders + const char* static_shader = + SafeGetAsset(AssetId::ASSET_SHADER_CNN_V2_STATIC); + const char* layer_shader = + SafeGetAsset(AssetId::ASSET_SHADER_CNN_V2_COMPUTE); + + if (!static_shader[0] || !layer_shader[0]) { + fprintf(stderr, "Error: CNN v2 shaders not available\n"); + wgpuTextureViewRelease(static_features_view); + wgpuTextureRelease(static_features_tex); + wgpuTextureViewRelease(layer_views[0]); + wgpuTextureViewRelease(layer_views[1]); + wgpuTextureRelease(layer_textures[0]); + wgpuTextureRelease(layer_textures[1]); + wgpuBufferRelease(weights_buffer); + wgpuTextureViewRelease(input_view); + return false; + } + + // Create static feature params buffer + WGPUBufferDescriptor static_params_desc = {}; + static_params_desc.size = sizeof(CNNv2StaticFeatureParams); + static_params_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + static_params_desc.mappedAtCreation = false; + + WGPUBuffer static_params_buffer = + wgpuDeviceCreateBuffer(device, &static_params_desc); + + CNNv2StaticFeatureParams static_params; + static_params.mip_level = mip_level; + static_params.padding[0] = 0; + static_params.padding[1] = 0; + static_params.padding[2] = 0; + wgpuQueueWriteBuffer(queue, static_params_buffer, 0, &static_params, + sizeof(static_params)); + + // Create static features compute pipeline + WGPUShaderSourceWGSL static_wgsl = {}; + static_wgsl.chain.sType = WGPUSType_ShaderSourceWGSL; + static_wgsl.code = str_view(static_shader); + + WGPUShaderModuleDescriptor static_module_desc = {}; + static_module_desc.nextInChain = &static_wgsl.chain; + + WGPUShaderModule static_module = + wgpuDeviceCreateShaderModule(device, &static_module_desc); + + // Bind group layout: 0=input, 1=input_mip1, 2=input_mip2, 3=depth, 4=output, + // 5=params + WGPUBindGroupLayoutEntry static_bgl_entries[6] = {}; + static_bgl_entries[0].binding = 0; + static_bgl_entries[0].visibility = WGPUShaderStage_Compute; + static_bgl_entries[0].texture.sampleType = WGPUTextureSampleType_Float; + static_bgl_entries[0].texture.viewDimension = WGPUTextureViewDimension_2D; + + static_bgl_entries[1].binding = 1; + static_bgl_entries[1].visibility = WGPUShaderStage_Compute; + static_bgl_entries[1].texture.sampleType = WGPUTextureSampleType_Float; + static_bgl_entries[1].texture.viewDimension = WGPUTextureViewDimension_2D; + + static_bgl_entries[2].binding = 2; + static_bgl_entries[2].visibility = WGPUShaderStage_Compute; + static_bgl_entries[2].texture.sampleType = WGPUTextureSampleType_Float; + static_bgl_entries[2].texture.viewDimension = WGPUTextureViewDimension_2D; + + static_bgl_entries[3].binding = 3; + static_bgl_entries[3].visibility = WGPUShaderStage_Compute; + static_bgl_entries[3].texture.sampleType = WGPUTextureSampleType_Float; + static_bgl_entries[3].texture.viewDimension = WGPUTextureViewDimension_2D; + + static_bgl_entries[4].binding = 4; + static_bgl_entries[4].visibility = WGPUShaderStage_Compute; + static_bgl_entries[4].storageTexture.access = + WGPUStorageTextureAccess_WriteOnly; + static_bgl_entries[4].storageTexture.format = WGPUTextureFormat_RGBA32Uint; + static_bgl_entries[4].storageTexture.viewDimension = + WGPUTextureViewDimension_2D; + + static_bgl_entries[5].binding = 5; + static_bgl_entries[5].visibility = WGPUShaderStage_Compute; + static_bgl_entries[5].buffer.type = WGPUBufferBindingType_Uniform; + static_bgl_entries[5].buffer.minBindingSize = + sizeof(CNNv2StaticFeatureParams); + + WGPUBindGroupLayoutDescriptor static_bgl_desc = {}; + static_bgl_desc.entryCount = 6; + static_bgl_desc.entries = static_bgl_entries; + + WGPUBindGroupLayout static_bgl = + wgpuDeviceCreateBindGroupLayout(device, &static_bgl_desc); + + WGPUPipelineLayoutDescriptor static_pl_desc = {}; + static_pl_desc.bindGroupLayoutCount = 1; + static_pl_desc.bindGroupLayouts = &static_bgl; + + WGPUPipelineLayout static_pl = + wgpuDeviceCreatePipelineLayout(device, &static_pl_desc); + + WGPUComputePipelineDescriptor static_pipeline_desc = {}; + static_pipeline_desc.compute.module = static_module; + static_pipeline_desc.compute.entryPoint = str_view("main"); + static_pipeline_desc.layout = static_pl; + + WGPUComputePipeline static_pipeline = + wgpuDeviceCreateComputePipeline(device, &static_pipeline_desc); + + wgpuShaderModuleRelease(static_module); + wgpuPipelineLayoutRelease(static_pl); + + // Create static bind group (use input as all mips for simplicity) + WGPUBindGroupEntry static_bg_entries[6] = {}; + static_bg_entries[0].binding = 0; + static_bg_entries[0].textureView = input_view; + static_bg_entries[1].binding = 1; + static_bg_entries[1].textureView = input_view; + static_bg_entries[2].binding = 2; + static_bg_entries[2].textureView = input_view; + static_bg_entries[3].binding = 3; + static_bg_entries[3].textureView = input_view; // Depth (use input) + static_bg_entries[4].binding = 4; + static_bg_entries[4].textureView = static_features_view; + static_bg_entries[5].binding = 5; + static_bg_entries[5].buffer = static_params_buffer; + static_bg_entries[5].size = sizeof(CNNv2StaticFeatureParams); + + WGPUBindGroupDescriptor static_bg_desc = {}; + static_bg_desc.layout = static_bgl; + static_bg_desc.entryCount = 6; + static_bg_desc.entries = static_bg_entries; + + WGPUBindGroup static_bg = wgpuDeviceCreateBindGroup(device, &static_bg_desc); + + wgpuBindGroupLayoutRelease(static_bgl); + + // Create layer compute pipeline + WGPUShaderSourceWGSL layer_wgsl = {}; + layer_wgsl.chain.sType = WGPUSType_ShaderSourceWGSL; + layer_wgsl.code = str_view(layer_shader); + + WGPUShaderModuleDescriptor layer_module_desc = {}; + layer_module_desc.nextInChain = &layer_wgsl.chain; + + WGPUShaderModule layer_module = + wgpuDeviceCreateShaderModule(device, &layer_module_desc); + + // Layer bind group layout: + // 0=static_features, 1=layer_input, 2=output, 3=weights, 4=params, + // 5=original + WGPUBindGroupLayoutEntry layer_bgl_entries[6] = {}; + layer_bgl_entries[0].binding = 0; + layer_bgl_entries[0].visibility = WGPUShaderStage_Compute; + layer_bgl_entries[0].texture.sampleType = WGPUTextureSampleType_Uint; + layer_bgl_entries[0].texture.viewDimension = WGPUTextureViewDimension_2D; + + layer_bgl_entries[1].binding = 1; + layer_bgl_entries[1].visibility = WGPUShaderStage_Compute; + layer_bgl_entries[1].texture.sampleType = WGPUTextureSampleType_Uint; + layer_bgl_entries[1].texture.viewDimension = WGPUTextureViewDimension_2D; + + layer_bgl_entries[2].binding = 2; + layer_bgl_entries[2].visibility = WGPUShaderStage_Compute; + layer_bgl_entries[2].storageTexture.access = + WGPUStorageTextureAccess_WriteOnly; + layer_bgl_entries[2].storageTexture.format = WGPUTextureFormat_RGBA32Uint; + layer_bgl_entries[2].storageTexture.viewDimension = + WGPUTextureViewDimension_2D; + + layer_bgl_entries[3].binding = 3; + layer_bgl_entries[3].visibility = WGPUShaderStage_Compute; + layer_bgl_entries[3].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + + layer_bgl_entries[4].binding = 4; + layer_bgl_entries[4].visibility = WGPUShaderStage_Compute; + layer_bgl_entries[4].buffer.type = WGPUBufferBindingType_Uniform; + layer_bgl_entries[4].buffer.minBindingSize = sizeof(CNNv2LayerParams); + + layer_bgl_entries[5].binding = 5; + layer_bgl_entries[5].visibility = WGPUShaderStage_Compute; + layer_bgl_entries[5].texture.sampleType = WGPUTextureSampleType_Float; + layer_bgl_entries[5].texture.viewDimension = WGPUTextureViewDimension_2D; + + WGPUBindGroupLayoutDescriptor layer_bgl_desc = {}; + layer_bgl_desc.entryCount = 6; + layer_bgl_desc.entries = layer_bgl_entries; + + WGPUBindGroupLayout layer_bgl = + wgpuDeviceCreateBindGroupLayout(device, &layer_bgl_desc); + + WGPUPipelineLayoutDescriptor layer_pl_desc = {}; + layer_pl_desc.bindGroupLayoutCount = 1; + layer_pl_desc.bindGroupLayouts = &layer_bgl; + + WGPUPipelineLayout layer_pl = + wgpuDeviceCreatePipelineLayout(device, &layer_pl_desc); + + WGPUComputePipelineDescriptor layer_pipeline_desc = {}; + layer_pipeline_desc.compute.module = layer_module; + layer_pipeline_desc.compute.entryPoint = str_view("main"); + layer_pipeline_desc.layout = layer_pl; + + WGPUComputePipeline layer_pipeline = + wgpuDeviceCreateComputePipeline(device, &layer_pipeline_desc); + + wgpuShaderModuleRelease(layer_module); + wgpuPipelineLayoutRelease(layer_pl); + + // Create layer params buffers + std::vector<WGPUBuffer> layer_params_buffers; + for (size_t i = 0; i < layer_info.size(); ++i) { + WGPUBufferDescriptor params_desc = {}; + params_desc.size = sizeof(CNNv2LayerParams); + params_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + params_desc.mappedAtCreation = false; + + WGPUBuffer buf = wgpuDeviceCreateBuffer(device, ¶ms_desc); + layer_params_buffers.push_back(buf); + } + + // Execute compute passes + WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device, nullptr); + + // Pass 1: Static features + printf("Computing static features...\n"); + WGPUComputePassEncoder static_pass = + wgpuCommandEncoderBeginComputePass(encoder, nullptr); + wgpuComputePassEncoderSetPipeline(static_pass, static_pipeline); + wgpuComputePassEncoderSetBindGroup(static_pass, 0, static_bg, 0, nullptr); + + uint32_t workgroups_x = (width + 7) / 8; + uint32_t workgroups_y = (height + 7) / 8; + wgpuComputePassEncoderDispatchWorkgroups(static_pass, workgroups_x, + workgroups_y, 1); + + wgpuComputePassEncoderEnd(static_pass); + wgpuComputePassEncoderRelease(static_pass); + + // Pass 2-N: CNN layers + for (size_t i = 0; i < layer_info.size(); ++i) { + const CNNv2LayerInfo& info = layer_info[i]; + + printf("Processing layer %zu/%zu (%ux%u, %u→%u channels)...\n", i + 1, + layer_info.size(), info.kernel_size, info.kernel_size, + info.in_channels, info.out_channels); + + // Update layer params + CNNv2LayerParams params; + params.kernel_size = info.kernel_size; + params.in_channels = info.in_channels; + params.out_channels = info.out_channels; + params.weight_offset = info.weight_offset; + params.is_output_layer = (i == layer_info.size() - 1) ? 1 : 0; + params.blend_amount = args.blend; + + wgpuQueueWriteBuffer(queue, layer_params_buffers[i], 0, ¶ms, + sizeof(params)); + + // Create bind group for this layer + WGPUBindGroupEntry layer_bg_entries[6] = {}; + layer_bg_entries[0].binding = 0; + layer_bg_entries[0].textureView = static_features_view; + + layer_bg_entries[1].binding = 1; + layer_bg_entries[1].textureView = + (i == 0) ? static_features_view : layer_views[i % 2]; + + layer_bg_entries[2].binding = 2; + layer_bg_entries[2].textureView = layer_views[(i + 1) % 2]; + + layer_bg_entries[3].binding = 3; + layer_bg_entries[3].buffer = weights_buffer; + layer_bg_entries[3].size = weights_size; + + layer_bg_entries[4].binding = 4; + layer_bg_entries[4].buffer = layer_params_buffers[i]; + layer_bg_entries[4].size = sizeof(CNNv2LayerParams); + + layer_bg_entries[5].binding = 5; + layer_bg_entries[5].textureView = input_view; + + WGPUBindGroupDescriptor layer_bg_desc = {}; + layer_bg_desc.layout = layer_bgl; + layer_bg_desc.entryCount = 6; + layer_bg_desc.entries = layer_bg_entries; + + WGPUBindGroup layer_bg = + wgpuDeviceCreateBindGroup(device, &layer_bg_desc); + + WGPUComputePassEncoder layer_pass = + wgpuCommandEncoderBeginComputePass(encoder, nullptr); + wgpuComputePassEncoderSetPipeline(layer_pass, layer_pipeline); + wgpuComputePassEncoderSetBindGroup(layer_pass, 0, layer_bg, 0, nullptr); + + wgpuComputePassEncoderDispatchWorkgroups(layer_pass, workgroups_x, + workgroups_y, 1); + + wgpuComputePassEncoderEnd(layer_pass); + wgpuComputePassEncoderRelease(layer_pass); + wgpuBindGroupRelease(layer_bg); + } + + WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, nullptr); + wgpuQueueSubmit(queue, 1, &commands); + wgpuCommandBufferRelease(commands); + wgpuCommandEncoderRelease(encoder); + + wgpuDevicePoll(device, true, nullptr); + + // Readback final result (from last layer's output texture) + printf("Reading pixels from GPU...\n"); + size_t final_layer_idx = (layer_info.size()) % 2; + std::vector<uint8_t> pixels = readback_rgba32uint_to_bgra8( + device, queue, layer_textures[final_layer_idx], width, height); + + if (pixels.empty()) { + fprintf(stderr, "Error: GPU readback failed\n"); + for (auto buf : layer_params_buffers) wgpuBufferRelease(buf); + wgpuComputePipelineRelease(layer_pipeline); + wgpuBindGroupLayoutRelease(layer_bgl); + wgpuBindGroupRelease(static_bg); + wgpuComputePipelineRelease(static_pipeline); + wgpuBufferRelease(static_params_buffer); + wgpuTextureViewRelease(static_features_view); + wgpuTextureRelease(static_features_tex); + wgpuTextureViewRelease(layer_views[0]); + wgpuTextureViewRelease(layer_views[1]); + wgpuTextureRelease(layer_textures[0]); + wgpuTextureRelease(layer_textures[1]); + wgpuBufferRelease(weights_buffer); + wgpuTextureViewRelease(input_view); + return false; + } + + // Debug hex dump + if (args.debug_hex) { + printf("First 8 pixels (BGRA hex):\n"); + for (int i = 0; i < 8 && i < width * height; ++i) { + const uint8_t b = pixels[i * 4 + 0]; + const uint8_t g = pixels[i * 4 + 1]; + const uint8_t r = pixels[i * 4 + 2]; + const uint8_t a = pixels[i * 4 + 3]; + printf(" [%d] 0x%02X%02X%02X%02X (RGBA)\n", i, r, g, b, a); + } + } + + // Save output + bool success; + if (args.output_png) { + printf("Saving PNG to '%s'...\n", args.output_path); + success = save_png(args.output_path, pixels, width, height); + } else { + printf("Saving PPM to '%s'...\n", args.output_path); + success = save_ppm(args.output_path, pixels, width, height); + } + + if (success) { + printf("Done! Output saved to '%s'\n", args.output_path); + } + + // Cleanup + for (auto buf : layer_params_buffers) wgpuBufferRelease(buf); + wgpuComputePipelineRelease(layer_pipeline); + wgpuBindGroupLayoutRelease(layer_bgl); + wgpuBindGroupRelease(static_bg); + wgpuComputePipelineRelease(static_pipeline); + wgpuBufferRelease(static_params_buffer); + wgpuTextureViewRelease(static_features_view); + wgpuTextureRelease(static_features_tex); + wgpuTextureViewRelease(layer_views[0]); + wgpuTextureViewRelease(layer_views[1]); + wgpuTextureRelease(layer_textures[0]); + wgpuTextureRelease(layer_textures[1]); + wgpuBufferRelease(weights_buffer); + wgpuTextureViewRelease(input_view); + + return success; +} + int main(int argc, char** argv) { // Parse arguments Args args; @@ -292,6 +945,19 @@ int main(int argc, char** argv) { printf("Loaded %dx%d image from '%s'\n", width, height, args.input_path); + // Branch based on CNN version + if (args.cnn_version == 2) { + bool success = process_cnn_v2(device, queue, instance, input_texture, + width, height, args); + wgpuTextureRelease(input_texture); + SamplerCache::Get().clear(); + fixture.shutdown(); + return success ? 0 : 1; + } + + // CNN v1 processing below + printf("Using CNN v1 (render pipeline architecture)\n"); + // Create input texture view const WGPUTextureViewDescriptor view_desc = { .format = WGPUTextureFormat_BGRA8Unorm, |