[FRONTEND] Added support for element-wise function defined in external LLVM bitcode (e.g., libdevice) (#562)

lib/codegen/extern_lib.cc

@@ -0,0 +1,63 @@
+#include "triton/codegen/extern_lib.h"
+
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Linker/Linker.h"
+#include "llvm/Transforms/IPO/PassManagerBuilder.h"
+#include "triton/codegen/pass.h"
+
+namespace triton {
+
+namespace codegen {
+
+std::unique_ptr<llvm::Module> ExternLib::load(llvm::LLVMContext& ctx) {
+  llvm::SMDiagnostic err;
+  auto mod = llvm::parseIRFile(this->path_, err, ctx);
+  if (!mod) {
+    throw std::runtime_error("Failed to load extern lib " + this->name_ +
+                             " at " + this->path_);
+  }
+  return mod;
+}
+
+void ExternLib::link(std::unique_ptr<llvm::Module>& llvm,
+                     std::unique_ptr<llvm::Module>& mod) {
+  // Set triple and data layout to match the target module
+  mod->setTargetTriple(llvm->getTargetTriple());
+  mod->setDataLayout(llvm->getDataLayout());
+  if (llvm::Linker::linkModules(*llvm, std::move(mod))) {
+    throw std::runtime_error("Failed to link extern lib " + this->name_ +
+                             " at " + this->path_);
+  }
+}
+
+void LibDevice::opt(llvm::LLVMContext& ctx, std::unique_ptr<llvm::Module>& llvm) {
+  // Add nvvm reflect flags to llvm module
+  // https://llvm.org/docs/LangRef.html#module-flags-metadata
+  // i32 4: Override the other module.
+  // i32 1: Emit an error
+  // If both modules specify Override, but the values differ, an error
+  // will be emitted.
+  llvm::Type* I32 = llvm::Type::getInt32Ty(ctx);
+  llvm::Metadata* md_four =
+      llvm::ConstantAsMetadata::get(llvm::ConstantInt::getSigned(I32, 4));
+  llvm::Metadata* md_name = llvm::MDString::get(ctx, "nvvm-reflect-ftz");
+  llvm::Metadata* md_one =
+      llvm::ConstantAsMetadata::get(llvm::ConstantInt::getSigned(I32, 1));
+  llvm::MDNode* reflect = llvm::MDNode::get(ctx, {md_four, md_name, md_one});
+  llvm->addModuleFlag(reflect);
+}
+
+std::unique_ptr<ExternLib> create_extern_lib(const std::string& lib_name,
+                                             const std::string& lib_path) {
+  if (lib_name == "libdevice") {
+    return std::make_unique<LibDevice>(lib_name, lib_path);
+  } else {
+    throw std::runtime_error("Unknown external library: " + lib_name);
+  }
+}
+
+}  // namespace codegen
+}  // namespace triton