Merge pull request #30 from facebookresearch/pcpu

Permit parallel CPU backend compilation (though currently disabled)

Author: prigoyal

include/tc/core/polyhedral/llvm_jit.h

@@ -34,11 +34,6 @@ class Jit {
   llvm::orc::RTDyldObjectLinkingLayer objectLayer_;
   llvm::orc::IRCompileLayer<decltype(objectLayer_), llvm::orc::SimpleCompiler>
       compileLayer_;
-  using OptimizeFunction = std::function<std::shared_ptr<llvm::Module>(
-      std::shared_ptr<llvm::Module>)>;
-
-  llvm::orc::IRTransformLayer<decltype(compileLayer_), OptimizeFunction>
-      optimizeLayer_;
 
  public:
   Jit();
@@ -47,17 +42,14 @@ class Jit {
       const std::string& specializedName,
       const polyhedral::Scop& scop);
 
-  using ModuleHandle = decltype(optimizeLayer_)::ModuleHandleT;
+  using ModuleHandle = decltype(compileLayer_)::ModuleHandleT;
   ModuleHandle addModule(std::unique_ptr<llvm::Module> M);
   void removeModule(ModuleHandle H);
 
   llvm::JITSymbol findSymbol(const std::string name);
   llvm::JITTargetAddress getSymbolAddress(const std::string name);
 
   llvm::TargetMachine& getTargetMachine();
-
- private:
-  std::shared_ptr<llvm::Module> optimizeModule(std::shared_ptr<llvm::Module> M);
 };
 
 } // namespace tc

src/core/polyhedral/codegen_llvm.cc

@@ -20,11 +20,17 @@
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Config/llvm-config.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/IPO.h"
+#include "llvm/Transforms/IPO/PassManagerBuilder.h"
+#include "llvm/Transforms/Tapir/CilkABI.h"
 
 #include "Halide/Halide.h"
 
@@ -37,6 +43,10 @@
 #include "tc/core/polyhedral/scop.h"
 #include "tc/core/scope_guard.h"
 
+#ifndef LLVM_VERSION_MAJOR
+#error LLVM_VERSION_MAJOR not set
+#endif
+
 using namespace Halide;
 
 namespace tc {
@@ -195,14 +205,17 @@ class IslAstExprInterpeter {
   }
 };
 
+DEFINE_bool(llvm_dump_before_opt, false, "Print IR before optimization");
+DEFINE_bool(llvm_dump_after_opt, false, "Print IR after optimization");
+static constexpr int kOptLevel = 3;
+
 class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
  public:
   const isl::pw_multi_aff* iteratorMap_;
   CodeGen_TC(Target t) : CodeGen_X86(t) {}
 
   using CodeGen_X86::codegen;
   using CodeGen_X86::llvm_type_of;
-  using CodeGen_X86::optimize_module;
   using CodeGen_X86::sym_get;
   using CodeGen_X86::sym_pop;
   using CodeGen_X86::sym_push;
@@ -294,6 +307,52 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
 
     value = sym_get(name);
   }
