[mlir][Target][LLVM] Adds an utility class for serializing operations to binary strings.

**For an explanation of these patches see D154153.**

Commit message:
This patch adds the utility base class `ModuleToObject`. This class provides an
interface for compiling module operations into binary strings, by default this
class serialize modules to LLVM bitcode.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D154100

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

@@ -0,0 +1,122 @@
+//===- ModuleToObject.h - Module to object base class -----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the base class for transforming operations into binary
+// objects.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_LLVM_MODULETOOBJECT_H
+#define MLIR_TARGET_LLVM_MODULETOOBJECT_H
+
+#include "mlir/IR/Operation.h"
+#include "llvm/IR/Module.h"
+
+namespace llvm {
+class TargetMachine;
+} // namespace llvm
+
+namespace mlir {
+namespace LLVM {
+class ModuleTranslation;
+/// Utility base class for transforming operations into binary objects, by
+/// default it returns the serialized LLVM bitcode for the module. The
+/// operations being transformed must be translatable into LLVM IR.
+class ModuleToObject {
+public:
+  ModuleToObject(Operation &module, StringRef triple, StringRef chip,
+                 StringRef features = {}, int optLevel = 3);
+  virtual ~ModuleToObject() = default;
+
+  /// Returns the operation being serialized.
+  Operation &getOperation();
+
+  /// Runs the serialization pipeline, returning `std::nullopt` on error.
+  virtual std::optional<SmallVector<char, 0>> run();
+
+protected:
+  // Hooks to be implemented by derived classes.
+
+  /// Hook for loading bitcode files, returns std::nullopt on failure.
+  virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
+  loadBitcodeFiles(llvm::Module &module, llvm::TargetMachine &targetMachine) {
+    return SmallVector<std::unique_ptr<llvm::Module>>();
+  }
+
+  /// Hook for performing additional actions on a loaded bitcode file.
+  virtual LogicalResult handleBitcodeFile(llvm::Module &module,
+                                          llvm::TargetMachine &targetMachine) {
+    return success();
+  }
+
+  /// Hook for performing additional actions on the llvmModule pre linking.
+  virtual void handleModulePreLink(llvm::Module &module,
+                                   llvm::TargetMachine &targetMachine) {}
+
+  /// Hook for performing additional actions on the llvmModule post linking.
+  virtual void handleModulePostLink(llvm::Module &module,
+                                    llvm::TargetMachine &targetMachine) {}
+
+  /// Serializes the LLVM IR bitcode to an object file, by default it serializes
+  /// to LLVM bitcode.
+  virtual std::optional<SmallVector<char, 0>>
+  moduleToObject(llvm::Module &llvmModule, llvm::TargetMachine &targetMachine);
+
+protected:
+  /// Create the target machine based on the target triple and chip.
+  std::unique_ptr<llvm::TargetMachine> createTargetMachine();
+
+  /// Loads a bitcode file from path.
+  std::unique_ptr<llvm::Module>
+  loadBitcodeFile(llvm::LLVMContext &context,
+                  llvm::TargetMachine &targetMachine, StringRef path);
+
+  /// Loads multiple bitcode files.
+  LogicalResult loadBitcodeFilesFromList(
+      llvm::LLVMContext &context, llvm::TargetMachine &targetMachine,
+      ArrayRef<std::string> fileList,
+      SmallVector<std::unique_ptr<llvm::Module>> &llvmModules,
+      bool failureOnError = true);
+
+  /// Translates the operation to LLVM IR.
+  std::unique_ptr<llvm::Module>
+  translateToLLVMIR(llvm::LLVMContext &llvmContext);
+
+  /// Link the llvmModule to other bitcode file.
+  LogicalResult linkFiles(llvm::Module &module,
+                          SmallVector<std::unique_ptr<llvm::Module>> &&libs);
+
+  /// Optimize the module.
+  LogicalResult optimizeModule(llvm::Module &module,
+                               llvm::TargetMachine &targetMachine, 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;
+
+  /// Target triple.
+  StringRef triple;
+
+  /// Target chip.
+  StringRef chip;
+
+  /// Target features.
+  StringRef features;
+
+  /// Optimization level.
+  int optLevel;
+};
+} // namespace LLVM
+} // namespace mlir
+
+#endif // MLIR_TARGET_LLVM_MODULETOOBJECT_H

mlir/lib/Target/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Cpp)
 add_subdirectory(SPIRV)
 add_subdirectory(LLVMIR)
