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

[fmorac]

Summary

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.

[fmorac]

Rebasing.

[mehdi_amini]

Does this really belong to the ExecutionEngine? Seems like a more independent utility, but I'm not entirely sure where to slot this either.

mlir/include/mlir/ExecutionEngine/ModuleToObject.h

37

Does this only work with gpu.module as the comment says?

[mehdi_amini]

Also this is dead code right now: could this be unit-tested somehow?

[fmorac]

In D154100#4464080, @mehdi_amini wrote:

Does this really belong to the ExecutionEngine? Seems like a more independent utility, but I'm not entirely sure where to slot this either.

I was not sure were to put it either. I add it here because I know the execution engine can dump object code, so maybe some of these utilities could be used by the Execution Engine in the future?

In D154100#4464082, @mehdi_amini wrote:

Also this is dead code right now: could this be unit-tested somehow?

The best test I could think of is serializing a module and reading it back and checking the contents, however I don't think that would be a particularly good test.

[krzysz00]

As a general thought re organization, it might make sense to move this and also stuff like OptUtils from the execution engine to ... let's call it something like LLVMUsageUtils or LLVMCompilerUtils ... and that'd either be a standalone thing or something we'd stick under Support/ (or Target/LLVM/) by analogy to how LLVM keeps a lot of the OS-interaction wrappers in its Support library..

But maybe that's a very silly thought.

[mehdi_amini]

Targets/LLVM may be a good fit (as a new CMake library target in this directory).

[fmorac]

Moves the ModuleToObject class to Target/LLVM. Also adds an unit test.

[krzysz00]

I'm confused about one of the dependencies, but this is a good refactor otherwise

mlir/lib/Target/LLVM/CMakeLists.txt

13 ↗

(On Diff #538780)

This shouldn't need the native codegen?

21 ↗

(On Diff #538780)

Nit: newline

mlir/unittests/Target/LLVM/CMakeLists.txt

4 ↗

(On Diff #538780)

The dependency on native codegen might need to be here?

[fmorac]

Added new line.
Moved nativecodegen to the test.

[mehdi_amini]

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

29 ↗

(On Diff #543106)

That assumes that the operation can be translated to LLVM? This isn't clearly said here.

What about a SPIRV target?

[fmorac]

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

29 ↗

(On Diff #543106)

I'll update the description.

SPIRV is out of scope for this class, that's also why we moved the class to Targets/LLVM.

[mehdi_amini]

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

29 ↗

(On Diff #543106)

Then shall it be in the mlir::LLVM namespace as well? The name is quite generic and could conflict.

[fmorac]

Moved ModuleToObject into the namespace mlir::LLVM.

[fmorac]

Updating commit message.

[fmorac]

Change the name of the test.

[mehdi_amini]

mlir/unittests/Target/LLVM/CMakeLists.txt

7 ↗

(On Diff #545487)

You marked @krzysz00 comment above "done", but we still have nativecodegen here, I guess it may be needed because you need to generate "some sort of object file" here?

What happens to this unit-test if I build LLVM with only the NVPTX target enabled? Or in general without a native target?

12 ↗

(On Diff #545487)

We should just link the dialects we need here.

[fmorac]

mlir/unittests/Target/LLVM/CMakeLists.txt

7 ↗

(On Diff #545487)

nativecodegen is needed because I'm creating a target machine and attaching it to the module, it fails if there's not a valid target.

What happens to this unit-test if I build LLVM with only the NVPTX target enabled? Or in general without a native target?

I would assume it fails and I didn't even think of that because I based the test on one from ExecutionEngine, I think the solution is skipping the test when the native target is not built.

12 ↗

(On Diff #545487)

I'll change it.

[krzysz00]

I'd like to see this working with just AMDGPU target enabled (and with X86 or the native target) off so I can use this in a library without needing to pull in extra targets

[fmorac]

In D154100#4545862, @krzysz00 wrote:

I'd like to see this working with just AMDGPU target enabled (and with X86 or the native target) off so I can use this in a library without needing to pull in extra targets

The only component that needs nativecodegen is the unit test, there are no dependencies in nativecodegen on the MLIRTargetLLVM library, the class compiles and works with any target. See the unit test in D154117 in there I work with NVPTX, I plan to add a similar test for AMDGPU.

[krzysz00]

Having gone back and re-read, yeah, the only thing here is turning the test off when native codegen isn't available

[fmorac]

Removed unnecessary libs & skip the test if the Native target is not being built.

[mehdi_amini]

mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp

38 ↗

(On Diff #546178)

Does this build without the native target? (I assumed you tried?)

Can we make it not a global and instead done in the constructor of a fixture?

[fmorac]

Switched to Fixtures for the test.

mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp

38 ↗

(On Diff #546178)

Yeah I tried with only NVPTX, those functions are always defined but have empty bodies if the target is not built, see https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Support/TargetSelect.h#L119

Done.

[fmorac]

Rebasing.

mlir/include/mlir/ExecutionEngine/ModuleToObject.h

+//===- 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_EXECUTIONENGINE_MODULETOOBJECT_H
+#define MLIR_EXECUTIONENGINE_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 bitcode for the module.
+class ModuleToObject {
+public:
+  ModuleToObject(Operation &module, StringRef triple, StringRef chip,
+                 StringRef features = {}, int optLevel = 3);
+  virtual ~ModuleToObject() = default;
+
+  /// Returns the gpu.module 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 gpu.module 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 mlir
+
+#endif // MLIR_EXECUTIONENGINE_MODULETOOBJECT_H

mlir/lib/ExecutionEngine/CMakeLists.txt

set(LLVM_OPTIONAL_SOURCES
   JitRunner.cpp
   )
 
 # Use a separate library for OptUtils, to avoid pulling in the entire JIT and
 # codegen infrastructure. Unlike MLIRExecutionEngine, this is part of
 # libMLIR.so.
 add_mlir_library(MLIRExecutionEngineUtils
   OptUtils.cpp
+  ModuleToObject.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/ExecutionEngine
 
   DEPENDS
   intrinsics_gen
 
   LINK_COMPONENTS

mlir/lib/ExecutionEngine/ModuleToObject.cpp

+//===- 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/ExecutionEngine/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/TargetParser/TargetParser.h"
+#include "llvm/Transforms/IPO/Internalize.h"
+
+using namespace mlir;
+
+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);
+}