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 | |
13 | namespace triton { |
14 | namespace codegen { |
15 | |
16 | /// |
17 | /// \brief ExternLib is a class that represents a library of external functions. |
18 | /// |
19 | class 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 | /// |
65 | typedef std::map<std::string, std::unique_ptr<ExternLib>> ExternLibMap; |
66 | |
67 | /// |
68 | /// \brief Concrete class for NVIDIA's libdevice library. |
69 | /// |
70 | class 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 | /// |
84 | std::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 | |