[Object] Change OffloadBinary::write to return SmallString<0>

[MaskRay]

Summary

SmallString<0> is more flexible and avoids an unneeded copy in
ObjectYAML/OffloadEmitter.cpp.

[MaskRay]

update clang-offload-packager

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

Expected<OffloadFile> getInputBitcodeLibrary(StringRef Input) {
 
   OffloadingImage Image{};
   Image.TheImageKind = IMG_Bitcode;
   Image.TheOffloadKind = getOffloadKind(Kind);
   Image.StringData["triple"] = Triple;
   Image.StringData["arch"] = Arch;
   Image.Image = std::move(*ImageOrError);
 
-  std::unique_ptr<MemoryBuffer> Binary = OffloadBinary::write(Image);
+  std::unique_ptr<MemoryBuffer> Binary =
+      MemoryBuffer::getMemBufferCopy(OffloadBinary::write(Image));
   auto NewBinaryOrErr = OffloadBinary::create(*Binary);
   if (!NewBinaryOrErr)
     return NewBinaryOrErr.takeError();
   return OffloadFile(std::move(*NewBinaryOrErr), std::move(Binary));
 }
 
 std::string getMainExecutable(const char *Name) {
   void *Ptr = (void *)(intptr_t)&getMainExecutable;
if (!FileOrErr)
     return FileOrErr.takeError();
   return *FileOrErr;
 }
 
 Expected<SmallVector<std::unique_ptr<MemoryBuffer>>>
 bundleOpenMP(ArrayRef<OffloadingImage> Images) {
   SmallVector<std::unique_ptr<MemoryBuffer>> Buffers;
   for (const OffloadingImage &Image : Images)
-    Buffers.emplace_back(OffloadBinary::write(Image));
+    Buffers.emplace_back(
+        MemoryBuffer::getMemBufferCopy(OffloadBinary::write(Image)));
 
   return std::move(Buffers);
 }
 
 Expected<SmallVector<std::unique_ptr<MemoryBuffer>>>
 bundleCuda(ArrayRef<OffloadingImage> Images, const ArgList &Args) {
   SmallVector<std::pair<StringRef, StringRef>, 4> InputFiles;
   for (const OffloadingImage &Image : Images)

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

for (const auto &[Key, Value] : Args) {
               getImageKind(sys::path::extension(Value).drop_front());
         ImageBinary.Image = std::move(*ObjectOrErr);
       } else if (Key == "kind") {
         ImageBinary.TheOffloadKind = getOffloadKind(Value);
       } else {
         ImageBinary.StringData[Key] = Value;
       }
     }
-    std::unique_ptr<MemoryBuffer> Buffer = OffloadBinary::write(ImageBinary);
-    if (Buffer->getBufferSize() % OffloadBinary::getAlignment() != 0)
+    llvm::SmallString<0> Buffer = OffloadBinary::write(ImageBinary);
+    if (Buffer.size() % OffloadBinary::getAlignment() != 0)
       return createStringError(inconvertibleErrorCode(),
                                "Offload binary has invalid size alignment");
-    OS << Buffer->getBuffer();
+    OS << Buffer;
   }
 
   if (Error E = writeFile(OutputFile,
                           StringRef(BinaryData.begin(), BinaryData.size())))
     return E;
   return Error::success();
 }
 

llvm/include/llvm/Object/OffloadBinary.h

 // complex it should be moved to a standard binary format like msgpack or ELF.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_OBJECT_OFFLOADBINARY_H
 #define LLVM_OBJECT_OFFLOADBINARY_H
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <memory>
 
 namespace llvm {
 
struct OffloadingImage {
     MapVector<StringRef, StringRef> StringData;
     std::unique_ptr<MemoryBuffer> Image;
   };
 
   /// Attempt to parse the offloading binary stored in \p Data.
   static Expected<std::unique_ptr<OffloadBinary>> create(MemoryBufferRef);
 
   /// Serialize the contents of \p File to a binary buffer to be read later.
-  static std::unique_ptr<MemoryBuffer> write(const OffloadingImage &);
+  static SmallString<0> write(const OffloadingImage &);
 
   static uint64_t getAlignment() { return 8; }
 
   ImageKind getImageKind() const { return TheEntry->TheImageKind; }
   OffloadKind getOffloadKind() const { return TheEntry->TheOffloadKind; }
   uint32_t getVersion() const { return TheHeader->Version; }
   uint32_t getFlags() const { return TheEntry->Flags; }
   uint64_t getSize() const { return TheHeader->Size; }

llvm/lib/Object/OffloadBinary.cpp

OffloadBinary::create(MemoryBufferRef Buf) {
   if (TheEntry->ImageOffset > Buf.getBufferSize() ||
       TheEntry->StringOffset > Buf.getBufferSize())
     return errorCodeToError(object_error::unexpected_eof);
 
   return std::unique_ptr<OffloadBinary>(
       new OffloadBinary(Buf, TheHeader, TheEntry));
 }
 
-std::unique_ptr<MemoryBuffer>
-OffloadBinary::write(const OffloadingImage &OffloadingData) {
+SmallString<0> OffloadBinary::write(const OffloadingImage &OffloadingData) {
   // Create a null-terminated string table with all the used strings.
   StringTableBuilder StrTab(StringTableBuilder::ELF);
   for (auto &KeyAndValue : OffloadingData.StringData) {
     StrTab.add(KeyAndValue.first);
     StrTab.add(KeyAndValue.second);
   }
   StrTab.finalize();
 
SmallString<0> OffloadBinary::write(const OffloadingImage &OffloadingData) {
   TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind;
   TheEntry.Flags = OffloadingData.Flags;
   TheEntry.StringOffset = sizeof(Header) + sizeof(Entry);
   TheEntry.NumStrings = OffloadingData.StringData.size();
 
   TheEntry.ImageOffset = BinaryDataSize;
   TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
 
-  SmallVector<char> Data;
+  SmallString<0> Data;
   Data.reserve(TheHeader.Size);
   raw_svector_ostream OS(Data);
   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
   OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
   for (auto &KeyAndValue : OffloadingData.StringData) {
     uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize;
     StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.first),
                     Offset + StrTab.getOffset(KeyAndValue.second)};
     OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
   }
   StrTab.write(OS);
   // Add padding to required image alignment.
   OS.write_zeros(TheEntry.ImageOffset - OS.tell());
   OS << OffloadingData.Image->getBuffer();
 
   // Add final padding to required alignment.
   assert(TheHeader.Size >= OS.tell() && "Too much data written?");
   OS.write_zeros(TheHeader.Size - OS.tell());
   assert(TheHeader.Size == OS.tell() && "Size mismatch");
 
-  return MemoryBuffer::getMemBufferCopy(OS.str());
+  return Data;
 }
 
 Error object::extractOffloadBinaries(MemoryBufferRef Buffer,
                                      SmallVectorImpl<OffloadFile> &Binaries) {
   file_magic Type = identify_magic(Buffer.getBuffer());
   switch (Type) {
   case file_magic::bitcode:
     return extractFromBitcode(Buffer, Binaries);

llvm/lib/ObjectYAML/OffloadEmitter.cpp

if (Member.StringEntries)
         Image.StringData[Entry.Key] = Entry.Value;
 
     SmallVector<char, 1024> Data;
     raw_svector_ostream OS(Data);
     if (Member.Content)
       Member.Content->writeAsBinary(OS);
     Image.Image = MemoryBuffer::getMemBufferCopy(OS.str());
 
-    std::unique_ptr<MemoryBuffer> Binary = object::OffloadBinary::write(Image);
-
     // Copy the data to a new buffer so we can modify the bytes directly.
-    SmallVector<char> NewBuffer;
-    std::copy(Binary->getBufferStart(), Binary->getBufferEnd(),
-              std::back_inserter(NewBuffer));
+    auto Buffer = object::OffloadBinary::write(Image);
     auto *TheHeader =
-        reinterpret_cast<object::OffloadBinary::Header *>(&NewBuffer[0]);
+        reinterpret_cast<object::OffloadBinary::Header *>(&Buffer[0]);
     if (Doc.Version)
       TheHeader->Version = *Doc.Version;
     if (Doc.Size)
       TheHeader->Size = *Doc.Size;
     if (Doc.EntryOffset)
       TheHeader->EntryOffset = *Doc.EntryOffset;
     if (Doc.EntrySize)
       TheHeader->EntrySize = *Doc.EntrySize;
 
-    Out.write(NewBuffer.begin(), NewBuffer.size());
+    Out.write(Buffer.begin(), Buffer.size());
   }
 
   return true;
 }
 
 } // namespace yaml
 } // namespace llvm

llvm/unittests/Object/OffloadingTest.cpp

TEST(OffloadingTest, checkOffloadingBinary) {
 
   OffloadBinary::OffloadingImage Data;
   Data.TheImageKind = static_cast<ImageKind>(KindDist(Rng));
   Data.TheOffloadKind = static_cast<OffloadKind>(KindDist(Rng));
   Data.Flags = KindDist(Rng);
   Data.StringData = StringData;
   Data.Image = std::move(ImageData);
 
-  auto BinaryBuffer = OffloadBinary::write(Data);
-
+  auto BinaryBuffer =
+      MemoryBuffer::getMemBufferCopy(OffloadBinary::write(Data));
   auto BinaryOrErr = OffloadBinary::create(*BinaryBuffer);
   if (!BinaryOrErr)
     FAIL();
 
   // Make sure we get the same data out.
   auto &Binary = **BinaryOrErr;
   ASSERT_EQ(Data.TheImageKind, Binary.getImageKind());
   ASSERT_EQ(Data.TheOffloadKind, Binary.getOffloadKind());