+
+ public:
+  void optimize_module() {
+    if (FLAGS_llvm_dump_before_opt) {
+      module->print(llvm::dbgs(), nullptr, false, true);
+    }
+
+    llvm::legacy::FunctionPassManager functionPassManager(module.get());
+    llvm::legacy::PassManager modulePassManager;
+
+    std::unique_ptr<llvm::TargetMachine> targetMachine =
+        Halide::Internal::make_target_machine(*module);
+    modulePassManager.add(llvm::createTargetTransformInfoWrapperPass(
+        targetMachine ? targetMachine->getTargetIRAnalysis()
+                      : llvm::TargetIRAnalysis()));
+    functionPassManager.add(llvm::createTargetTransformInfoWrapperPass(
+        targetMachine ? targetMachine->getTargetIRAnalysis()
+                      : llvm::TargetIRAnalysis()));
+
+    llvm::PassManagerBuilder b;
+    b.OptLevel = kOptLevel;
+    b.tapirTarget = new llvm::CilkABI();
+    b.Inliner = llvm::createFunctionInliningPass(b.OptLevel, 0, false);
+    b.LoopVectorize = true;
+    b.SLPVectorize = true;
+
+    if (targetMachine) {
+      targetMachine->adjustPassManager(b);
+    }
+
+    b.populateFunctionPassManager(functionPassManager);
+    b.populateModulePassManager(modulePassManager);
+
+    // Run optimization passes
+    functionPassManager.doInitialization();
+    for (llvm::Module::iterator i = module->begin(); i != module->end(); i++) {
+      functionPassManager.run(*i);
+
+      functionPassManager.doFinalization();
+      modulePassManager.run(*module);
+
+      if (FLAGS_llvm_dump_after_opt) {
+        module->print(llvm::dbgs(), nullptr, false, true);
+      }
+    }
+  }
 };
 
 class LLVMCodegen {
@@ -451,6 +510,19 @@ class LLVMCodegen {
         llvm::BasicBlock::Create(llvmCtx, "loop_latch", function);
     auto* loopExitBB = llvm::BasicBlock::Create(llvmCtx, "loop_exit", function);
 
+    // TODO: integrate query ISL as to whether the relevant loop ought be
+    // parallelized
+    bool parallel = false;
+
+    llvm::Value* SyncRegion = nullptr;
+    if (parallel) {
+      SyncRegion = halide_cg.get_builder().CreateCall(
+          llvm::Intrinsic::getDeclaration(
+              function->getParent(), llvm::Intrinsic::syncregion_start),
+          {},
+          "syncreg");
+    }
+
     halide_cg.get_builder().CreateBr(headerBB);
 
     llvm::PHINode* phi = nullptr;
@@ -498,9 +570,22 @@ class LLVMCodegen {
     // Create Body
     {
       halide_cg.get_builder().SetInsertPoint(loopBodyBB);
+
+      if (parallel) {
+        auto* detachedBB =
+            llvm::BasicBlock::Create(llvmCtx, "det.achd", function);
+        halide_cg.get_builder().CreateDetach(
+            detachedBB, loopLatchBB, SyncRegion);
+        halide_cg.get_builder().SetInsertPoint(detachedBB);
+      }
       auto* currentBB = emitAst(node.for_get_body());
       halide_cg.get_builder().SetInsertPoint(currentBB);
-      halide_cg.get_builder().CreateBr(loopLatchBB);
+
+      if (parallel) {
+        halide_cg.get_builder().CreateReattach(loopLatchBB, SyncRegion);
+      } else {
+        halide_cg.get_builder().CreateBr(loopLatchBB);
+      }
     }
 
     // Create Latch
@@ -516,6 +601,11 @@ class LLVMCodegen {
 
     halide_cg.get_builder().SetInsertPoint(loopExitBB);
     halide_cg.sym_pop(node.for_get_iterator().get_id().get_name());
+    if (parallel) {
+      auto* syncBB = llvm::BasicBlock::Create(llvmCtx, "synced", function);
+      halide_cg.get_builder().CreateSync(syncBB, SyncRegion);
+      halide_cg.get_builder().SetInsertPoint(syncBB);
+    }
     return halide_cg.get_builder().GetInsertBlock();
   }
 
@@ -583,7 +673,8 @@ class LLVMCodegen {
   CodeGen_TC halide_cg;
 };
 
-// Create a list of isl ids to be used as loop iterators when building the AST.
+// Create a list of isl ids to be used as loop iterators when building the
+// AST.
 //
 // Note that this function can be scrapped as ISL can generate some default
 // iterator names.  However, it may come handy for associating extra info with

src/core/polyhedral/llvm_jit.cc

@@ -13,6 +13,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include <stdexcept>
+
 #include "tc/core/polyhedral/llvm_jit.h"
 
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
@@ -29,17 +31,55 @@
 
 using namespace llvm;
 
+// Parse through ldconfig to find the path of a particular
+// shared library. This is an unfortunate way to have to
+// find it, but I couldn't immediately find something in
+// imported libraries that would resolve this for us.
+std::string find_library_path(std::string library) {
+  std::string command = "ldconfig -p | grep " + library;
+
+  FILE* fpipe = popen(command.c_str(), "r");
+
+  if (fpipe == nullptr) {
+    throw std::runtime_error("Failed to popen()");
+  }
+
+  std::string output;
+  char buffer[512];
+
+  while (1) {
+    int charactersRead = fread(buffer, 1, sizeof(buffer), fpipe);
+    if (charactersRead == 0)
+      break;
+    output += std::string(buffer, charactersRead);
+  }
+  pclose(fpipe);
+
+  int idx = output.rfind("=> ");
+  if (idx == std::string::npos) {
+    throw std::runtime_error("Failed locate library: " + library);
+  }
+  output = output.substr(idx + 3);
+  if (output.length() > 0 && output[output.length() - 1] == '\n') {
+    output = output.substr(0, output.length() - 1);
+  }
+  return output;
+}
+
 namespace tc {
 
 Jit::Jit()
     : TM_(EngineBuilder().selectTarget()),
       DL_(TM_->createDataLayout()),
       objectLayer_([]() { return std::make_shared<SectionMemoryManager>(); }),
-      compileLayer_(objectLayer_, orc::SimpleCompiler(*TM_)),
-      optimizeLayer_(compileLayer_, [this](std::shared_ptr<Module> M) {
-        return optimizeModule(std::move(M));
-      }) {
-  sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
+      compileLayer_(objectLayer_, orc::SimpleCompiler(*TM_)) {
+  std::string err;
+
+  auto path = find_library_path("libcilkrts.so");
+  sys::DynamicLibrary::LoadLibraryPermanently(path.c_str(), &err);
+  if (err != "") {
+    throw std::runtime_error("Failed to find cilkrts: " + err);
+  }
 }
 
 void Jit::codegenScop(
@@ -57,7 +97,7 @@ Jit::ModuleHandle Jit::addModule(std::unique_ptr<Module> M) {
   M->setTargetTriple(TM_->getTargetTriple().str());
   auto Resolver = orc::createLambdaResolver(
       [&](const std::string& Name) {
-        if (auto Sym = optimizeLayer_.findSymbol(Name, false))
+        if (auto Sym = compileLayer_.findSymbol(Name, false))
           return Sym;
         return JITSymbol(nullptr);
       },
@@ -67,7 +107,7 @@ Jit::ModuleHandle Jit::addModule(std::unique_ptr<Module> M) {
         return JITSymbol(nullptr);
       });
 
-  auto res = optimizeLayer_.addModule(std::move(M), std::move(Resolver));
+  auto res = compileLayer_.addModule(std::move(M), std::move(Resolver));
   CHECK(res) << "Failed to jit compile.";
   return *res;
 }
@@ -76,7 +116,7 @@ JITSymbol Jit::findSymbol(const std::string Name) {
   std::string MangledName;
   raw_string_ostream MangledNameStream(MangledName);
   Mangler::getNameWithPrefix(MangledNameStream, Name, DL_);
-  return optimizeLayer_.findSymbol(MangledNameStream.str(), true);
+  return compileLayer_.findSymbol(MangledNameStream.str(), true);
 }
 
 JITTargetAddress Jit::getSymbolAddress(const std::string Name) {
@@ -85,44 +125,4 @@ JITTargetAddress Jit::getSymbolAddress(const std::string Name) {
   return *res;
 }
 
-DEFINE_bool(llvm_no_opt, false, "Disable LLVM optimizations");
-DEFINE_bool(llvm_debug_passes, false, "Print pass debug info");
-DEFINE_bool(llvm_dump_optimized_ir, false, "Print optimized IR");
-
-std::shared_ptr<Module> Jit::optimizeModule(std::shared_ptr<Module> M) {
-  if (FLAGS_llvm_no_opt) {
-    return M;
-  }
-
-  PassBuilder PB(TM_.get());
-  AAManager AA;
-  CHECK(PB.parseAAPipeline(AA, "default"))
-      << "Unable to parse AA pipeline description.";
-  LoopAnalysisManager LAM(FLAGS_llvm_debug_passes);
-  FunctionAnalysisManager FAM(FLAGS_llvm_debug_passes);
-  CGSCCAnalysisManager CGAM(FLAGS_llvm_debug_passes);
-  ModuleAnalysisManager MAM(FLAGS_llvm_debug_passes);
-  FAM.registerPass([&] { return std::move(AA); });
-  PB.registerModuleAnalyses(MAM);
-  PB.registerCGSCCAnalyses(CGAM);
-  PB.registerFunctionAnalyses(FAM);
-  PB.registerLoopAnalyses(LAM);
-  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
-
-  ModulePassManager MPM(FLAGS_llvm_debug_passes);
-  MPM.addPass(VerifierPass());
-  CHECK(PB.parsePassPipeline(MPM, "default<O3>", true, FLAGS_llvm_debug_passes))
-      << "Unable to parse pass pipline description.";
-  MPM.addPass(VerifierPass());
-
-  MPM.run(*M, MAM);
-
-  if (FLAGS_llvm_dump_optimized_ir) {
-    // M->dump(); // does not link
-    M->print(llvm::errs(), nullptr);
-  }
-
-  return M;
-}
-
 } // namespace tc