Merge pull request #219 from facebookresearch/pcpumap

Parallel CPU mapper

Author: nicolasvasilache

build.sh

@@ -17,6 +17,7 @@ else
 fi
 WITH_PYTHON_C2=${WITH_PYTHON_C2:=OFF}
 WITH_NNPACK=${WITH_NNPACK:=OFF}
+WITH_TAPIR=${WITH_TAPIR:=ON}
 PYTHON=${PYTHON:="`which python3`"}
 PROTOC=${PROTOC:="`which protoc`"}
 CORES=${CORES:=32}
@@ -401,6 +402,7 @@ function install_tc() {
     rm -rf *
     VERBOSE=${VERBOSE} ${CMAKE_VERSION} -DWITH_CAFFE2=${WITH_CAFFE2} \
         -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
+        -DWITH_TAPIR=${WITH_TAPIR} \
         -DPYTHON_EXECUTABLE=${PYTHON} \
         -DHALIDE_PREFIX=${INSTALL_PREFIX} \
         -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \

include/tc/core/polyhedral/codegen_llvm.h

@@ -20,11 +20,29 @@
 
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 
 #include "Halide.h"
 
 namespace tc {
+
+static inline std::string toString(llvm::Value* llvmObject) {
+  std::string output;
+  llvm::raw_string_ostream rso(output);
+  llvmObject->print(rso);
+  rso.str();
+  return output;
+}
+
+static inline std::string toString(llvm::Module* llvmObject) {
+  std::string output;
+  llvm::raw_string_ostream rso(output);
+  llvmObject->print(rso, nullptr, false, true);
+  rso.str();
+  return output;
+}
+
 namespace polyhedral {
 struct Scop;
 

include/tc/core/polyhedral/llvm_jit.h

@@ -38,12 +38,11 @@ class Jit {
  public:
   Jit();
 
-  void codegenScop(
+  using ModuleHandle = decltype(compileLayer_)::ModuleHandleT;
+  std::shared_ptr<llvm::Module> codegenScop(
       const std::string& specializedName,
       const polyhedral::Scop& scop);
-
-  using ModuleHandle = decltype(compileLayer_)::ModuleHandleT;
-  ModuleHandle addModule(std::unique_ptr<llvm::Module> M);
+  ModuleHandle addModule(std::shared_ptr<llvm::Module> M);
   void removeModule(ModuleHandle H);
 
   llvm::JITSymbol findSymbol(const std::string name);

include/tc/core/polyhedral/scop.h

@@ -341,7 +341,6 @@ struct Scop {
   static std::unique_ptr<Scop> makeScheduled(
       const Scop& scop,
       const SchedulerOptionsView& schedulerOptions);
-
   // Tile the outermost band.
   // Splits the band into tile loop band and point loop band where point loops
   // have fixed trip counts specified in "tiling", and returns a pointer to the

src/core/polyhedral/codegen_llvm.cc

@@ -55,17 +55,6 @@ using namespace Halide;
 
 namespace tc {
 
-namespace {
-template <typename T>
-std::string toString(T* llvmObject) {
-  std::string output;
-  llvm::raw_string_ostream rso(output);
-  llvmObject->print(rso, nullptr, false, true);
-  rso.str();
-  return output;
-}
-} // namespace
-
 namespace halide2isl {
 isl::aff makeIslAffFromExpr(isl::space space, const Halide::Expr& e);
 }
@@ -217,6 +206,9 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
   using CodeGen_X86::sym_push;
 
   void init_module() override {
+    const char* llvm_args[] = {"tc (LLVM argument parsing)", nullptr};
+    llvm::cl::ParseCommandLineOptions(
+        sizeof(llvm_args) / sizeof(*llvm_args) - 1, llvm_args);
     init_context();
     module =
         llvm::make_unique<llvm::Module>("TensorComprehensionsModule", *context);
@@ -311,14 +303,13 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
     functionPassManager.doInitialization();
     for (llvm::Module::iterator i = module->begin(); i != module->end(); i++) {
       functionPassManager.run(*i);
+    }
 
-      functionPassManager.doFinalization();
-      modulePassManager.run(*module);
+    functionPassManager.doFinalization();
+    modulePassManager.run(*module);
 
-      LOG_IF(INFO, FLAGS_llvm_dump_after_opt)
-          << "[LLVM-IR] After optimization:\n"
-          << toString(module.get());
-    }
+    LOG_IF(INFO, FLAGS_llvm_dump_after_opt) << "[LLVM-IR] After optimization:\n"
+                                            << toString(module.get());
   }
 };
 
@@ -492,8 +483,7 @@ class LLVMCodegen {
 
     // TODO: integrate query ISL as to whether the relevant loop ought be
     // parallelized
-    bool parallel = false;
-
+    bool parallel = isl_ast_node_for_is_coincident(node.get());
     llvm::Value* SyncRegion = nullptr;
 
 #ifdef TAPIR_VERSION_MAJOR

src/core/polyhedral/llvm_jit.cc

@@ -82,18 +82,20 @@ Jit::Jit()
   }
 }
 
-void Jit::codegenScop(
+std::shared_ptr<Module> Jit::codegenScop(
     const std::string& specializedName,
     const polyhedral::Scop& scop) {
-  addModule(emitLLVMKernel(
-      specializedName, scop, getTargetMachine().createDataLayout()));
+  std::shared_ptr<Module> mod = emitLLVMKernel(
+      specializedName, scop, getTargetMachine().createDataLayout());
+  addModule(mod);
+  return mod;
 }
 
 TargetMachine& Jit::getTargetMachine() {
   return *TM_;
 }
 
-Jit::ModuleHandle Jit::addModule(std::unique_ptr<Module> M) {
+Jit::ModuleHandle Jit::addModule(std::shared_ptr<Module> M) {
   M->setTargetTriple(TM_->getTargetTriple().str());
   auto Resolver = orc::createLambdaResolver(
       [&](const std::string& Name) {
@@ -107,7 +109,7 @@ Jit::ModuleHandle Jit::addModule(std::unique_ptr<Module> M) {
         return JITSymbol(nullptr);
       });
 
-  auto res = compileLayer_.addModule(std::move(M), std::move(Resolver));
+  auto res = compileLayer_.addModule(M, std::move(Resolver));
   CHECK(res) << "Failed to jit compile.";
   return *res;
 }

test/CMakeLists.txt

@@ -100,7 +100,6 @@ foreach(i ${CORE_TEST_FILES})
   target_link_libraries(${i} ${GOOGLE_LIBS} tc_core_cuda_no_sdk)
 endforeach()
 
-
 add_executable(test_mapper_llvm test_mapper_llvm.cc)
 add_test(test_mapper_llvm test_mapper_llvm)
 target_link_libraries(
@@ -112,6 +111,19 @@ target_link_libraries(
 
   tc_core_cpu tc_lang)
 
+if (WITH_TAPIR)
+  add_executable(test_mapper_tapir test_mapper_tapir.cc)
+  add_test(test_mapper_tapir test_mapper_tapir)
+  target_link_libraries(
+    test_mapper_tapir
+
+    ${GOOGLE_LIBS}
+    ${ATEN_LIBRARIES}
+    -lLLVM
+
+    tc_core_cpu tc_lang)
+endif()
+
 ################################################################################
 # TensorComprehensions tests
 # No real need for NVCC if we only use NVRTC

test/test_mapper_tapir.cc

@@ -0,0 +1,98 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+
+#include <llvm/IR/InstIterator.h>
+#include <llvm/IR/Instructions.h>
+
+#include "tc/aten/utils.h"
+#include "tc/core/cpu/cpu_tc_executor.h"
+#include "tc/core/execution_engine.h"
+#include "tc/core/mapping_options.h"
+#include "tc/core/polyhedral/codegen_llvm.h"
+#include "tc/core/polyhedral/llvm_jit.h"
+#include "tc/core/polyhedral/scop.h"
+#include "tc/core/scope_guard.h"
+
+#include "test_harness_aten.h"
+
+using namespace std;
+
+using namespace tc;
+using namespace tc::polyhedral;
+using namespace tc::polyhedral::detail;
+
+TEST(TapirCodegen, BasicParallel) {
+  string tc = R"TC(
+def fun(float(N, M) A, float(N, M) B) -> (C) {
+  C(n, m) = A(n, m) + B(n, m)
+}
+)TC";
+  auto N = 40;
+  auto M = 24;
+
+  auto ctx = isl::with_exceptions::globalIslCtx();
+  auto scop = polyhedral::Scop::makeScop(ctx, tc);
+  auto context = scop->makeContext(
+      std::unordered_map<std::string, int>{{"N", N}, {"M", M}});
+  scop = Scop::makeSpecializedScop(*scop, context);
+  SchedulerOptionsProto sop;
+  SchedulerOptionsView sov(sop);
+  scop = Scop::makeScheduled(*scop, sov);
+  Jit jit;
+  auto mod = jit.codegenScop("kernel_anon", *scop);
+  auto fn = mod->getFunction("kernel_anon");
+
+  std::set<string> calledFunctions;
+  for (llvm::inst_iterator I = llvm::inst_begin(fn), E = llvm::inst_end(fn);
+       I != E;
+       ++I) {
+    if (llvm::CallInst* c = llvm::dyn_cast<llvm::CallInst>(&*I)) {
+      if (auto called = c->getCalledFunction()) {
+        calledFunctions.insert(called->getName());
+      }
+    }
+  }
+
+  ASSERT_NE(0, calledFunctions.count("__cilkrts_get_tls_worker"));
+  ASSERT_NE(0, calledFunctions.count("__cilkrts_bind_thread_1"));
+  ASSERT_NE(0, calledFunctions.count("llvm.stacksave"));
+  ASSERT_NE(0, calledFunctions.count("__cilkrts_sync"));
+
+  auto fptr =
+      (void (*)(float*, float*, float*))jit.getSymbolAddress("kernel_anon");
+
+  at::Tensor A = at::CPU(at::kFloat).rand({N, M});
+  at::Tensor B = at::CPU(at::kFloat).rand({N, M});
+  at::Tensor C = at::CPU(at::kFloat).rand({N, M});
+  at::Tensor Cc = A + B;
+  fptr(A.data<float>(), B.data<float>(), C.data<float>());
+
+  checkRtol(Cc - C, {A, B}, N * M);
+}
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  ::gflags::ParseCommandLineFlags(&argc, &argv, true);
+  ::google::InitGoogleLogging(argv[0]);
+  initialize_llvm();
+  return RUN_ALL_TESTS();
+}