Handle more broad pixel transformations

This handles pixel data transformation across a wider array of pixel types.

At some point this may make sense ot move to a GPU kernel for more efficient exectuion.

Author: llvm-beanz

tools/offloader/WritePNG.cpp

@@ -40,24 +40,42 @@ llvm::Error WritePNG(llvm::StringRef OutputPath, const Resource &R) {
   // 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 * (Depth / 8));
+  std::unique_ptr<uint8_t[]> TranslatedPixels;
+  bool PixelsNeedTranslation =
+      R.Format == DataFormat::Float32 || R.Format == DataFormat::Float64;
+  if (PixelsNeedTranslation)
+    TranslatedPixels =
+        std::make_unique<uint8_t[]>(PixelComponents * (Depth / 8));
 
   switch (R.Format) {
   case DataFormat::Float32:
     if (Depth == 16)
-      TranslatePixelData(reinterpret_cast<uint16_t *>(PixelData.get()),
+      TranslatePixelData(reinterpret_cast<uint16_t *>(TranslatedPixels.get()),
                          reinterpret_cast<const float *>(R.Data.get()),
                          PixelComponents);
     else
-      TranslatePixelData(reinterpret_cast<uint8_t *>(PixelData.get()),
+      TranslatePixelData(reinterpret_cast<uint8_t *>(TranslatedPixels.get()),
                          reinterpret_cast<const float *>(R.Data.get()),
                          PixelComponents);
     break;
+  case DataFormat::Float64:
+    if (Depth == 16)
+      TranslatePixelData(reinterpret_cast<uint16_t *>(TranslatedPixels.get()),
+                         reinterpret_cast<const double *>(R.Data.get()),
+                         PixelComponents);
+    else
+      TranslatePixelData(reinterpret_cast<uint8_t *>(TranslatedPixels.get()),
+                         reinterpret_cast<const double *>(R.Data.get()),
+                         PixelComponents);
+    break;
   default:
     llvm_unreachable("Unsupported format for png output");
   }
 
+  uint8_t *PixelData = PixelsNeedTranslation
+                           ? TranslatedPixels.get()
+                           : reinterpret_cast<uint8_t *>(R.Data.get());
+
   png_structp PNG =
       png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL); // 8
   if (!PNG)
@@ -95,7 +113,7 @@ llvm::Error WritePNG(llvm::StringRef OutputPath, const Resource &R) {
   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 * DepthBytes) - RowSize;
+  uint8_t *Row = PixelData + (PixelComponents * DepthBytes) - RowSize;
   for (int I = 0; I < Height; ++I, Row -= RowSize)
     Rows[I] = (png_bytep)Row;