+add_subdirectory(LLVM)

mlir/lib/Target/LLVM/CMakeLists.txt

@@ -0,0 +1,22 @@
+add_mlir_library(MLIRTargetLLVM
+  ModuleToObject.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Target/LLVM
+
+  DEPENDS
+  intrinsics_gen
+
+  LINK_COMPONENTS
+  Core
+  IPO
+  IRReader
+  Linker
+  MC
+  Passes
+  Support
+  Target
+  LINK_LIBS PUBLIC
+  MLIRExecutionEngineUtils
+  MLIRTargetLLVMIRExport
+)

mlir/lib/Target/LLVM/ModuleToObject.cpp

@@ -0,0 +1,227 @@
+//===- ModuleToObject.cpp - Module to object base class ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the base class for transforming Operations into binary
+// objects.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVM/ModuleToObject.h"
+
+#include "mlir/ExecutionEngine/OptUtils.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Export.h"
+#include "mlir/Target/LLVMIR/ModuleTranslation.h"
+
+#include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Linker/Linker.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/IPO/Internalize.h"
+
+using namespace mlir;
+using namespace mlir::LLVM;
+
+ModuleToObject::ModuleToObject(Operation &module, StringRef triple,
+                               StringRef chip, StringRef features, int optLevel)
+    : module(module), triple(triple), chip(chip), features(features),
+      optLevel(optLevel) {}
+
+Operation &ModuleToObject::getOperation() { return module; }
+
+std::unique_ptr<llvm::TargetMachine> ModuleToObject::createTargetMachine() {
+  std::string error;
+  // Load the target.
+  const llvm::Target *target =
+      llvm::TargetRegistry::lookupTarget(triple, error);
+  if (!target) {
+    getOperation().emitError() << "Failed to lookup target: " << error;
+    return {};
+  }
+
+  // Create the target machine using the target.
+  llvm::TargetMachine *machine =
+      target->createTargetMachine(triple, chip, features, {}, {});
+  if (!machine) {
+    getOperation().emitError() << "Failed to create the target machine.";
+    return {};
+  }
+  return std::unique_ptr<llvm::TargetMachine>{machine};
+}
+
+std::unique_ptr<llvm::Module>
+ModuleToObject::loadBitcodeFile(llvm::LLVMContext &context,
+                                llvm::TargetMachine &targetMachine,
+                                StringRef path) {
+  llvm::SMDiagnostic error;
+  std::unique_ptr<llvm::Module> library =
+      llvm::getLazyIRFileModule(path, error, context);
+  if (!library) {
+    getOperation().emitError() << "Failed loading file from " << path
+                               << ", error: " << error.getMessage();
+    return nullptr;
+  }
+  if (failed(handleBitcodeFile(*library, targetMachine))) {
+    return nullptr;
+  }
+  return library;
+}
+
+LogicalResult ModuleToObject::loadBitcodeFilesFromList(
+    llvm::LLVMContext &context, llvm::TargetMachine &targetMachine,
+    ArrayRef<std::string> fileList,
+    SmallVector<std::unique_ptr<llvm::Module>> &llvmModules,
+    bool failureOnError) {
+  for (const std::string &str : fileList) {
+    // Test if the path exists, if it doesn't abort.
+    StringRef pathRef = StringRef(str.data(), str.size());
+    if (!llvm::sys::fs::is_regular_file(pathRef)) {
+      getOperation().emitError()
+          << "File path: " << pathRef << " does not exist or is not a file.\n";
+      return failure();
+    }
+    // Load the file or abort on error.
+    if (auto bcFile = loadBitcodeFile(context, targetMachine, pathRef))
+      llvmModules.push_back(std::move(bcFile));
+    else if (failureOnError)
+      return failure();
+  }
+  return success();
+}
+
+std::unique_ptr<llvm::Module>
+ModuleToObject::translateToLLVMIR(llvm::LLVMContext &llvmContext) {
+  return translateModuleToLLVMIR(&getOperation(), llvmContext);
+}
+
+LogicalResult
+ModuleToObject::linkFiles(llvm::Module &module,
+                          SmallVector<std::unique_ptr<llvm::Module>> &&libs) {
+  if (libs.empty())
+    return success();
+  llvm::Linker linker(module);
+  for (std::unique_ptr<llvm::Module> &libModule : libs) {
+    // This bitcode linking imports the library functions into the module,
+    // allowing LLVM optimization passes (which must run after linking) to
+    // optimize across the libraries and the module's code. We also only import
+    // symbols if they are referenced by the module or a previous library since
+    // there will be no other source of references to those symbols in this
+    // compilation and since we don't want to bloat the resulting code object.
+    bool err = linker.linkInModule(
+        std::move(libModule), llvm::Linker::Flags::LinkOnlyNeeded,
+        [](llvm::Module &m, const StringSet<> &gvs) {
+          llvm::internalizeModule(m, [&gvs](const llvm::GlobalValue &gv) {
+            return !gv.hasName() || (gvs.count(gv.getName()) == 0);
+          });
+        });
+    // True is linker failure
+    if (err) {
+      getOperation().emitError("Unrecoverable failure during bitcode linking.");
+      // We have no guaranties about the state of `ret`, so bail
+      return failure();
+    }
+  }
+  return success();
+}
+
+LogicalResult ModuleToObject::optimizeModule(llvm::Module &module,
+                                             llvm::TargetMachine &targetMachine,
+                                             int optLevel) {
+  if (optLevel < 0 || optLevel > 3)
+    return getOperation().emitError()
+           << "Invalid optimization level: " << optLevel << ".";
+
+  targetMachine.setOptLevel(static_cast<llvm::CodeGenOpt::Level>(optLevel));
+
+  auto transformer =
+      makeOptimizingTransformer(optLevel, /*sizeLevel=*/0, &targetMachine);
+  auto error = transformer(&module);
+  if (error) {
+    InFlightDiagnostic mlirError = getOperation().emitError();
+    llvm::handleAllErrors(
+        std::move(error), [&mlirError](const llvm::ErrorInfoBase &ei) {
+          mlirError << "Could not optimize LLVM IR: " << ei.message() << "\n";
+        });
+    return mlirError;
+  }
+  return success();
+}
+
+std::optional<std::string>
+ModuleToObject::translateToISA(llvm::Module &llvmModule,
+                               llvm::TargetMachine &targetMachine) {
+  std::string targetISA;
+  llvm::raw_string_ostream stream(targetISA);
+
+  { // Drop pstream after this to prevent the ISA from being stuck buffering
+    llvm::buffer_ostream pstream(stream);
+    llvm::legacy::PassManager codegenPasses;
+
+    if (targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
+                                          llvm::CGFT_AssemblyFile))
+      return std::nullopt;
+
+    codegenPasses.run(llvmModule);
+  }
+  return stream.str();
+}
+
+std::optional<SmallVector<char, 0>>
+ModuleToObject::moduleToObject(llvm::Module &llvmModule,
+                               llvm::TargetMachine &targetMachine) {
+  SmallVector<char, 0> binaryData;
+  // Write the LLVM module bitcode to a buffer.
+  llvm::raw_svector_ostream outputStream(binaryData);
+  llvm::WriteBitcodeToFile(llvmModule, outputStream);
+  return binaryData;
+}
+
+std::optional<SmallVector<char, 0>> ModuleToObject::run() {
+  // Translate the module to LLVM IR.
+  llvm::LLVMContext llvmContext;
+  std::unique_ptr<llvm::Module> llvmModule = translateToLLVMIR(llvmContext);
+  if (!llvmModule) {
+    getOperation().emitError() << "Failed creating the llvm::Module.";
+    return std::nullopt;
+  }
+
+  // Create the target machine.
+  std::unique_ptr<llvm::TargetMachine> targetMachine = createTargetMachine();
+  if (!targetMachine)
+    return std::nullopt;
+
+  // Set the data layout and target triple of the module.
+  llvmModule->setDataLayout(targetMachine->createDataLayout());
+  llvmModule->setTargetTriple(targetMachine->getTargetTriple().getTriple());
+
+  // Link bitcode files.
+  handleModulePreLink(*llvmModule, *targetMachine);
+  {
+    auto libs = loadBitcodeFiles(*llvmModule, *targetMachine);
+    if (!libs)
+      return std::nullopt;
+    if (libs->size())
+      if (failed(linkFiles(*llvmModule, std::move(*libs))))
+        return std::nullopt;
+    handleModulePostLink(*llvmModule, *targetMachine);
+  }
+
+  // Optimize the module.
+  if (failed(optimizeModule(*llvmModule, *targetMachine, optLevel)))
+    return std::nullopt;
+
+  // Return the serialized object.
+  return moduleToObject(*llvmModule, *targetMachine);
+}

