ExecutionEngine.h

Go to the documentation of this file.

//===- ExecutionEngine.h - MLIR Execution engine and utils -----*- C++ -*--===//
//
// Part of the MLIR 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 provides a JIT-backed execution engine for MLIR modules.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_
#define MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_

#include "mlir/Support/LLVM.h"
#include "llvm/ExecutionEngine/ObjectCache.h"
#include "llvm/ExecutionEngine/Orc/LLJIT.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/Error.h"

#include <functional>
#include <memory>

namespace llvm {
template <typename T> class Expected;
class Module;
class ExecutionEngine;
class MemoryBuffer;
} // namespace llvm

namespace mlir {

class ModuleOp;

class SimpleObjectCache : public llvm::ObjectCache {
public:
  void notifyObjectCompiled(const llvm::Module *M,
                            llvm::MemoryBufferRef ObjBuffer) override;
  std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *M) override;

  void dumpToObjectFile(StringRef filename);

private:
  llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> cachedObjects;
};

class ExecutionEngine {
public:
  ExecutionEngine(bool enableObjectCache);

  static llvm::Expected<std::unique_ptr<ExecutionEngine>> create(
      ModuleOp m, std::function<llvm::Error(llvm::Module *)> transformer = {},
      Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
      ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = false);

  llvm::Expected<void (*)(void **)> lookup(StringRef name) const;

  template <typename... Args>
  llvm::Error invoke(StringRef name, Args &... args);

  llvm::Error invoke(StringRef name, MutableArrayRef<void *> args);

  static bool setupTargetTriple(llvm::Module *llvmModule);

  void dumpToObjectFile(StringRef filename);

private:
  // Ordering of llvmContext and jit is important for destruction purposes: the
  // jit must be destroyed before the context.
  llvm::LLVMContext llvmContext;

  // Underlying LLJIT.
  std::unique_ptr<llvm::orc::LLJIT> jit;

  // Underlying cache.
  std::unique_ptr<SimpleObjectCache> cache;
};

template <typename... Args>
llvm::Error ExecutionEngine::invoke(StringRef name, Args &... args) {
  auto expectedFPtr = lookup(name);
  if (!expectedFPtr)
    return expectedFPtr.takeError();
  auto fptr = *expectedFPtr;

  SmallVector<void *, 8> packedArgs{static_cast<void *>(&args)...};
  (*fptr)(packedArgs.data());

  return llvm::Error::success();
}

} // end namespace mlir

#endif // MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_