[LLVM] Enable same target option in JITModule (#778)

* [LLVM] Enable same target option in JITModule

* not set mcpu explicitly

src/codegen/llvm/llvm_common.cc

@@ -36,9 +36,11 @@ void InitializeLLVM() {
   }
 }
 
-llvm::TargetMachine*
-GetLLVMTargetMachine(const std::string& target_str,
-                     bool allow_null) {
+void ParseLLVMTargetOptions(const std::string& target_str,
+                            std::string* triple,
+                            std::string* mcpu,
+                            std::string* mattr,
+                            llvm::TargetOptions* options) {
   // setup target triple
   size_t start = 0;
   if (target_str.length() >= 4 &&
@@ -46,9 +48,10 @@ GetLLVMTargetMachine(const std::string& target_str,
     start = 4;
   }
   // simple parser
-  std::string target_triple = "";
-  std::string cpu = "generic";
-  std::string attr = "";
+  triple->resize(0);
+  mcpu->resize(0);
+  mattr->resize(0);
+
   bool soft_float_abi = false;
   std::string key, value;
   std::istringstream is(target_str.substr(start, target_str.length() - start));
@@ -69,11 +72,11 @@ GetLLVMTargetMachine(const std::string& target_str,
     }
     if (key == "-target" ||
         key == "-mtriple") {
-      target_triple = value;
+      *triple = value;
     } else if (key == "-mcpu") {
-      cpu = value;
+      *mcpu = value;
     } else if (key == "-mattr") {
-      attr = value;
+      *mattr = value;
     } else if (key == "-mfloat-abi") {
       if (value == "hard") {
         soft_float_abi = false;
@@ -89,19 +92,13 @@ GetLLVMTargetMachine(const std::string& target_str,
     }
   }
 
-  if (target_triple.length() == 0 ||
-      target_triple == "default") {
-    target_triple = llvm::sys::getDefaultTargetTriple();
-  }
-  std::string err;
-  const llvm::Target* target =
-      llvm::TargetRegistry::lookupTarget(target_triple, err);
-  if (target == nullptr) {
-    CHECK(allow_null) << err << " target_triple=" << target_triple;
-    return nullptr;
+  if (triple->length() == 0 ||
+      *triple == "default") {
+    *triple = llvm::sys::getDefaultTargetTriple();
   }
   // set target option
-  llvm::TargetOptions opt;
+  llvm::TargetOptions& opt = *options;
+  opt = llvm::TargetOptions();
   #if TVM_LLVM_VERSION < 50
   opt.LessPreciseFPMADOption = true;
   #endif
@@ -114,8 +111,38 @@ GetLLVMTargetMachine(const std::string& target_str,
   } else {
     opt.FloatABIType = llvm::FloatABI::Hard;
   }
+}
+
+
+llvm::TargetMachine*
+GetLLVMTargetMachine(const std::string& target_str,
+                     bool allow_null) {
+  std::string target_triple, mcpu, mattr;
+  llvm::TargetOptions opt;
+
+  ParseLLVMTargetOptions(target_str,
+                         &target_triple,
+                         &mcpu,
+                         &mattr,
+                         &opt);
+
+  if (target_triple.length() == 0 ||
+      target_triple == "default") {
+    target_triple = llvm::sys::getDefaultTargetTriple();
+  }
+  if (mcpu.length() == 0) {
+    mcpu = "generic";
+  }
+
+  std::string err;
+  const llvm::Target* target =
+      llvm::TargetRegistry::lookupTarget(target_triple, err);
+  if (target == nullptr) {
+    CHECK(allow_null) << err << " target_triple=" << target_triple;
+    return nullptr;
+  }
   llvm::TargetMachine* tm = target->createTargetMachine(
-      target_triple, cpu, attr, opt, llvm::Reloc::PIC_);
+      target_triple, mcpu, mattr, opt, llvm::Reloc::PIC_);
   return tm;
 }
 

src/codegen/llvm/llvm_common.h

@@ -58,6 +58,20 @@ namespace codegen {
 void InitializeLLVM();
 
 /*!
+ * \brief Parse target options
+ * \param target_str Target string, in format "llvm -target=xxx -mcpu=xxx"
+ * \param triple Target triple
+ * \param mcpu cpu info
+ * \param options the options
+ * \param mattr The attributes
+ */
+void ParseLLVMTargetOptions(const std::string& target_str,
+                            std::string* triple,
+                            std::string* mcpu,
+                            std::string* mattr,
+                            llvm::TargetOptions* options);
+
+/*!
  * \brief Get target machine from target_str string.
  * \param target_str Target string, in format "llvm -target=xxx -mcpu=xxx"
  * \param allow_null Whether allow null to be returned.

src/codegen/llvm/llvm_module.cc

@@ -120,6 +120,10 @@ class LLVMModuleNode final : public runtime::ModuleNode {
     }
     cg->AddMainFunction(funcs[0]->name);
     module_ = cg->Finish();
+    module_->addModuleFlag(
+        llvm::Module::Warning, "tvm_target",
+        llvm::MDString::get(*ctx_, target));
+    target_ = target;
     mptr_ = module_.get();
   }
 
@@ -133,11 +137,19 @@ class LLVMModuleNode final : public runtime::ModuleNode {
       LOG(FATAL) << "Fail to load ir file " << file_name << "\n"
                  << "line " << err.getLineNo() << ":" << msg;
     }
-    std::string target = module_->getTargetTriple();
-    mptr_ = module_.get();
+    std::string target_;
+    llvm::Metadata* mtarget = module_->getModuleFlag("tvm_target");
+    if (mtarget != nullptr) {
+      llvm::MDString* pstr = llvm::dyn_cast<llvm::MDString>(mtarget);
+      CHECK(pstr != nullptr);
+      target_ = pstr->getString();
+    } else {
       std::ostringstream os;
-    os << "llvm -target " << target;
-    tm_ = GetLLVMTargetMachine(os.str());
+      os << "llvm -target " << module_->getTargetTriple();
+      target_ = os.str();
+    }
+    mptr_ = module_.get();
+    tm_ = GetLLVMTargetMachine(target_);
   }
 
  private:
@@ -145,8 +157,19 @@ class LLVMModuleNode final : public runtime::ModuleNode {
     CHECK(ee_ == nullptr);
     std::lock_guard<std::mutex> lock(mutex_);
     llvm::EngineBuilder builder(std::move(module_));
+    std::string triple, mcpu, mattr;
+    llvm::TargetOptions opt;
+    ParseLLVMTargetOptions(target_, &triple, &mcpu, &mattr, &opt);
     builder.setEngineKind(llvm::EngineKind::JIT);
     builder.setOptLevel(llvm::CodeGenOpt::Aggressive);
+    if (mcpu.length() != 0) {
+      builder.setMCPU(mcpu);
+    }
+    if (mattr.length() != 0) {
+      std::vector<std::string> mattrs{mattr};
+      builder.setMAttrs(mattrs);
+    }
+    builder.setTargetOptions(opt);
     llvm::TargetMachine *tm = builder.selectTarget();
     llvm::TargetMachine *tm_sys = GetLLVMTargetMachine("llvm");
     if (tm_sys->getTargetTriple().getArch() != tm->getTargetTriple().getArch()) {

topi/tests/python/test_topi_bnn.py

@@ -38,9 +38,9 @@ def verify_binary_dense(batch, in_dim, out_dim):
     bnn_a = tvm.nd.array(np.zeros(get_const_tuple(bnn_A.shape), dtype=bnn_A.dtype), ctx)
     bnn_b = tvm.nd.array(np.zeros(get_const_tuple(bnn_B.shape), dtype=bnn_B.dtype), ctx)
     bnn_c = tvm.nd.array(np.zeros(get_const_tuple(bnn_C.shape), dtype=bnn_C.dtype), ctx)
-    f1 = tvm.build(s1, [A, bnn_A], 'llvm -mcpu=core-avx2')
-    f2 = tvm.build(s2, [B, bnn_B], 'llvm -mcpu=core-avx2')
-    f3 = tvm.build(s3, [bnn_A1, bnn_B1, bnn_C], 'llvm -mcpu=core-avx2')
+    f1 = tvm.build(s1, [A, bnn_A], 'llvm')
+    f2 = tvm.build(s2, [B, bnn_B], 'llvm')
+    f3 = tvm.build(s3, [bnn_A1, bnn_B1, bnn_C], 'llvm')
     f1(a, bnn_a)
     f2(b, bnn_b)
     f3(bnn_a, bnn_b, bnn_c)