1#ifndef _TRITON_CODE_GEN_EXTERN_LIB_H_
2#define _TRITON_CODE_GEN_EXTERN_LIB_H_
3
4#include <memory>
5#include <string>
6#include <map>
7
8#include "llvm/IR/LLVMContext.h"
9#include "llvm/IR/Module.h"
10#include "llvm/IRReader/IRReader.h"
11#include "llvm/Support/SourceMgr.h"
12
13namespace triton {
14namespace codegen {
15
16///
17/// \brief ExternLib is a class that represents a library of external functions.
18///
19class ExternLib {
20 public:
21 ExternLib(const std::string &name, const std::string &path)
22 : name_(name), path_(path) {}
23
24 virtual ~ExternLib() = default;
25
26 virtual const std::string &name() const { return name_; }
27
28 virtual const std::string &path() const { return path_; }
29
30 ///
31 /// \brief Load the library and return the module.
32 ///
33 std::unique_ptr<llvm::Module> load(llvm::LLVMContext &ctx);
34
35 ///
36 /// \brief Link the module into the given module.
37 ///
38 void link(std::unique_ptr<llvm::Module> &llvm,
39 std::unique_ptr<llvm::Module> &mod);
40
41 ///
42 /// \brief Run load, link, and opt on the module.
43 ///
44 virtual void install(llvm::LLVMContext &ctx,
45 std::unique_ptr<llvm::Module> &llvm) {
46 auto mod = load(ctx);
47 link(llvm, mod);
48 opt(ctx, llvm);
49 }
50
51 ///
52 /// \brief Run opt on the module.
53 ///
54 virtual void opt(llvm::LLVMContext &ctx,
55 std::unique_ptr<llvm::Module> &llvm) = 0;
56
57 private:
58 std::string name_;
59 std::string path_;
60};
61
62///
63/// \brief ExternLibMap is a map of ExternLibs from their names to their paths.
64///
65typedef std::map<std::string, std::unique_ptr<ExternLib>> ExternLibMap;
66
67///
68/// \brief Concrete class for NVIDIA's libdevice library.
69///
70class LibDevice final : public ExternLib {
71 public:
72 LibDevice(const std::string &name, const std::string &path)
73 : ExternLib(name, path) {}
74
75 virtual ~LibDevice() = default;
76
77 virtual void opt(llvm::LLVMContext &ctx,
78 std::unique_ptr<llvm::Module> &llvm) override;
79};
80
81///
82/// \brief Create an ExternLib instance based on the name and path.
83///
84std::unique_ptr<ExternLib> create_extern_lib(const std::string &lib_name,
85 const std::string &lib_path);
86
87} // namespace codegen
88} // namespace triton
89
90#endif
91