[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

72–78

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

77

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()) {
 
     Offset += Binary.getSize();
   }
 
   return Error::success();
 }
 
 // Extract offloading binaries from an Object file \p Obj.
-Error extractFromBinary(const ObjectFile &Obj,
+Error extractFromObject(const ObjectFile &Obj,
                         SmallVectorImpl<OffloadFile> &Binaries) {
-  for (ELFSectionRef Sec : Obj.sections()) {
-    if (Sec.getType() != ELF::SHT_LLVM_OFFLOADING)
+  assert((Obj.isELF() || Obj.isCOFF()) && "Invalid file type");
+
+  for (SectionRef Sec : Obj.sections()) {
+    // ELF files contain a section with the LLVM_OFFLOADING type.
+    if (Obj.isELF() &&
+        static_cast<ELFSectionRef>(Sec).getType() != ELF::SHT_LLVM_OFFLOADING)
       continue;
 
+    // COFF has no section types so we rely on the name of the section.
+    if (Obj.isCOFF()) {
+      Expected<StringRef> NameOrErr = Sec.getName();
+      if (!NameOrErr)
+        return NameOrErr.takeError();
+
+      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;
   }

 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);
   case file_magic::elf_relocatable:
   case file_magic::elf_executable:
-  case file_magic::elf_shared_object: {
+  case file_magic::elf_shared_object:
+  case file_magic::coff_object: {
     Expected<std::unique_ptr<ObjectFile>> ObjFile =
         ObjectFile::createObjectFile(Buffer, Type);
     if (!ObjFile)
       return ObjFile.takeError();
-    return extractFromBinary(*ObjFile->get(), Binaries);
+    return extractFromObject(*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)));