mlir/unittests/CMakeLists.txt

@@ -16,6 +16,7 @@ add_subdirectory(Pass)
 add_subdirectory(Support)
 add_subdirectory(Rewrite)
 add_subdirectory(TableGen)
+add_subdirectory(Target)
 add_subdirectory(Transforms)
 
 if(MLIR_ENABLE_EXECUTION_ENGINE)

mlir/unittests/Target/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(LLVM)

mlir/unittests/Target/LLVM/CMakeLists.txt

@@ -0,0 +1,26 @@
+add_mlir_unittest(MLIRTargetLLVMTests
+  SerializeToLLVMBitcode.cpp
+)
+
+llvm_map_components_to_libnames(llvm_libs nativecodegen)
+
+target_link_libraries(MLIRTargetLLVMTests
+  PRIVATE
+  MLIRTargetLLVM
+  MLIRLLVMDialect
+  MLIRLLVMToLLVMIRTranslation
+  MLIRBuiltinToLLVMIRTranslation
+  ${llvm_libs}
+)
+
+if (DEFINED LLVM_NATIVE_TARGET)
+  target_compile_definitions(MLIRTargetLLVMTests
+    PRIVATE
+    -DLLVM_NATIVE_TARGET_TEST_ENABLED=1
+  )
+else()
+  target_compile_definitions(MLIRTargetLLVMTests
+    PRIVATE
+    -DLLVM_NATIVE_TARGET_TEST_ENABLED=0
+  )
+endif()

mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp

@@ -0,0 +1,76 @@
+//===- SerializeToLLVMBitcode.cpp -------------------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Target/LLVM/ModuleToObject.h"
+#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/MemoryBufferRef.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/TargetParser/Host.h"
+
+#include "gmock/gmock.h"
+
+using namespace mlir;
+
+// Skip the test if the native target was not built.
+#if LLVM_NATIVE_TARGET_TEST_ENABLED == 0
+#define SKIP_WITHOUT_NATIVE(x) DISABLED_##x
+#else
+#define SKIP_WITHOUT_NATIVE(x) x
+#endif
+
+class MLIRTargetLLVM : public ::testing::Test {
+protected:
+  virtual void SetUp() {
+    llvm::InitializeNativeTarget();
+    llvm::InitializeNativeTargetAsmPrinter();
+  }
+};
+
+TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
+  std::string moduleStr = R"mlir(
+  llvm.func @foo(%arg0 : i32) {
+    llvm.return
+  }
+  )mlir";
+
+  DialectRegistry registry;
+  registerBuiltinDialectTranslation(registry);
+  registerLLVMDialectTranslation(registry);
+  MLIRContext context(registry);
+
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+
+  // Serialize the module.
+  std::string targetTriple = llvm::sys::getDefaultTargetTriple();
+  LLVM::ModuleToObject serializer(*(module->getOperation()), targetTriple, "",
+                                  "");
+  std::optional<SmallVector<char, 0>> serializedModule = serializer.run();
+  ASSERT_TRUE(!!serializedModule);
+  ASSERT_TRUE(serializedModule->size() > 0);
+
+  // Read the serialized module.
+  llvm::MemoryBufferRef buffer(
+      StringRef(serializedModule->data(), serializedModule->size()), "module");
+  llvm::LLVMContext llvmContext;
+  llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =
+      llvm::getLazyBitcodeModule(buffer, llvmContext);
+  ASSERT_TRUE(!!llvmModule);
+  ASSERT_TRUE(!!*llvmModule);
+
+  // Check that it has a function named `foo`.
+  ASSERT_TRUE((*llvmModule)->getFunction("foo") != nullptr);
+}