[mlir][gpu] Expose some utility functions from gpu-to-binary infra

mlir/include/mlir/Target/LLVM/ModuleToObject.h

@@ -44,6 +44,13 @@ class ModuleToObject {
   /// Runs the serialization pipeline, returning `std::nullopt` on error.
   virtual std::optional<SmallVector<char, 0>> run();
 
+  /// Translate LLVM module to textual ISA.
+  /// TODO: switch to SmallString
+  static FailureOr<std::string>
+  translateModuleToISA(llvm::Module &llvmModule,
+                       llvm::TargetMachine &targetMachine,
+                       function_ref<InFlightDiagnostic()> emitError);
+
 protected:
   // Hooks to be implemented by derived classes.
 
@@ -98,11 +105,6 @@ class ModuleToObject {
   /// Optimize the module.
   virtual LogicalResult optimizeModule(llvm::Module &module, int optL);
 
-  /// Utility function for translating to ISA, returns `std::nullopt` on
-  /// failure.
-  static std::optional<std::string>
-  translateToISA(llvm::Module &llvmModule, llvm::TargetMachine &targetMachine);
-
 protected:
   /// Module to transform to a binary object.
   Operation &module;

mlir/include/mlir/Target/LLVM/ROCDL/Utils.h

@@ -43,6 +43,15 @@ enum class AMDGCNLibraries : uint32_t {
   All = (LastLib << 1) - 1
 };
 
+/// Assembles ISA to an object code.
+FailureOr<SmallVector<char, 0>>
+assembleIsa(StringRef isa, StringRef targetTriple, StringRef chip,
+            StringRef features, function_ref<InFlightDiagnostic()> emitError);
+
+FailureOr<SmallVector<char, 0>>
+linkObjectCode(ArrayRef<char> objectCode, StringRef toolkitPath,
+               function_ref<InFlightDiagnostic()> emitError);
+
 /// Base class for all ROCDL serializations from GPU modules into binary
 /// strings. By default this class serializes into LLVM bitcode.
 class SerializeGPUModuleBase : public LLVM::ModuleToObject {
@@ -98,9 +107,6 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
   moduleToObjectImpl(const gpu::TargetOptions &targetOptions,
                      llvm::Module &llvmModule);
 
-  /// Returns the assembled ISA.
-  std::optional<SmallVector<char, 0>> assembleIsa(StringRef isa);
-
   /// ROCDL target attribute.
   ROCDLTargetAttr target;
 

mlir/lib/Target/LLVM/ModuleToObject.cpp

@@ -205,9 +205,9 @@ LogicalResult ModuleToObject::optimizeModule(llvm::Module &module,
   return success();
 }
 
-std::optional<std::string>
-ModuleToObject::translateToISA(llvm::Module &llvmModule,
-                               llvm::TargetMachine &targetMachine) {
+FailureOr<std::string> ModuleToObject::translateModuleToISA(
+    llvm::Module &llvmModule, llvm::TargetMachine &targetMachine,
+    function_ref<InFlightDiagnostic()> emitError) {
   std::string targetISA;
   llvm::raw_string_ostream stream(targetISA);
 
@@ -217,7 +217,7 @@ ModuleToObject::translateToISA(llvm::Module &llvmModule,
 
     if (targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
                                           llvm::CodeGenFileType::AssemblyFile))
-      return std::nullopt;
+      return emitError() << "Target machine cannot emit assembly";
 
     codegenPasses.run(llvmModule);
   }

mlir/lib/Target/LLVM/NVVM/Target.cpp

@@ -697,18 +697,19 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
     return std::nullopt;
   }
   moduleToObjectTimer.startTimer();
-  std::optional<std::string> serializedISA =
-      translateToISA(llvmModule, **targetMachine);
+  FailureOr<std::string> serializedISA =
+      translateModuleToISA(llvmModule, **targetMachine,
+                           [&]() { return getOperation().emitError(); });
   moduleToObjectTimer.stopTimer();
   llvmToISATimeInMs = moduleToObjectTimer.getTotalTime().getWallTime() * 1000;
   moduleToObjectTimer.clear();
-  if (!serializedISA) {
+  if (failed(serializedISA)) {
     getOperation().emitError() << "Failed translating the module to ISA.";
     return std::nullopt;
   }
 
   if (isaCallback)
-    isaCallback(serializedISA.value());
+    isaCallback(*serializedISA);
 
 #define DEBUG_TYPE "serialize-to-isa"
   LDBG() << "PTX for module: " << getOperation().getNameAttr() << "\n"

mlir/lib/Target/LLVM/ROCDL/Target.cpp

@@ -277,23 +277,19 @@ void SerializeGPUModuleBase::addControlVariables(
   }
 }
 
-std::optional<SmallVector<char, 0>>
-SerializeGPUModuleBase::assembleIsa(StringRef isa) {
-  auto loc = getOperation().getLoc();
-
-  StringRef targetTriple = this->triple;
-
+FailureOr<SmallVector<char, 0>>
+mlir::ROCDL::assembleIsa(StringRef isa, StringRef targetTriple, StringRef chip,
+                         StringRef features,
+                         function_ref<InFlightDiagnostic()> emitError) {
   SmallVector<char, 0> result;
   llvm::raw_svector_ostream os(result);
 
   llvm::Triple triple(llvm::Triple::normalize(targetTriple));
   std::string error;
   const llvm::Target *target =
       llvm::TargetRegistry::lookupTarget(triple, error);
-  if (!target) {
-    emitError(loc, Twine("failed to lookup target: ") + error);
-    return std::nullopt;
-  }
+  if (!target)
+    return emitError() << "failed to lookup target: " << error;
 
   llvm::SourceMgr srcMgr;
   srcMgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(isa), SMLoc());
@@ -330,76 +326,79 @@ SerializeGPUModuleBase::assembleIsa(StringRef isa) {
   std::unique_ptr<llvm::MCTargetAsmParser> tap(
       target->createMCAsmParser(*sti, *parser, *mcii, mcOptions));
 
-  if (!tap) {
-    emitError(loc, "assembler initialization error");
-    return std::nullopt;
-  }
+  if (!tap)
+    return emitError() << "assembler initialization error";
 
   parser->setTargetParser(*tap);
   parser->Run(false);
   return std::move(result);
 }
 
-std::optional<SmallVector<char, 0>>
-SerializeGPUModuleBase::compileToBinary(const std::string &serializedISA) {
-  // Assemble the ISA.
-  std::optional<SmallVector<char, 0>> isaBinary = assembleIsa(serializedISA);
-
-  if (!isaBinary) {
-    getOperation().emitError() << "failed during ISA assembling";
-    return std::nullopt;
-  }
-
+FailureOr<SmallVector<char, 0>>
+mlir::ROCDL::linkObjectCode(ArrayRef<char> objectCode, StringRef toolkitPath,
+                            function_ref<InFlightDiagnostic()> emitError) {
   // Save the ISA binary to a temp file.
   int tempIsaBinaryFd = -1;
   SmallString<128> tempIsaBinaryFilename;
   if (llvm::sys::fs::createTemporaryFile("kernel%%", "o", tempIsaBinaryFd,
-                                         tempIsaBinaryFilename)) {
-    getOperation().emitError()
-        << "failed to create a temporary file for dumping the ISA binary";
-    return std::nullopt;
-  }
+                                         tempIsaBinaryFilename))
+    return emitError()
+           << "failed to create a temporary file for dumping the ISA binary";
+
   llvm::FileRemover cleanupIsaBinary(tempIsaBinaryFilename);
   {
     llvm::raw_fd_ostream tempIsaBinaryOs(tempIsaBinaryFd, true);
-    tempIsaBinaryOs << StringRef(isaBinary->data(), isaBinary->size());
+    tempIsaBinaryOs << StringRef(objectCode.data(), objectCode.size());
     tempIsaBinaryOs.flush();
   }
 
   // Create a temp file for HSA code object.
   SmallString<128> tempHsacoFilename;
-  if (llvm::sys::fs::createTemporaryFile("kernel", "hsaco",
-                                         tempHsacoFilename)) {
-    getOperation().emitError()
-        << "failed to create a temporary file for the HSA code object";
-    return std::nullopt;
-  }
+  if (llvm::sys::fs::createTemporaryFile("kernel", "hsaco", tempHsacoFilename))
+    return emitError()
+           << "failed to create a temporary file for the HSA code object";
+
   llvm::FileRemover cleanupHsaco(tempHsacoFilename);
 
   llvm::SmallString<128> lldPath(toolkitPath);
   llvm::sys::path::append(lldPath, "llvm", "bin", "ld.lld");
   int lldResult = llvm::sys::ExecuteAndWait(
       lldPath,
       {"ld.lld", "-shared", tempIsaBinaryFilename, "-o", tempHsacoFilename});
-  if (lldResult != 0) {
-    getOperation().emitError() << "lld invocation failed";
-    return std::nullopt;
-  }
+  if (lldResult != 0)
+    return emitError() << "lld invocation failed";
 
   // Load the HSA code object.
   auto hsacoFile =
       llvm::MemoryBuffer::getFile(tempHsacoFilename, /*IsText=*/false);
-  if (!hsacoFile) {
-    getOperation().emitError()
-        << "failed to read the HSA code object from the temp file";
-    return std::nullopt;
-  }
+  if (!hsacoFile)
+    return emitError()
+           << "failed to read the HSA code object from the temp file";
 
   StringRef buffer = (*hsacoFile)->getBuffer();
 
   return SmallVector<char, 0>(buffer.begin(), buffer.end());
 }
 
+std::optional<SmallVector<char, 0>>
+SerializeGPUModuleBase::compileToBinary(const std::string &serializedISA) {
+  auto errCallback = [&]() { return getOperation().emitError(); };
+  // Assemble the ISA.
+  FailureOr<SmallVector<char, 0>> isaBinary = ROCDL::assembleIsa(
+      serializedISA, this->triple, this->chip, this->features, errCallback);
+
+  if (failed(isaBinary))
+    return std::nullopt;
+
+  // Link the object code.
+  FailureOr<SmallVector<char, 0>> linkedCode =
+      ROCDL::linkObjectCode(*isaBinary, toolkitPath, errCallback);
+  if (failed(linkedCode))
+    return std::nullopt;
+
+  return linkedCode;
+}
+
 std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
     const gpu::TargetOptions &targetOptions, llvm::Module &llvmModule) {
   // Return LLVM IR if the compilation target is offload.
@@ -422,9 +421,10 @@ std::optional<SmallVector<char, 0>> SerializeGPUModuleBase::moduleToObjectImpl(
   }
 
   // Translate the Module to ISA.
-  std::optional<std::string> serializedISA =
-      translateToISA(llvmModule, **targetMachine);
-  if (!serializedISA) {
+  FailureOr<std::string> serializedISA =
+      translateModuleToISA(llvmModule, **targetMachine,
+                           [&]() { return getOperation().emitError(); });
+  if (failed(serializedISA)) {
     getOperation().emitError() << "failed translating the module to ISA";
     return std::nullopt;
   }

mlir/lib/Target/LLVM/XeVM/Target.cpp

@@ -301,9 +301,10 @@ SPIRVSerializer::moduleToObject(llvm::Module &llvmModule) {
   // Return SPIRV if the compilation target is `assembly`.
   if (targetOptions.getCompilationTarget() ==
       gpu::CompilationTarget::Assembly) {
-    std::optional<std::string> serializedISA =
-        translateToISA(llvmModule, **targetMachine);
-    if (!serializedISA) {
+    FailureOr<std::string> serializedISA =
+        translateModuleToISA(llvmModule, **targetMachine,
+                             [&]() { return getGPUModuleOp().emitError(); });
+    if (failed(serializedISA)) {
       getGPUModuleOp().emitError() << "Failed translating the module to ISA."
                                    << triple << ", can't compile with LLVM\n";
       return std::nullopt;