1 #pragma once
2
3 #ifdef TORCH_ENABLE_LLVM
4 #include <c10/macros/Macros.h>
5 #include <c10/util/Exception.h>
6 #include <torch/csrc/Export.h>
7 #include <optional>
8
9 C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wsuggest-override")
10 #include <llvm/ExecutionEngine/JITSymbol.h>
C10_DIAGNOSTIC_POP()11 C10_DIAGNOSTIC_POP()
12 #include <llvm/ExecutionEngine/Orc/Core.h>
13 #include <llvm/ExecutionEngine/Orc/ThreadSafeModule.h>
14 #include <llvm/Target/TargetMachine.h>
15
16 #include <memory>
17 #include <string>
18
19 namespace torch {
20 namespace jit {
21 namespace tensorexpr {
22
23 inline std::string formatError(llvm::Error&& err, const char* msg) {
24 static constexpr const char* defaultErrorMsg =
25 "Unexpected failure in LLVM JIT";
26 std::string errorMsg(msg ? msg : defaultErrorMsg);
27 llvm::raw_string_ostream ss(errorMsg);
28 ss << ": " << err;
29 return ss.str();
30 }
31
32 template <typename T>
33 T assertSuccess(llvm::Expected<T> valOrErr, const char* msg = nullptr) {
34 TORCH_INTERNAL_ASSERT(valOrErr, formatError(valOrErr.takeError(), msg));
35 return std::move(*valOrErr);
36 }
37
38 inline void assertSuccess(llvm::Error err, const char* msg = nullptr) {
39 TORCH_INTERNAL_ASSERT(!err, formatError(std::move(err), msg));
40 }
41
42 } // namespace tensorexpr
43 } // namespace jit
44 } // namespace torch
45
46 namespace llvm {
47 namespace orc {
48
49 class PytorchLLVMJITImpl;
50
51 class TORCH_API PytorchLLVMJIT {
52 public:
53 PytorchLLVMJIT(
54 std::optional<std::string> triple,
55 std::optional<std::string> cpu,
56 std::optional<std::string> attrs);
57 ~PytorchLLVMJIT();
58
59 void addModule(std::unique_ptr<Module> M, std::unique_ptr<LLVMContext> C);
60
61 JITSymbol findSymbol(const std::string Name);
62
63 bool hasSymbol(const std::string& Name);
64
65 TargetMachine& getTargetMachine();
66
67 const DataLayout& getDataLayout();
68
69 private:
70 // Use the PImpl idiom here to hide the no-rtti parts of the JIT structure.
71 std::unique_ptr<PytorchLLVMJITImpl> impl_;
72 };
73
74 } // end namespace orc
75 } // end namespace llvm
76
77 #endif // ENABLE LLVM
78