1 files changed, 46 insertions, 2 deletions

diff --git a/tools/cnn_test.cc b/tools/cnn_test.cc
index 39ed436..59f1d22 100644
--- a/tools/cnn_test.cc
+++ b/tools/cnn_test.cc
@@ -42,6 +42,8 @@ struct Args {
   float blend = 1.0f;
   bool output_png = true; // Default to PNG
   const char* save_intermediates = nullptr;
+  int num_layers = 3; // Default to 3 layers
+  bool debug_hex = false; // Print first 8 pixels as hex
 };
 
 // Parse command-line arguments
@@ -73,6 +75,14 @@ static bool parse_args(int argc, char** argv, Args* args) {
       }
     } else if (strcmp(argv[i], "--save-intermediates") == 0 && i + 1 < argc) {
       args->save_intermediates = argv[++i];
+    } else if (strcmp(argv[i], "--layers") == 0 && i + 1 < argc) {
+      args->num_layers = atoi(argv[++i]);
+      if (args->num_layers < 1 || args->num_layers > 10) {
+        fprintf(stderr, "Error: layers must be in range [1, 10]\n");
+        return false;
+      }
+    } else if (strcmp(argv[i], "--debug-hex") == 0) {
+      args->debug_hex = true;
     } else if (strcmp(argv[i], "--help") == 0) {
       return false;
     } else {
@@ -90,7 +100,9 @@ static void print_usage(const char* prog) {
   fprintf(stderr, "\nOPTIONS:\n");
   fprintf(stderr, "  --blend F                Final blend amount (0.0-1.0, default: 1.0)\n");
   fprintf(stderr, "  --format ppm|png         Output format (default: png)\n");
+  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, "  --help                   Show this help\n");
 }
 
@@ -273,6 +285,7 @@ int main(int argc, char** argv) {
   WGPUTexture input_texture =
       load_texture(device, queue, args.input_path, &width, &height);
   if (!input_texture) {
+    SamplerCache::Get().clear();
     fixture.shutdown();
     return 1;
   }
@@ -303,6 +316,7 @@ int main(int argc, char** argv) {
     if (pipeline_final) wgpuRenderPipelineRelease(pipeline_final);
     wgpuTextureViewRelease(input_view);
     wgpuTextureRelease(input_texture);
+    SamplerCache::Get().clear();
     fixture.shutdown();
     return 1;
   }
@@ -360,8 +374,8 @@ int main(int argc, char** argv) {
   WGPUSampler sampler =
       SamplerCache::Get().get_or_create(device, SamplerCache::clamp());
 
-  // Multi-layer processing (fixed 3 layers)
-  const int NUM_LAYERS = 3;
+  // Multi-layer processing
+  const int NUM_LAYERS = args.num_layers;
   int dst_idx = 0; // Index of texture to render to
 
   // First layer reads from input, subsequent layers read from previous output
@@ -427,6 +441,18 @@ int main(int argc, char** argv) {
       printf("Reading pixels from GPU...\n");
       std::vector<uint8_t> pixels = rt.read_pixels();
 
+      // Debug: print first 8 pixels as hex
+      if (args.debug_hex && !pixels.empty()) {
+        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);
+        }
+      }
+
       if (pixels.empty()) {
         fprintf(stderr, "Error: GPU readback failed\n");
         wgpuTextureViewRelease(intermediate_views[0]);
@@ -440,6 +466,7 @@ int main(int argc, char** argv) {
         wgpuBindGroupLayoutRelease(bgl);
         wgpuRenderPipelineRelease(pipeline_final);
         wgpuRenderPipelineRelease(pipeline_intermediate);
+        SamplerCache::Get().clear();
         fixture.shutdown();
         return 1;
       }
@@ -466,6 +493,7 @@ int main(int argc, char** argv) {
         wgpuBindGroupLayoutRelease(bgl);
         wgpuRenderPipelineRelease(pipeline_final);
         wgpuRenderPipelineRelease(pipeline_intermediate);
+        SamplerCache::Get().clear();
         fixture.shutdown();
         return 1;
       }
@@ -501,6 +529,18 @@ int main(int argc, char** argv) {
         std::vector<uint8_t> pixels = texture_readback_fp16_to_u8(
             device, queue, intermediate_textures[dst_idx], width, height);
 
+        // Debug: print first 8 pixels as hex
+        if (args.debug_hex && !pixels.empty()) {
+          printf("Layer %d first 8 pixels (BGRA hex):\n", layer);
+          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);
+          }
+        }
+
         if (!pixels.empty()) {
           save_png(layer_path, pixels, width, height);
         } else {
@@ -517,6 +557,9 @@ int main(int argc, char** argv) {
     }
   }
 
+  // Wait for all GPU work to complete before cleanup
+  wgpuDevicePoll(device, true, nullptr);
+
   // Cleanup
   wgpuTextureViewRelease(intermediate_views[0]);
   wgpuTextureViewRelease(intermediate_views[1]);
@@ -529,6 +572,7 @@ int main(int argc, char** argv) {
   wgpuRenderPipelineRelease(pipeline_final);
   wgpuTextureViewRelease(input_view);
   wgpuTextureRelease(input_texture);
+  SamplerCache::Get().clear();
   fixture.shutdown();
 
   return 0;