Capture image properties in YAML

llvm-beanz · llvm-beanz · commit 076b2aa77203 · 2024-12-21T13:28:20.000-06:00
This refactors some code throughout the offloader to read the output properties from the YAML. This enables naming the output and also specifying color depth and image dimensions for the PNG writer.
diff --git a/include/API/Pipeline.h b/include/API/Pipeline.h
@@ -14,6 +14,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <memory>
+#include <string>
 
 namespace hlsltest {
 
@@ -43,6 +44,13 @@ struct DirectXBinding {
   uint32_t Space;
 };
 
+struct OutputProperties {
+  std::string Name;
+  int Height;
+  int Width;
+  int Depth;
+};
+
 struct Resource {
   DataFormat Format;
   int Channels;
@@ -51,6 +59,7 @@ struct Resource {
   size_t Size;
   std::unique_ptr<char[]> Data;
   DirectXBinding DXBinding;
+  OutputProperties OutputProps;
 
   bool isRaw() const {
     return RawSize > 0;
@@ -124,6 +133,10 @@ template <> struct MappingTraits<hlsltest::DirectXBinding> {
   static void mapping(IO &I, hlsltest::DirectXBinding &B);
 };
 
+template <> struct MappingTraits<hlsltest::OutputProperties> {
+  static void mapping(IO &I, hlsltest::OutputProperties &P);
+};
+
 template <> struct ScalarEnumerationTraits<hlsltest::DataFormat> {
   static void enumeration(IO &I, hlsltest::DataFormat &V) {
 #define ENUM_CASE(Val) I.enumCase(V, #Val, hlsltest::DataFormat::Val)
diff --git a/lib/API/Pipeline.cpp b/lib/API/Pipeline.cpp
@@ -69,13 +69,23 @@ void MappingTraits<hlsltest::Resource>::mapping(IO &I, hlsltest::Resource &R) {
     DATA_CASE(Float32, float)
     DATA_CASE(Float64, double)
   }
+
   I.mapRequired("DirectXBinding", R.DXBinding);
+  I.mapOptional("OutputProps", R.OutputProps);
 }
 
 void MappingTraits<hlsltest::DirectXBinding>::mapping(
     IO &I, hlsltest::DirectXBinding &B) {
   I.mapRequired("Register", B.Register);
   I.mapRequired("Space", B.Space);
 }
+
+void MappingTraits<hlsltest::OutputProperties>::mapping(
+    IO &I, hlsltest::OutputProperties &P) {
+  I.mapRequired("Name", P.Name);
+  I.mapRequired("Height", P.Height);
+  I.mapRequired("Width", P.Width);
+  I.mapRequired("Depth", P.Depth);
+}
 } // namespace yaml
 } // namespace llvm
diff --git a/test/Basic/Mandelbrot.test b/test/Basic/Mandelbrot.test
@@ -64,17 +64,23 @@ DescriptorSets:
       DirectXBinding:
         Register: 0
         Space: 0
+      OutputProps:
+        Name: Output
+        Height: 4096
+        Width: 4096
+        Depth: 16
 ...
 #--- end
 
-# REQUIRES: imagediff
 # this test doesn't yet do anything... so we should skip it.
 
+# REQUIRES: goldenimage
 # RUN: split-file %s %t
 # RUN: %if DirectX %{ dxc -T cs_6_0 -Fo %t.dxil %t/source.hlsl %}
-# RUN: %if DirectX %{ %offloader %t/pipeline.yaml %t.dxil | FileCheck %s %}
+# RUN: %if DirectX %{ %offloader %t/pipeline.yaml %t.dxil -r Output -o %t/output.png %}
 # RUN: %if Vulkan %{ dxc -T cs_6_0 -spirv -Fo %t.spv %t/source.hlsl %}
-# RUN: %if Vulkan %{ %offloader %t/pipeline.yaml %t.spv | FileCheck %s %}
+# RUN: %if Vulkan %{ %offloader %t/pipeline.yaml %t.spv -r Output -o %t/output.png %}
 # RUN: %if Metal %{ dxc -T cs_6_0 -Fo %t.dxil %t/source.hlsl %}
 # RUN: %if Metal %{ metal-shaderconverter %t.dxil -o=%t.metallib %}
-# RUN: %if Metal %{ %offloader %t/pipeline.yaml %t.metallib | FileCheck %s %}
+# RUN: %if Metal %{ %offloader %t/pipeline.yaml %t.metallib -r Output -o %t/output.png %}
+# RUN: cmp %t/output.png %goldenimage_dir/hlsl/Basic/Mandelbrot.png
diff --git a/tools/offloader/WritePNG.cpp b/tools/offloader/WritePNG.cpp
@@ -13,30 +13,46 @@
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/SwapByteOrder.h"
 
+#include <limits.h>
 #include <png.h>
 #include <stdio.h>
 
 using namespace hlsltest;
 
+template <typename DstType, typename SrcType>
+void TranslatePixelData(DstType *Dst, const SrcType *Src, uint64_t Count) {
+  for (uint64_t I = 0; I < Count; ++I, ++Src, ++Dst) {
+    SrcType Val = *Src;
+    Val *= std::numeric_limits<DstType>::max();
+    *Dst = static_cast<DstType>(Val);
+    if (sizeof(DstType) > 1 && !llvm::sys::IsBigEndianHost)
+      llvm::sys::swapByteOrder(*Dst);
+  }
+}
+
 llvm::Error WritePNG(llvm::StringRef OutputPath, const Resource &R) {
-  // TODO: Fixme
-  int Height = 4096;
-  int Width = 4096;
+  int Height = R.OutputProps.Height;
+  int Width = R.OutputProps.Width;
+  int Depth = R.OutputProps.Depth;
+  assert(Depth == 16 || Depth == 8 && "Color depth should be 8 or 16 bits.");
   // Translate pixel data into uint8_t.
   uint64_t Stride = R.getSingleElementSize();
   uint64_t PixelComponents = R.Size / Stride;
   std::unique_ptr<uint8_t[]> PixelData =
-      std::make_unique<uint8_t[]>(PixelComponents);
+      std::make_unique<uint8_t[]>(PixelComponents * (Depth / 8));
 
-  const uint8_t *Src = reinterpret_cast<const uint8_t *>(R.Data.get());
   switch (R.Format) {
   case DataFormat::Float32:
-    for (uint64_t I = 0; I < PixelComponents; ++I, Src += Stride) {
-      float Val = *reinterpret_cast<const float *>(Src);
-      Val *= 255.f;
-      PixelData[I] = static_cast<uint8_t>(Val);
-    }
+    if (Depth == 16)
+      TranslatePixelData(reinterpret_cast<uint16_t *>(PixelData.get()),
+                         reinterpret_cast<const float *>(R.Data.get()),
+                         PixelComponents);
+    else
+      TranslatePixelData(reinterpret_cast<uint8_t *>(PixelData.get()),
+                         reinterpret_cast<const float *>(R.Data.get()),
+                         PixelComponents);
     break;
   default:
     llvm_unreachable("Unsupported format for png output");
@@ -63,7 +79,7 @@ llvm::Error WritePNG(llvm::StringRef OutputPath, const Resource &R) {
   if (!F)
     return llvm::createStringError(std::errc::io_error, "Failed openiong file");
   png_init_io(PNG, F);
-  png_set_IHDR(PNG, PNGInfo, Width, Height, 8, PNG_COLOR_TYPE_RGB_ALPHA,
+  png_set_IHDR(PNG, PNGInfo, Width, Height, Depth, PNG_COLOR_TYPE_RGB_ALPHA,
                PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE,
                PNG_FILTER_TYPE_BASE);
   png_colorp ColorP =
@@ -76,9 +92,10 @@ llvm::Error WritePNG(llvm::StringRef OutputPath, const Resource &R) {
   png_set_packing(PNG);
 
   png_bytepp Rows = (png_bytepp)png_malloc(PNG, Height * sizeof(png_bytep));
-  uint64_t RowSize = Width * R.Channels;
+  uint64_t DepthBytes = Depth / 8;
+  uint64_t RowSize = Width * R.Channels * DepthBytes;
   // Step one row back from the end
-  uint8_t *Row = PixelData.get() + PixelComponents - RowSize;
+  uint8_t *Row = PixelData.get() + (PixelComponents * DepthBytes) - RowSize;
   for (int I = 0; I < Height; ++I, Row -= RowSize)
     Rows[I] = (png_bytep)Row;
 
diff --git a/tools/offloader/offloader.cpp b/tools/offloader/offloader.cpp
@@ -44,16 +44,14 @@ static cl::opt<std::string> OutputFilename("o", cl::desc("Output filename"),
                                            cl::value_desc("filename"),
                                            cl::init("-"));
 
-static cl::opt<std::string> ImageOutput("r",
-                                        cl::desc("Resource index to output"),
-                                        cl::value_desc("set,resource"),
-                                        cl::init(""));
+static cl::opt<std::string>
+    ImageOutput("r", cl::desc("Resource name to output as png"),
+                cl::value_desc("<name>"), cl::init(""));
 
 static cl::opt<bool>
     Quiet("quiet", cl::desc("Suppress printing the pipeline as output"));
 
-static cl::opt<bool>
-    UseWarp("warp", cl::desc("Use warp"));
+static cl::opt<bool> UseWarp("warp", cl::desc("Use warp"));
 
 llvm::Error WritePNG(llvm::StringRef, const Resource &);
 
@@ -101,9 +99,8 @@ int run() {
   }
 
   if (UseWarp && APIToUse != GPUAPI::DirectX)
-    ExitOnErr(
-        createStringError(std::errc::executable_format_error,
-                          "WARP required DirectX API"));
+    ExitOnErr(createStringError(std::errc::executable_format_error,
+                                "WARP required DirectX API"));
 
   if (APIToUse == GPUAPI::Unknown)
     ExitOnErr(
@@ -139,28 +136,17 @@ int run() {
       Out->keep();
       return 0;
     }
-    llvm::Regex R("^([0-9]+),([0-9]+)$");
-    llvm::SmallVector<llvm::StringRef, 2> Matches;
-    if (!R.match(ImageOutput, &Matches))
-      ExitOnErr(
-          createStringError(std::errc::invalid_argument,
-                            "Image output argument must be specified as "
-                            "<set>,<resource> (e.g. \"0,1\", \"3,0\", etc)"));
-    uint64_t Set = 0, Resource = 0;
-    if (!to_integer(Matches[1], Set) || !to_integer(Matches[2], Resource))
-      ExitOnErr(
-          createStringError(std::errc::invalid_argument,
-                            "Image output argument must be specified as "
-                            "<set>,<resource> (e.g. \"0,1\", \"3,0\", etc)"));
-    if (Set >= PipelineDesc.Sets.size())
-      ExitOnErr(createStringError(std::errc::invalid_argument,
-                                  "Specified descriptor set out of range"));
-    if (Resource >= PipelineDesc.Sets[Set].Resources.size())
-      ExitOnErr(createStringError(std::errc::invalid_argument,
-                                  "Specified descriptor index out of range"));
+    for (const auto &S : PipelineDesc.Sets) {
+      for (const auto &R : S.Resources) {
+        if (R.OutputProps.Name == ImageOutput) {
+          ExitOnErr(WritePNG(OutputFilename, R));
+          return 0;
+        }
+      }
+    }
+
     ExitOnErr(
-        WritePNG(OutputFilename, PipelineDesc.Sets[Set].Resources[Resource]));
-    return 0;
+        createStringError(Twine("No descriptor with name ") + ImageOutput));
   }
   return 1;
 }
