[ELF] Attempt to set the OS when using 'makeTriple()' (#76992)

Summary:
This patch fixes up the `makeTriple()` interface to emit append the
operating system information when it is readily avaialble from the ELF.
The main motivation for this is so the GPU architectures can be easily
identified correctly when given and ELF. E.g. we want
`amdgpu-amd-amdhsa` as the output and not `amdgpu--`.

This required adding support for the CUDA OS/ABI, which is easily found
to be `0x33` when using `readelf`.

6 files changed, 80 insertions, 9 deletions

diff --git a/llvm/include/llvm/BinaryFormat/ELF.h b/llvm/include/llvm/BinaryFormat/ELF.h
index 0f968eac36e7..9b8128a9ec40 100644
--- a/llvm/include/llvm/BinaryFormat/ELF.h
+++ b/llvm/include/llvm/BinaryFormat/ELF.h
@@ -356,6 +356,7 @@ enum {
   ELFOSABI_AROS = 15,          // AROS
   ELFOSABI_FENIXOS = 16,       // FenixOS
   ELFOSABI_CLOUDABI = 17,      // Nuxi CloudABI
+  ELFOSABI_CUDA = 51,          // NVIDIA CUDA architecture.
   ELFOSABI_FIRST_ARCH = 64,    // First architecture-specific OS ABI
   ELFOSABI_AMDGPU_HSA = 64,    // AMD HSA runtime
   ELFOSABI_AMDGPU_PAL = 65,    // AMD PAL runtime
diff --git a/llvm/include/llvm/Object/ELFObjectFile.h b/llvm/include/llvm/Object/ELFObjectFile.h
index da78e11b678d..da72b0e335b9 100644
--- a/llvm/include/llvm/Object/ELFObjectFile.h
+++ b/llvm/include/llvm/Object/ELFObjectFile.h
@@ -454,6 +454,7 @@ public:
   uint8_t getBytesInAddress() const override;
   StringRef getFileFormatName() const override;
   Triple::ArchType getArch() const override;
+  Triple::OSType getOS() const override;
   Expected<uint64_t> getStartAddress() const override;
 
   unsigned getPlatformFlags() const override { return EF.getHeader().e_flags; }
@@ -1382,6 +1383,35 @@ template <class ELFT> Triple::ArchType ELFObjectFile<ELFT>::getArch() const {
   }
 }
 
+template <class ELFT> Triple::OSType ELFObjectFile<ELFT>::getOS() const {
+  switch (EF.getHeader().e_ident[ELF::EI_OSABI]) {
+  case ELF::ELFOSABI_NETBSD:
+    return Triple::NetBSD;
+  case ELF::ELFOSABI_LINUX:
+    return Triple::Linux;
+  case ELF::ELFOSABI_HURD:
+    return Triple::Hurd;
+  case ELF::ELFOSABI_SOLARIS:
+    return Triple::Solaris;
+  case ELF::ELFOSABI_AIX:
+    return Triple::AIX;
+  case ELF::ELFOSABI_FREEBSD:
+    return Triple::FreeBSD;
+  case ELF::ELFOSABI_OPENBSD:
+    return Triple::OpenBSD;
+  case ELF::ELFOSABI_CUDA:
+    return Triple::CUDA;
+  case ELF::ELFOSABI_AMDGPU_HSA:
+    return Triple::AMDHSA;
+  case ELF::ELFOSABI_AMDGPU_PAL:
+    return Triple::AMDPAL;
+  case ELF::ELFOSABI_AMDGPU_MESA3D:
+    return Triple::Mesa3D;
+  default:
+    return Triple::UnknownOS;
+  }
+}
+
 template <class ELFT>
 Expected<uint64_t> ELFObjectFile<ELFT>::getStartAddress() const {
   return EF.getHeader().e_entry;
diff --git a/llvm/include/llvm/Object/ObjectFile.h b/llvm/include/llvm/Object/ObjectFile.h
index c254fc2ccfde..8c868c7643ed 100644
--- a/llvm/include/llvm/Object/ObjectFile.h
+++ b/llvm/include/llvm/Object/ObjectFile.h
@@ -337,6 +337,7 @@ public:
 
   virtual StringRef getFileFormatName() const = 0;
   virtual Triple::ArchType getArch() const = 0;
+  virtual Triple::OSType getOS() const { return Triple::UnknownOS; }
   virtual Expected<SubtargetFeatures> getFeatures() const = 0;
   virtual std::optional<StringRef> tryGetCPUName() const {
     return std::nullopt;
diff --git a/llvm/lib/Object/ObjectFile.cpp b/llvm/lib/Object/ObjectFile.cpp
index ca921836b7f6..c05eb0a0468e 100644
--- a/llvm/lib/Object/ObjectFile.cpp
+++ b/llvm/lib/Object/ObjectFile.cpp
@@ -111,6 +111,10 @@ Triple ObjectFile::makeTriple() const {
   auto Arch = getArch();
   TheTriple.setArch(Triple::ArchType(Arch));
 
+  auto OS = getOS();
+  if (OS != Triple::UnknownOS)
+    TheTriple.setOS(OS);
+
   // For ARM targets, try to use the build attributes to build determine
   // the build target. Target features are also added, but later during
   // disassembly.
@@ -129,10 +133,13 @@ Triple ObjectFile::makeTriple() const {
     // XCOFF implies AIX.
     TheTriple.setOS(Triple::AIX);
     TheTriple.setObjectFormat(Triple::XCOFF);
-  }
-  else if (isGOFF()) {
+  } else if (isGOFF()) {
     TheTriple.setOS(Triple::ZOS);
     TheTriple.setObjectFormat(Triple::GOFF);
+  } else if (TheTriple.isAMDGPU()) {
+    TheTriple.setVendor(Triple::AMD);
+  } else if (TheTriple.isNVPTX()) {
+    TheTriple.setVendor(Triple::NVIDIA);
   }
 
   return TheTriple;
diff --git a/llvm/tools/llvm-readobj/ELFDumper.cpp b/llvm/tools/llvm-readobj/ELFDumper.cpp
index abf7ba6ba1c3..10797b83d3d9 100644
--- a/llvm/tools/llvm-readobj/ELFDumper.cpp
+++ b/llvm/tools/llvm-readobj/ELFDumper.cpp
@@ -1084,6 +1084,7 @@ const EnumEntry<unsigned> ElfOSABI[] = {
   {"AROS",         "AROS",                 ELF::ELFOSABI_AROS},
   {"FenixOS",      "FenixOS",              ELF::ELFOSABI_FENIXOS},
   {"CloudABI",     "CloudABI",             ELF::ELFOSABI_CLOUDABI},
+  {"CUDA",         "NVIDIA - CUDA",        ELF::ELFOSABI_CUDA},
   {"Standalone",   "Standalone App",       ELF::ELFOSABI_STANDALONE}
 };
 
@@ -1093,6 +1094,10 @@ const EnumEntry<unsigned> AMDGPUElfOSABI[] = {
   {"AMDGPU_MESA3D", "AMDGPU - MESA3D", ELF::ELFOSABI_AMDGPU_MESA3D}
 };
 
+const EnumEntry<unsigned> NVPTXElfOSABI[] = {
+  {"NVIDIA_CUDA", "NVIDIA - CUDA", ELF::ELFOSABI_CUDA},
+};
+
 const EnumEntry<unsigned> ARMElfOSABI[] = {
   {"ARM", "ARM", ELF::ELFOSABI_ARM}
 };
diff --git a/llvm/unittests/Object/ELFObjectFileTest.cpp b/llvm/unittests/Object/ELFObjectFileTest.cpp
index 2878ca088cd7..e7619ded494e 100644
--- a/llvm/unittests/Object/ELFObjectFileTest.cpp
+++ b/llvm/unittests/Object/ELFObjectFileTest.cpp
@@ -36,7 +36,8 @@ struct DataForTest {
 
   template <typename T>
   std::vector<uint8_t> makeElfData(uint8_t Class, uint8_t Encoding,
-                                   uint16_t Machine) {
+                                   uint16_t Machine, uint8_t OS,
+                                   uint16_t Flags) {
     T Ehdr{}; // Zero-initialise the header.
     Ehdr.e_ident[ELF::EI_MAG0] = 0x7f;
     Ehdr.e_ident[ELF::EI_MAG1] = 'E';
@@ -45,9 +46,11 @@ struct DataForTest {
     Ehdr.e_ident[ELF::EI_CLASS] = Class;
     Ehdr.e_ident[ELF::EI_DATA] = Encoding;
     Ehdr.e_ident[ELF::EI_VERSION] = 1;
+    Ehdr.e_ident[ELF::EI_OSABI] = OS;
     Ehdr.e_type = ELF::ET_REL;
     Ehdr.e_machine = Machine;
     Ehdr.e_version = 1;
+    Ehdr.e_flags = Flags;
     Ehdr.e_ehsize = sizeof(T);
 
     bool IsLittleEndian = Encoding == ELF::ELFDATA2LSB;
@@ -64,12 +67,13 @@ struct DataForTest {
     return Bytes;
   }
 
-  DataForTest(uint8_t Class, uint8_t Encoding, uint16_t Machine) {
+  DataForTest(uint8_t Class, uint8_t Encoding, uint16_t Machine,
+              uint8_t OS = ELF::ELFOSABI_NONE, uint16_t Flags = 0) {
     if (Class == ELF::ELFCLASS64)
-      Data = makeElfData<ELF::Elf64_Ehdr>(Class, Encoding, Machine);
+      Data = makeElfData<ELF::Elf64_Ehdr>(Class, Encoding, Machine, OS, Flags);
     else {
       assert(Class == ELF::ELFCLASS32);
-      Data = makeElfData<ELF::Elf32_Ehdr>(Class, Encoding, Machine);
+      Data = makeElfData<ELF::Elf32_Ehdr>(Class, Encoding, Machine, OS, Flags);
     }
   }
 };
@@ -287,6 +291,30 @@ TEST(ELFObjectFileTest, MachineTestForXtensa) {
     checkFormatAndArch(Data, Formats[Idx], Triple::xtensa);
 }
 
+TEST(ELFObjectFileTest, CheckOSAndTriple) {
+  std::tuple<uint16_t, uint8_t, StringRef> Formats[] = {
+      {ELF::EM_AMDGPU, ELF::ELFOSABI_AMDGPU_HSA, "amdgcn-amd-amdhsa"},
+      {ELF::EM_X86_64, ELF::ELFOSABI_LINUX, "x86_64--linux"},
+      {ELF::EM_X86_64, ELF::ELFOSABI_NETBSD, "x86_64--netbsd"},
+      {ELF::EM_X86_64, ELF::ELFOSABI_HURD, "x86_64--hurd"},
+      {ELF::EM_X86_64, ELF::ELFOSABI_SOLARIS, "x86_64--solaris"},
+      {ELF::EM_X86_64, ELF::ELFOSABI_AIX, "x86_64--aix"},
+      {ELF::EM_X86_64, ELF::ELFOSABI_FREEBSD, "x86_64--freebsd"},
+      {ELF::EM_X86_64, ELF::ELFOSABI_OPENBSD, "x86_64--openbsd"},
+      {ELF::EM_CUDA, ELF::ELFOSABI_CUDA, "nvptx64-nvidia-cuda"}};
+  for (auto [Machine, OS, Triple] : Formats) {
+    const DataForTest D(ELF::ELFCLASS64, ELF::ELFDATA2LSB, Machine, OS,
+                        ELF::EF_AMDGPU_MACH_AMDGCN_LAST);
+    Expected<std::unique_ptr<ObjectFile>> ELFObjOrErr =
+        object::ObjectFile::createELFObjectFile(
+            MemoryBufferRef(toStringRef(D.Data), "dummyELF"));
+    ASSERT_THAT_EXPECTED(ELFObjOrErr, Succeeded());
+
+    auto &ELFObj = **ELFObjOrErr;
+    EXPECT_EQ(Triple, ELFObj.makeTriple().getTriple());
+  }
+}
+
 // ELF relative relocation type test.
 TEST(ELFObjectFileTest, RelativeRelocationTypeTest) {
   EXPECT_EQ(ELF::R_CKCORE_RELATIVE, getELFRelativeRelocationType(ELF::EM_CSKY));
@@ -1273,13 +1301,12 @@ Sections:
 )";
 
   auto ErroringMatcher = [](const Elf_Shdr &Sec) -> Expected<bool> {
-    if(Sec.sh_type == ELF::SHT_PROGBITS)
+    if (Sec.sh_type == ELF::SHT_PROGBITS)
       return createError("This was supposed to fail.");
     return false;
   };
 
-  DoCheckFails(OneTextSection, ErroringMatcher,
-               "This was supposed to fail.");
+  DoCheckFails(OneTextSection, ErroringMatcher, "This was supposed to fail.");
 
   StringRef MissingRelocatableContent = R"(
 Sections: