From c4f4ff2966aba1a4db6ae661583742d5d5f5b48e Mon Sep 17 00:00:00 2001
From: Diego Caballero <diego.caballero@intel.com>
Date: Mon, 22 Jul 2019 13:08:00 -0700
Subject: [PATCH] Pass reference target machine to OrcJit compiler.

PR#49 (https://github.com/tensorflow/mlir/pull/49) adds TTI
initialization pass to pass managers for a given target machine.
However, this may cause an inconsistency if the target machine used by
the JIT compiler is not the same as the one used by TTI. This PR allows
passing a reference target machine to be used by JIT compiler. If no
reference target machine is provided, a default host target machine is
used, including all its sub-target features.
---
 .../mlir/ExecutionEngine/ExecutionEngine.h    |  6 +-
 lib/ExecutionEngine/ExecutionEngine.cpp       | 63 ++++++++++++++++---
 2 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/include/mlir/ExecutionEngine/ExecutionEngine.h b/include/mlir/ExecutionEngine/ExecutionEngine.h
index 69f6c2e72f33..2023d0b75d15 100644
--- a/include/mlir/ExecutionEngine/ExecutionEngine.h
+++ b/include/mlir/ExecutionEngine/ExecutionEngine.h
@@ -32,6 +32,7 @@
 namespace llvm {
 template <typename T> class Expected;
 class Module;
+class TargetMachine;
 } // namespace llvm
 
 namespace mlir {
@@ -61,10 +62,13 @@ class ExecutionEngine {
   /// can be used, e.g., for reporting or optimization.
   /// If `sharedLibPaths` are provided, the underlying JIT-compilation will open
   /// and link the shared libraries for symbol resolution.
+  /// If `refTM` is provided, the underlying JIT-compilation will use that
+  /// target machine as reference to build its target machine.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
   create(ModuleOp m,
          std::function<llvm::Error(llvm::Module *)> transformer = {},
-         ArrayRef<StringRef> sharedLibPaths = {});
+         ArrayRef<StringRef> sharedLibPaths = {},
+         llvm::TargetMachine *refTM = nullptr);
 
   /// Looks up a packed-argument function with the given name and returns a
   /// pointer to it.  Propagates errors in case of failure.
diff --git a/lib/ExecutionEngine/ExecutionEngine.cpp b/lib/ExecutionEngine/ExecutionEngine.cpp
index 0317a92c43fa..8cfd0b769f1a 100644
--- a/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -144,10 +144,14 @@ class OrcJIT {
     loadLibraries(sharedLibPaths);
   }
 
-  // Create a JIT engine for the current host.
+  // Create a JIT engine for the reference target machine `refTM` or the current
+  // host if `refTM` is not provided.
   static Expected<std::unique_ptr<OrcJIT>>
-  createDefault(IRTransformer transformer, ArrayRef<StringRef> sharedLibPaths) {
-    auto machineBuilder = llvm::orc::JITTargetMachineBuilder::detectHost();
+  create(IRTransformer transformer, ArrayRef<StringRef> sharedLibPaths,
+         llvm::TargetMachine *refTM) {
+
+    auto machineBuilder =
+        refTM ? createRefMachineBuilder(refTM) : createDefaultMachineBuilder();
     if (!machineBuilder)
       return machineBuilder.takeError();
 
@@ -173,6 +177,50 @@ class OrcJIT {
   }
 
 private:
+  // Create a JIT engine for the reference target machine `refTM`.
+  static Expected<llvm::orc::JITTargetMachineBuilder>
+  createRefMachineBuilder(llvm::TargetMachine *refTM) {
+    assert(refTM && "Expected reference target machine!");
+
+    Expected<llvm::orc::JITTargetMachineBuilder> machineBuilder =
+        llvm::orc::JITTargetMachineBuilder(refTM->getTargetTriple());
+    if (!machineBuilder)
+      return machineBuilder.takeError();
+
+    machineBuilder->setCPU(refTM->getTargetCPU());
+    machineBuilder->setRelocationModel(refTM->getRelocationModel());
+    machineBuilder->setCodeModel(refTM->getCodeModel());
+    machineBuilder->setCodeGenOptLevel(refTM->getOptLevel());
+
+    std::vector<std::string> features;
+    llvm::SubtargetFeatures::Split(features, refTM->getTargetFeatureString());
+    machineBuilder->addFeatures(features);
+
+    return machineBuilder;
+  }
+
+  // Create a JIT engine for the current host.
+  static Expected<llvm::orc::JITTargetMachineBuilder>
+  createDefaultMachineBuilder() {
+    auto machineBuilder = llvm::orc::JITTargetMachineBuilder::detectHost();
+    if (!machineBuilder)
+      return machineBuilder.takeError();
+
+    // Retrieve host CPU sub-target features.
+    llvm::SubtargetFeatures subtargetFeatures;
+    llvm::StringMap<bool> featureMap;
+    llvm::sys::getHostCPUFeatures(featureMap);
+    for (auto &feature : featureMap)
+      subtargetFeatures.AddFeature(feature.first(), feature.second);
+
+    // Relocation model, code model and codegen opt level are kept to default
+    // values.
+    machineBuilder->setCPU(llvm::sys::getHostCPUName());
+    machineBuilder->addFeatures(subtargetFeatures.getFeatures());
+
+    return machineBuilder;
+  }
+
   // Wrap the `irTransformer` into a function that can be called by the
   // IRTranformLayer.  If `irTransformer` is not set up, return the module as
   // is without errors.
@@ -321,12 +369,11 @@ void packFunctionArguments(llvm::Module *module) {
 // Out of line for PIMPL unique_ptr.
 ExecutionEngine::~ExecutionEngine() = default;
 
-Expected<std::unique_ptr<ExecutionEngine>>
-ExecutionEngine::create(ModuleOp m,
-                        std::function<llvm::Error(llvm::Module *)> transformer,
-                        ArrayRef<StringRef> sharedLibPaths) {
+Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
+    ModuleOp m, std::function<llvm::Error(llvm::Module *)> transformer,
+    ArrayRef<StringRef> sharedLibPaths, llvm::TargetMachine *refTM) {
   auto engine = llvm::make_unique<ExecutionEngine>();
-  auto expectedJIT = impl::OrcJIT::createDefault(transformer, sharedLibPaths);
+  auto expectedJIT = impl::OrcJIT::create(transformer, sharedLibPaths, refTM);
   if (!expectedJIT)
     return expectedJIT.takeError();