[Comgr] Add new Action to compile SPIR-V to Relocatable (llvm#1252)

In this patch, we add a new action:

AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE

That accepts a set of .spv files, translates them to .bc files, extracts any embedded @llvm.cmdline flags, and then compiles to a set of relocatable .o files.

Author: lamb-j

amd/comgr/include/amd_comgr.h.in

@@ -1794,6 +1794,23 @@ typedef enum amd_comgr_action_kind_s {
    */
   AMD_COMGR_ACTION_UNBUNDLE = 0xF,
 
+   /**
+   * Compile each source SPIR-V object in @p input into a relocatable.
+   * For each successful compilation, add a relocatable object to @p result
+   *
+   * We accomplish this by first translating the .spv files to .bc via the
+   * SPIR-V translator. We then extract any relevant -cc1 flags from the embedded
+   * @llvm.cmdline variable. Finally, we compile the bitcode to a reloctable,
+   * appending any extracted flags.
+   *
+   * Return @p AMD_COMGR_STATUS_ERROR if any translation, flag extraction, or
+   * compilation fails.
+   *
+   * Return @p AMD_COMGR_STATUS_ERROR_INVALID_ARGUMENT
+   * if any input is not SPIR-V.
+   */
+  AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE = 0x10,
+
    /**
    * Translate each source SPIR-V object in @p input into LLVM IR Bitcode.
    * For each successful translation, add a bc object to @p result   *

amd/comgr/src/comgr-compiler.cpp

@@ -58,6 +58,7 @@
 #include "clang/FrontendTool/Utils.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
@@ -877,7 +878,8 @@ amd_comgr_status_t AMDGPUCompiler::removeTmpDirs() {
 #endif
 }
 
-amd_comgr_status_t AMDGPUCompiler::processFile(const char *InputFilePath,
+amd_comgr_status_t AMDGPUCompiler::processFile(DataObject *Input,
+                                               const char *InputFilePath,
                                                const char *OutputFilePath) {
   SmallVector<const char *, 128> Argv = Args;
 
@@ -899,6 +901,12 @@ amd_comgr_status_t AMDGPUCompiler::processFile(const char *InputFilePath,
     Argv.push_back("-save-temps=obj");
   }
 
+  // Add SPIR-V flags
+  for (auto Flag : Input->SpirvFlags) {
+    Argv.push_back("-Xclang");
+    Argv.push_back(Flag);
+  }
+
   Argv.push_back(InputFilePath);
 
   Argv.push_back("-o");
@@ -910,6 +918,12 @@ amd_comgr_status_t AMDGPUCompiler::processFile(const char *InputFilePath,
 amd_comgr_status_t
 AMDGPUCompiler::processFiles(amd_comgr_data_kind_t OutputKind,
                              const char *OutputSuffix) {
+  return processFiles(OutputKind, OutputSuffix, InSet);
+}
+
+amd_comgr_status_t
+AMDGPUCompiler::processFiles(amd_comgr_data_kind_t OutputKind,
+                             const char *OutputSuffix, DataSet *InSet) {
   for (auto *Input : InSet->DataObjects) {
     if (Input->DataKind != AMD_COMGR_DATA_KIND_INCLUDE) {
       continue;
@@ -947,7 +961,7 @@ AMDGPUCompiler::processFiles(amd_comgr_data_kind_t OutputKind,
     auto OutputFilePath = getFilePath(Output, OutputDir);
 
     if (auto Status =
-            processFile(InputFilePath.c_str(), OutputFilePath.c_str())) {
+            processFile(Input, InputFilePath.c_str(), OutputFilePath.c_str())) {
       return Status;
     }
 
@@ -1888,11 +1902,106 @@ amd_comgr_status_t AMDGPUCompiler::linkToExecutable() {
   return amd_comgr_data_set_add(OutSetT, OutputT);
 }
 
+// TODO: Generalize this list to include all -cc1 flags and arguments that are
+//   still valid in a bitcode compilation context
+static inline const std::unordered_set<std::string_view> ValidSpirvFlags{
+    "-fapprox-func",
+    "-fcolor-diagnostics",
+    "-fconvergent-functions",
+    "-fexceptions",
+    "-ffast-math",
+    "-ffinite-math-only",
+    "-ffp-contract=fast",
+    "-ffp-contract=fast-honor-pragmas",
+    "-ffp-contract=on",
+    "-fgpu-rdc",
+    "-finline-functions",
+    "-fno-autolink",
+    "-fno-experimental-relative-c++-abi-vtables",
+    "-fno-rounding-math",
+    "-fno-signed-zeros",
+    "-fno-threadsafe-statics",
+    "-freciprocal-math",
+    "-funsafe-math-optimizations",
+    "-fvisibility=hidden",
+    "-O0",
+    "-O1",
+    "-O2",
+    "-O3",
+    "--save-temps"};
+
+amd_comgr_status_t AMDGPUCompiler::extractSpirvFlags(DataSet *BcSet) {
+
+  for (auto *Bc : BcSet->DataObjects) {
+    // Create SPIR-V IR Module from Bitcode Buffer
+    SMDiagnostic SMDiag;
+    LLVMContext Context;
+    Context.setDiagnosticHandler(
+        std::make_unique<AMDGPUCompilerDiagnosticHandler>(this->LogS), true);
+
+    auto Mod = getLazyIRModule(
+        MemoryBuffer::getMemBuffer(StringRef(Bc->Data, Bc->Size), "", false),
+        SMDiag, Context, true);
+
+    if (!Mod) {
+      SMDiag.print("SPIR-V Bitcode", LogS, /* ShowColors */ false);
+      return AMD_COMGR_STATUS_ERROR;
+    }
+
+    if (verifyModule(*Mod, &LogS))
+      return AMD_COMGR_STATUS_ERROR;
+
+    // Fetch @llvm.cmdline
+    GlobalVariable *CmdLine = Mod->getNamedGlobal("llvm.cmdline");
+
+    // Return if no @llvm.cmdline
+    if (!CmdLine)
+      return AMD_COMGR_STATUS_SUCCESS;
+
+    if (ConstantDataSequential *CDS =
+            dyn_cast<ConstantDataSequential>(CmdLine->getInitializer())) {
+
+      // Add each valid null-terminated '\0' string to Flags
+      std::string Tmp;
+      StringRef CmdLineRaw = CDS->getRawDataValues();
+      std::stringstream ss(CmdLineRaw.str());
+      while (getline(ss, Tmp, '\0')) {
+        if (Tmp == "--hipstdpar" || Tmp == "-amdgpu-enable-hipstdpar") {
+          Bc->SpirvFlags.push_back("-mllvm");
+          Bc->SpirvFlags.push_back("-amdgpu-enable-hipstdpar");
+        } else if (Tmp == "-amdgpu-spill-cfi-saved-regs") {
+          Bc->SpirvFlags.push_back("-mllvm");
+          Bc->SpirvFlags.push_back("-amdgpu-spill-cfi-saved-regs");
+        } else if (ValidSpirvFlags.count(Tmp)) {
+          Bc->SpirvFlags.push_back(Saver.save(Tmp.c_str()).data());
+        }
+      }
+    }
+
+    // COV5 required for SPIR-V
+    Bc->SpirvFlags.push_back("-mcode-object-version=5");
+
+    if (env::shouldEmitVerboseLogs()) {
+      LogS << "        SPIR-V Flags: " << Bc->Name << "\n";
+      for (auto Flag : Bc->SpirvFlags)
+        LogS << "          " << Flag << "\n";
+    }
+  }
+
+  return AMD_COMGR_STATUS_SUCCESS;
+}
+
 amd_comgr_status_t AMDGPUCompiler::translateSpirvToBitcode() {
+  return translateSpirvToBitcodeImpl(InSet, DataSet::convert(OutSetT));
+}
+
+amd_comgr_status_t
+AMDGPUCompiler::translateSpirvToBitcodeImpl(DataSet *SpirvInSet,
+                                            DataSet *BcOutSet) {
 #ifdef COMGR_DISABLE_SPIRV
-  LogS << "Calling AMDGPUCompiler::translateSpirvToBitcode() not supported. "
-       << "Comgr is built with -DCOMGR_DISABLE_SPIRV. Re-build LLVM and Comgr "
-       << "with LLVM-SPIRV-Translator support to continue.\n";
+  LogS << "Calling AMDGPUCompiler::translateSpirvToBitcodeImpl() not "
+       << "supported. Comgr is built with -DCOMGR_DISABLE_SPIRV. Re-build LLVM "
+       << "and Comgr with LLVM-SPIRV-Translator support to continue.\n";
   return AMD_COMGR_STATUS_ERROR;
 #else
   if (auto Status = createTmpDirs()) {
@@ -1901,7 +2010,7 @@ amd_comgr_status_t AMDGPUCompiler::translateSpirvToBitcode() {
 
   auto Cache = CommandCache::get(LogS);
 
-  for (auto *Input : InSet->DataObjects) {
+  for (auto *Input : SpirvInSet->DataObjects) {
 
     if (env::shouldSaveTemps()) {
       if (auto Status = outputToFile(Input, getFilePath(Input, InputDir))) {
@@ -1939,7 +2048,8 @@ amd_comgr_status_t AMDGPUCompiler::translateSpirvToBitcode() {
     Output->setName(std::string(Input->Name) + std::string(".bc"));
     Output->setData(OutBuf);
 
-    if (auto Status = amd_comgr_data_set_add(OutSetT, OutputT)) {
+    if (auto Status =
+            amd_comgr_data_set_add(DataSet::convert(BcOutSet), OutputT)) {
       return Status;
     }
 
@@ -1960,6 +2070,50 @@ amd_comgr_status_t AMDGPUCompiler::translateSpirvToBitcode() {
 #endif
 }
 
+amd_comgr_status_t AMDGPUCompiler::compileSpirvToRelocatable() {
+  if (auto Status = createTmpDirs()) {
+    return Status;
+  }
+
+  for (auto *Input : InSet->DataObjects) {
+    if (Input->DataKind != AMD_COMGR_DATA_KIND_SPIRV)
+      return AMD_COMGR_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  // Translate .spv to .bc
+  amd_comgr_data_set_t TranslatedSpirvT;
+  if (auto Status = amd_comgr_create_data_set(&TranslatedSpirvT))
+    return Status;
+  DataSet *TranslatedSpirv = DataSet::convert(TranslatedSpirvT);
+
+  if (auto Status = translateSpirvToBitcodeImpl(InSet, TranslatedSpirv))
+    return Status;
+
+  // Extract relevant -cc1 flags from @llvm.cmdline
+  if (auto Status = extractSpirvFlags(TranslatedSpirv))
+    return Status;
+
+  // Compile bitcode to relocatable
+  if (ActionInfo->IsaName) {
+    if (auto Status = addTargetIdentifierFlags(ActionInfo->IsaName)) {
+      return Status;
+    }
+  }
+
+  if (ActionInfo->ShouldLinkDeviceLibs) {
+    if (auto Status = addDeviceLibraries()) {
+      return Status;
+    }
+  }
+
+  Args.push_back("-c");
+
+  Args.push_back("-mllvm");
+  Args.push_back("-amdgpu-internalize-symbols");
+
+  return processFiles(AMD_COMGR_DATA_KIND_RELOCATABLE, ".o", TranslatedSpirv);
+}
+
 AMDGPUCompiler::AMDGPUCompiler(DataAction *ActionInfo, DataSet *InSet,
                                DataSet *OutSet, raw_ostream &LogS)
     : ActionInfo(ActionInfo), InSet(InSet), OutSetT(DataSet::convert(OutSet)),

amd/comgr/src/comgr-compiler.h

@@ -69,19 +69,25 @@ class AMDGPUCompiler {
 
   amd_comgr_status_t createTmpDirs();
   amd_comgr_status_t removeTmpDirs();
-  amd_comgr_status_t processFile(const char *InputFilePath,
+  amd_comgr_status_t processFile(DataObject *Input, const char *InputFilePath,
                                  const char *OutputFilePath);
   /// Process each file in @c InSet individually, placing output in @c OutSet.
   amd_comgr_status_t processFiles(amd_comgr_data_kind_t OutputKind,
                                   const char *OutputSuffix);
+  amd_comgr_status_t processFiles(amd_comgr_data_kind_t OutputKind,
+                                  const char *OutputSuffix, DataSet *InSet);
   amd_comgr_status_t addIncludeFlags();
   amd_comgr_status_t addTargetIdentifierFlags(llvm::StringRef IdentStr,
                                               bool CompilingSrc);
   amd_comgr_status_t addCompilationFlags();
   amd_comgr_status_t addDeviceLibraries();
+  amd_comgr_status_t extractSpirvFlags(DataSet *BcSet);
 
   amd_comgr_status_t executeInProcessDriver(llvm::ArrayRef<const char *> Args);
 
+  amd_comgr_status_t translateSpirvToBitcodeImpl(DataSet *SpirvInSet,
+                                                 DataSet *BcOutSet);
+
 public:
   AMDGPUCompiler(DataAction *ActionInfo, DataSet *InSet, DataSet *OutSet,
                  llvm::raw_ostream &LogS);
@@ -98,6 +104,7 @@ class AMDGPUCompiler {
   amd_comgr_status_t linkToRelocatable();
   amd_comgr_status_t linkToExecutable();
   amd_comgr_status_t compileToExecutable();
+  amd_comgr_status_t compileSpirvToRelocatable();
   amd_comgr_status_t translateSpirvToBitcode();
 
   amd_comgr_language_t getLanguage() const { return ActionInfo->Language; }

amd/comgr/src/comgr.cpp

@@ -187,6 +187,8 @@ amd_comgr_status_t dispatchCompilerAction(amd_comgr_action_kind_t ActionKind,
     return Compiler.compileToBitcode(true);
   case AMD_COMGR_ACTION_COMPILE_SOURCE_TO_EXECUTABLE:
     return Compiler.compileToExecutable();
+  case AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE:
+    return Compiler.compileSpirvToRelocatable();
   case AMD_COMGR_ACTION_TRANSLATE_SPIRV_TO_BC:
     return Compiler.translateSpirvToBitcode();
 
@@ -291,6 +293,8 @@ StringRef getActionKindName(amd_comgr_action_kind_t ActionKind) {
     return "AMD_COMGR_ACTION_COMPILE_SOURCE_TO_EXECUTABLE";
   case AMD_COMGR_ACTION_UNBUNDLE:
     return "AMD_COMGR_ACTION_UNBUNDLE";
+  case AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE:
+    return "AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE";
   case AMD_COMGR_ACTION_TRANSLATE_SPIRV_TO_BC:
     return "AMD_COMGR_ACTION_TRANSLATE_SPIRV_TO_BC";
   }
@@ -1376,6 +1380,7 @@ amd_comgr_status_t AMD_COMGR_API
     case AMD_COMGR_ACTION_COMPILE_SOURCE_TO_RELOCATABLE:
     case AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC:
     case AMD_COMGR_ACTION_COMPILE_SOURCE_TO_EXECUTABLE:
+    case AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE:
     case AMD_COMGR_ACTION_TRANSLATE_SPIRV_TO_BC:
       ActionStatus = dispatchCompilerAction(ActionKind, ActionInfoP, InputSetP,
                                             ResultSetP, *LogP);

amd/comgr/src/comgr.h

@@ -131,6 +131,7 @@ struct DataObject {
   DataSymbol *DataSym;
   std::vector<std::string> MangledNames;
   std::map<std::string, std::string> NameExpressionMap;
+  llvm::SmallVector<const char *, 128> SpirvFlags;
 
 private:
   std::unique_ptr<llvm::MemoryBuffer> Buffer;

amd/comgr/test-lit/CMakeLists.txt

@@ -41,5 +41,6 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
 add_comgr_lit_binary(source-to-bc-with-dev-libs c)
 add_comgr_lit_binary(spirv-translator c)
 add_comgr_lit_binary(compile-minimal-test c)
+add_comgr_lit_binary(spirv-to-reloc c)
 
 add_dependencies(check-comgr test-lit)

amd/comgr/test-lit/comgr-sources/spirv-to-reloc.c

@@ -0,0 +1,57 @@
+#include "amd_comgr.h"
+#include "common.h"
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+int main(int argc, char *argv[]) {
+  char *BufSpv;
+  size_t SizeSpv;
+  amd_comgr_data_t DataSpv;
+  amd_comgr_data_set_t DataSetSpv, DataSetReloc;
+  amd_comgr_action_info_t DataAction;
+  size_t Count;
+
+  if (argc != 3) {
+    fprintf(stderr, "Usage: spirv-to-reloc file.spv file.o\n");
+    exit(1);
+  }
+
+  SizeSpv = setBuf(argv[1], &BufSpv);
+
+  amd_comgr_(create_data(AMD_COMGR_DATA_KIND_SPIRV, &DataSpv));
+  amd_comgr_(set_data(DataSpv, SizeSpv, BufSpv));
+  amd_comgr_(set_data_name(DataSpv, "file.spv"));
+
+  amd_comgr_(create_data_set(&DataSetSpv));
+  amd_comgr_(data_set_add(DataSetSpv, DataSpv));
+
+  amd_comgr_(create_action_info(&DataAction));
+  amd_comgr_(action_info_set_language(DataAction, AMD_COMGR_LANGUAGE_HIP));
+  amd_comgr_(action_info_set_isa_name(DataAction, "amdgcn-amd-amdhsa--gfx900"));
+
+  amd_comgr_(create_data_set(&DataSetReloc));
+  amd_comgr_(do_action(AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE,
+                       DataAction, DataSetSpv, DataSetReloc));
+
+  amd_comgr_(
+      action_data_count(DataSetReloc, AMD_COMGR_DATA_KIND_RELOCATABLE, &Count));
+
+  if (Count != 1) {
+    printf("AMD_COMGR_ACTION_COMPILE_SPIRV_TO_RELOCATABLE Failed: "
+           "produced %zu RELOCATABLE objects (expected 1)\n",
+           Count);
+    exit(1);
+  }
+
+  amd_comgr_data_t DataReloc;
+  amd_comgr_(action_data_get_data(DataSetReloc, AMD_COMGR_DATA_KIND_RELOCATABLE,
+                                  0, &DataReloc));
+  dumpData(DataReloc, argv[2]);
+
+  amd_comgr_(release_data(DataSpv));
+  amd_comgr_(destroy_data_set(DataSetSpv));
+  amd_comgr_(destroy_data_set(DataSetReloc));
+  amd_comgr_(destroy_action_info(DataAction));
+  free(BufSpv);
+}