[Binary] Support extracting offloading files from COFF

[jhuber6]

Summary

This patch adds initial support for extracting offloading binaries from
COFF objects. This is a first step to allow building offloading files
on Windows targets with the new driver.

Depends on D136796

[jhenderson]

To my eyes, this looks reasonale, but I don't know enough about COFF to be able to comfortably review this. Best get another opinion.

[tra]

Added @rnk for Windows know-how.

[jhuber6]

In D136855#3899604, @tra wrote:

Added @rnk for Windows know-how.

A question I have for COFF, is if the only way to identify these sections would be via the name. For ELF we use a custom section type to easily identify the fat binary section, but it doesn't seem that there is a COFF equivalent. I thought about checking for the IMAGE_SCN_LNK_REMOVE and IMAGE_SCN_MEM_DISCARDABLE flags, but that's not guaranteed to be correct.

[rnk]

Seems reasonable to me.

llvm/lib/Object/OffloadBinary.cpp

75–76

Rather than duplicating the code, I think it makes sense to do the ELF check inside the loop. You can construct an ELFSectionRef from a regular SectionRef if the object is an ELF file.

[jhuber6]

Addressing comments.

[rnk]

llvm/lib/Object/OffloadBinary.cpp

80

static_cast doesn't seem idiomatic, calling the constructor like ELFSectionRef(Sec).getType() != ... seems more straightforward.

llvm/lib/Object/OffloadBinary.cpp

 #include "llvm/BinaryFormat/Magic.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/MC/StringTableBuilder.h"
 #include "llvm/Object/Archive.h"
 #include "llvm/Object/ArchiveWriter.h"
 #include "llvm/Object/Binary.h"
+#include "llvm/Object/COFF.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/Error.h"
 #include "llvm/Object/IRObjectFile.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/FileOutputBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 
while (Offset < Contents.getBuffer().size()) {
     Binaries.emplace_back(std::move(*NewBinaryOrErr), std::move(BufferCopy));
 
     Offset += Binary.getSize();
   }
 
   return Error::success();
 }
 
