[Triton-MLIR][BACKEND] Add elementwise ops and tests (#804)

Co-authored-by: Keren Zhou <kerenzhou@openai.com>

lib/Target/LLVMIR/LLVMIRTranslation.cpp

@@ -16,6 +16,9 @@
 #include "triton/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.h"
 #include "triton/tools/sys/getenv.hpp"
 #include "llvm/IR/Constants.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Linker/Linker.h"
+#include "llvm/Support/SourceMgr.h"
 
 namespace mlir {
 namespace triton {
@@ -148,13 +151,80 @@ translateTritonGPUToLLVMIR(llvm::LLVMContext *llvmContext,
     return nullptr;
   }
 
+  std::map<std::string, std::string> extern_libs;
+  SmallVector<LLVM::LLVMFuncOp> funcs;
+  module.walk([&](LLVM::LLVMFuncOp func) {
+    if (func.isExternal())
+      funcs.push_back(func);
+  });
+
+  for (auto &func : funcs) {
+    if (func.getOperation()->hasAttr("libname")) {
+      auto name =
+          func.getOperation()->getAttr("libname").dyn_cast<StringAttr>();
+      auto path =
+          func.getOperation()->getAttr("libpath").dyn_cast<StringAttr>();
+      if (name) {
+        std::string lib_name = name.str();
+        extern_libs[lib_name] = path.str();
+      }
+    }
+  }
+
+  if (module.getOperation()->hasAttr("triton_gpu.externs")) {
+    auto dict = module.getOperation()
+                    ->getAttr("triton_gpu.externs")
+                    .dyn_cast<DictionaryAttr>();
+    for (auto &attr : dict) {
+      extern_libs[attr.getName().strref().trim().str()] =
+          attr.getValue().dyn_cast<StringAttr>().strref().trim().str();
+    }
+  }
+
   auto llvmir = translateLLVMToLLVMIR(llvmContext, module);
   if (!llvmir) {
     llvm::errs() << "Translate to LLVM IR failed";
+    return nullptr;
+  }
+
+  llvm::SMDiagnostic err;
+  for (auto &lib : extern_libs) {
+    auto ext_mod = llvm::parseIRFile(lib.second, err, *llvmContext);
+    if (!ext_mod) {
+      llvm::errs() << "Failed to load extern lib " << lib.first;
+      return nullptr;
+    }
+    ext_mod->setTargetTriple(llvmir->getTargetTriple());
+    ext_mod->setDataLayout(llvmir->getDataLayout());
+
+    if (llvm::Linker::linkModules(*llvmir, std::move(ext_mod))) {
+      llvm::errs() << "Failed to link extern lib " << lib.first;
+      return nullptr;
+    }
   }
 
   return llvmir;
 }
 
+void addExternalLibs(mlir::ModuleOp &module,
+                     const std::vector<std::string> &names,
+                     const std::vector<std::string> &paths) {
+  if (names.empty() || names.size() != paths.size())
+    return;
+
+  llvm::SmallVector<NamedAttribute, 2> attrs;
+
+  for (size_t i = 0; i < names.size(); ++i) {
+    auto name = StringAttr::get(module->getContext(), names[i]);
+    auto path = StringAttr::get(module->getContext(), paths[i]);
+    NamedAttribute attr(name, path);
+    attrs.push_back(attr);
+  }
+
+  DictionaryAttr dict = DictionaryAttr::get(module->getContext(), attrs);
+  module.getOperation()->setAttr("triton_gpu.externs", dict);
+  return;
+}
+
 } // namespace triton
 } // namespace mlir