-// Extract offloading binaries from an Object file \p Obj.
-Error extractFromBinary(const ObjectFile &Obj,
-                        SmallVectorImpl<OffloadFile> &Binaries) {
+// Extract offloading binaries from an ELF Object file \p Obj.
+Error extractFromELF(const ObjectFile &Obj,
+  assert(Obj.isELF() && "Invalid file type");
+
   for (ELFSectionRef Sec : Obj.sections()) {
     if (Sec.getType() != ELF::SHT_LLVM_OFFLOADING)
       continue;
 
     Expected<StringRef> Buffer = Sec.getContents();
     if (!Buffer)
       return Buffer.takeError();
 
     MemoryBufferRef Contents(*Buffer, Obj.getFileName());
     if (Error Err = extractOffloadFiles(Contents, Binaries))
       return Err;
   }
 
   return Error::success();
 }
 
+// Extract offloading binaries from a COFF Object file \p Obj.
+Error extractFromCOFF(const ObjectFile &Obj,
+                      SmallVectorImpl<OffloadFile> &Binaries) {
+  assert(Obj.isCOFF() || Obj.isCOFFImportFile() && "Invalid file type");
+
+  for (SectionRef Sec : Obj.sections()) {
+    Expected<StringRef> NameOrErr = Sec.getName();
+    if (!NameOrErr)
+      return NameOrErr.takeError();
+
+    // COFF has no section types so we rely on the name of the section.
+    if (!NameOrErr->equals(".llvm.offloading"))
+      continue;
+
+    Expected<StringRef> Buffer = Sec.getContents();
+    if (!Buffer)
+      return Buffer.takeError();
+
+    MemoryBufferRef Contents(*Buffer, Obj.getFileName());
+    if (Error Err = extractOffloadFiles(Contents, Binaries))
+      return Err;
+  }
+
+  return Error::success();
+}
+
 Error extractFromBitcode(MemoryBufferRef Buffer,
                          SmallVectorImpl<OffloadFile> &Binaries) {
   LLVMContext Context;
   SMDiagnostic Err;
   std::unique_ptr<Module> M = getLazyIRModule(
       MemoryBuffer::getMemBuffer(Buffer, /*RequiresNullTerminator=*/false), Err,
       Context);
   if (!M)
case file_magic::bitcode:
     return extractFromBitcode(Buffer, Binaries);
   case file_magic::elf_relocatable:
   case file_magic::elf_executable:
   case file_magic::elf_shared_object: {
     Expected<std::unique_ptr<ObjectFile>> ObjFile =
         ObjectFile::createObjectFile(Buffer, Type);
     if (!ObjFile)
       return ObjFile.takeError();
-    return extractFromBinary(*ObjFile->get(), Binaries);
+    return extractFromELF(*ObjFile->get(), Binaries);
+  }
+  case file_magic::coff_object: {
+    Expected<std::unique_ptr<ObjectFile>> ObjFile =
+        ObjectFile::createObjectFile(Buffer, Type);
+    if (!ObjFile)
+      return ObjFile.takeError();
+    return extractFromCOFF(*ObjFile->get(), Binaries);
   }
   case file_magic::archive: {
     Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
         object::Archive::create(Buffer);
     if (!LibFile)
       return LibFile.takeError();
     return extractFromArchive(*LibFile->get(), Binaries);
   }

llvm/test/tools/llvm-objdump/Offloading/coff.test

+# RUN: yaml2obj %S/Inputs/binary.yaml -o %t.bin
+# RUN: yaml2obj %s -o %t.coff
+# RUN: llvm-objcopy --add-section .llvm.offloading=%t.bin %t.coff
+# RUN: llvm-objdump --offloading %t.coff | FileCheck %s --match-full-lines --strict-whitespace --implicit-check-not={{.}}
+
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: []
+sections:
+  - Name:            .rdata
+    Characteristics: []
+  - Name:            .llvm.offloading
+    Characteristics: [ IMAGE_SCN_LNK_REMOVE, IMAGE_SCN_MEM_DISCARDABLE ]
+    Alignment:       8
+symbols:
+
+#      CHECK:{{.*}}file format coff-x86-64
+# CHECK-EMPTY:
+# CHECK-NEXT:OFFLOADING IMAGE [0]:
+# CHECK-NEXT:kind            llvm ir
+# CHECK-NEXT:arch            gfx908
+# CHECK-NEXT:triple          amdgcn-amd-amdhsa
+# CHECK-NEXT:producer        openmp
+# CHECK-EMPTY:
+# CHECK-NEXT:OFFLOADING IMAGE [1]:
+# CHECK-NEXT:kind            llvm ir
+# CHECK-NEXT:arch            gfx90a
+# CHECK-NEXT:triple          amdgcn-amd-amdhsa
+# CHECK-NEXT:producer        openmp
+# CHECK-EMPTY:
+# CHECK-NEXT:OFFLOADING IMAGE [2]:
+# CHECK-NEXT:kind            cubin
+# CHECK-NEXT:arch            sm_52
+# CHECK-NEXT:triple          nvptx64-nvidia-cuda
+# CHECK-NEXT:producer        openmp
+# CHECK-EMPTY:
+# CHECK-NEXT:OFFLOADING IMAGE [3]:
+# CHECK-NEXT:kind            <none>
+# CHECK-NEXT:arch            sm_70
+# CHECK-NEXT:triple          nvptx64-nvidia-cuda
+# CHECK-NEXT:producer        none

llvm/test/tools/llvm-objdump/Offloading/non-elf.test

 # RUN: yaml2obj %s -o %t
 # RUN: llvm-objdump --offloading %t 2>&1 | FileCheck -DFILENAME=%t %s
 
---- !COFF
-header:
-  Machine:         IMAGE_FILE_MACHINE_AMD64
-  Characteristics: []
-sections:
-  - Name:            .rdata
-    Characteristics: []
-    SectionData:     00
-symbols:
+--- !mach-o
+FileHeader:
+  magic:      0xFEEDFACE
+  cputype:    0x00000007
+  cpusubtype: 0x00000003
+  filetype:   0x00000001
+  ncmds:      0
+  sizeofcmds: 0
+  flags:      0x00002000
 
-# CHECK: warning: '[[FILENAME]]': --offloading is currently only supported for ELF targets
+# CHECK: warning: '[[FILENAME]]': --offloading is currently only supported for COFF and ELF targets

llvm/tools/llvm-objdump/OffloadDump.cpp

static void printBinary(const OffloadBinary &OB, uint64_t Index) {
   outs() << left_justify("arch", 16) << OB.getArch() << "\n";
   outs() << left_justify("triple", 16) << OB.getTriple() << "\n";
   outs() << left_justify("producer", 16)
          << getOffloadKindName(OB.getOffloadKind()) << "\n";
 }
 
 /// Print the embedded offloading contents of an ObjectFile \p O.
 void llvm::dumpOffloadBinary(const ObjectFile &O) {
-  if (!O.isELF()) {
-    reportWarning("--offloading is currently only supported for ELF targets",
-                  O.getFileName());
+  if (!O.isELF() && !O.isCOFF()) {
+    reportWarning(
+        "--offloading is currently only supported for COFF and ELF targets",
     return;
   }
 
   SmallVector<OffloadFile> Binaries;
   if (Error Err = extractOffloadBinaries(O.getMemoryBufferRef(), Binaries))
     reportError(O.getFileName(), "while extracting offloading files: " +
                                      toString(std::move(Err)));