Proteus
Programmable JIT compilation and optimization for C/C++ using LLVM
Loading...
Searching...
No Matches
Func.hpp
Go to the documentation of this file.
1#ifndef PROTEUS_FRONTEND_FUNC_HPP
2#define PROTEUS_FRONTEND_FUNC_HPP
3
4#include <deque>
5#include <memory>
6
7#include <llvm/IR/IRBuilder.h>
8#include <llvm/IR/Module.h>
9
11#include "proteus/Error.h"
15
16namespace proteus {
17
18struct Var;
19class JitModule;
20class LoopBoundInfo;
21template <typename... ForLoopBuilders> class LoopNestBuilder;
22template <typename BodyLambda> class ForLoopBuilder;
23
24// Helper struct to represent the signature of a function.
25// Useful to partially-specialize function templates.
26template <typename... ArgTs> struct ArgTypeList {};
27template <typename T> struct FnSig;
28template <typename RetT_, typename... ArgT> struct FnSig<RetT_(ArgT...)> {
30 using RetT = RetT_;
31};
32
33using namespace llvm;
34
36 void operator()() const {}
37};
38
39class FuncBase {
40protected:
44 IRBuilderBase::InsertPoint IP;
45
46 std::deque<std::unique_ptr<Var>> Arguments;
47 std::deque<std::unique_ptr<Var>> Variables;
48 std::deque<std::unique_ptr<Var>> RuntimeConstants;
49
50 std::string Name;
51
52 enum class ScopeKind { FUNCTION, IF, FOR };
53 struct Scope {
54 std::string File;
55 int Line;
57 IRBuilderBase::InsertPoint ContIP;
58
59 explicit Scope(const char *File, int Line, ScopeKind Kind,
60 IRBuilderBase::InsertPoint ContIP)
62 };
63 std::vector<Scope> Scopes;
64
65 std::string toString(ScopeKind Kind) {
66 switch (Kind) {
68 return "FUNCTION";
69 case ScopeKind::IF:
70 return "IF";
71 case ScopeKind::FOR:
72 return "FOR";
73 default:
74 PROTEUS_FATAL_ERROR("Unsupported Kind " +
75 std::to_string(static_cast<int>(Kind)));
76 }
77 }
78
79 Var &emitAtomic(AtomicRMWInst::BinOp Op, Var &Addr, Var &Val);
80
81public:
83
85
88
90
92
93 // Create a variable. If PointerElemType is non-null, a PointerVar is created
94 // with an alloca-backed pointer slot; otherwise a ScalarVar is created.
96 Type *PointerElemType = nullptr);
97
98 template <typename T> Var &declVar(StringRef Name = "var") {
99 static_assert(!std::is_array_v<T>, "Expected non-array type");
100
101 Function *F = getFunction();
102 return declVarInternal(Name, TypeMap<T>::get(F->getContext()));
103 }
104
105 template <typename T>
107 StringRef Name = "array_var") {
108 static_assert(std::is_array_v<T>, "Expected array type");
109
110 Function *F = getFunction();
111 auto *BasePointer =
112 emitArrayCreate(TypeMap<T>::get(F->getContext(), NElem), AS, Name);
113
114 auto *ArrTy = cast<ArrayType>(TypeMap<T>::get(F->getContext(), NElem));
115 auto &Ref = *Variables.emplace_back(
116 std::make_unique<ArrayVar>(BasePointer, *this, ArrTy));
117 return Ref;
118 }
119
120 template <typename T> Var &defVar(T Val, StringRef Name = "var") {
121 Function *F = getFunction();
122 Var &VarRef = declVarInternal(Name, TypeMap<T>::get(F->getContext()));
123 VarRef = Val;
124 return VarRef;
125 }
126
127 Var &defVar(const Var &Val, StringRef Name = "var") {
128 if (&Val.Fn != this)
129 PROTEUS_FATAL_ERROR("Variables should belong to the same function");
130 Var &VarRef = declVarInternal(Name, Val.getValueType());
131 VarRef = Val;
132 return VarRef;
133 }
134
135 template <typename T>
136 Var &defRuntimeConst(T Val, StringRef Name = "run.const.var") {
137 Function *F = getFunction();
138 Var &VarRef = declVarInternal(Name, TypeMap<T>::get(F->getContext()));
139 VarRef = Val;
140 return VarRef;
141 }
142
143 template <typename... ArgT> auto defRuntimeConsts(ArgT &&...Args) {
144 return std::tie(defRuntimeConst(std::forward<ArgT>(Args))...);
145 }
146
147 template <typename... Ts> void declArgs() {
148 Function *F = getFunction();
149 auto &EntryBB = F->getEntryBlock();
150 IP = IRBuilderBase::InsertPoint(&EntryBB, EntryBB.end());
151 IRB.restoreIP(IP);
152
153 (
154 [&]() {
155 // Determine if this argument is a pointer-like type.
156 auto &Ctx = F->getContext();
157 Type *ArgTy = TypeMap<Ts>::get(Ctx);
159
160 // Allocate a slot to hold the incoming argument value (pointer or
161 // scalar).
162 auto *Alloca =
163 emitAlloca(ArgTy, "arg." + std::to_string(Arguments.size()));
164
165 auto *Arg = F->getArg(Arguments.size());
166 IRB.CreateStore(Arg, Alloca);
167
168 if (PtrElemTy) {
169 // Pointer argument: create a PointerVar with the correct element
170 // type.
171 Arguments.emplace_back(
172 std::make_unique<PointerVar>(Alloca, *this, PtrElemTy));
173 } else {
174 // Scalar argument.
175 Arguments.emplace_back(std::make_unique<ScalarVar>(Alloca, *this));
176 }
177 }(),
178 ...);
179 IRB.ClearInsertionPoint();
180 }
181
182 Var &getArg(unsigned int ArgNo);
183
184 void beginFunction(const char *File = __builtin_FILE(),
185 int Line = __builtin_LINE());
186 void endFunction();
187
188 void beginIf(Var &CondVar, const char *File = __builtin_FILE(),
189 int Line = __builtin_LINE());
190 void endIf();
191
192 void beginFor(Var &IterVar, Var &InitVar, Var &UpperBound, Var &IncVar,
193 const char *File = __builtin_FILE(),
194 int Line = __builtin_LINE());
195 void endFor();
196
197 template <typename Sig>
198 std::enable_if_t<!std::is_void_v<typename FnSig<Sig>::RetT>, Var &>
200
201 template <typename Sig>
202 std::enable_if_t<std::is_void_v<typename FnSig<Sig>::RetT>, void>
204
205 template <typename Sig, typename... ArgVars>
206 std::enable_if_t<!std::is_void_v<typename FnSig<Sig>::RetT>, Var &>
208
209 template <typename Sig, typename... ArgVars>
210 std::enable_if_t<std::is_void_v<typename FnSig<Sig>::RetT>, void>
212
213 template <typename BuiltinFuncT>
215 using RetT = std::invoke_result_t<BuiltinFuncT &, FuncBase &>;
216 if constexpr (std::is_void_v<RetT>) {
217 std::invoke(std::forward<BuiltinFuncT>(BuiltinFunc), *this);
218 } else {
219 return std::invoke(std::forward<BuiltinFuncT>(BuiltinFunc), *this);
220 }
221 }
222
227
228 template <typename BodyLambda = EmptyLambda>
229 auto forLoop(const LoopBoundInfo &Bounds, BodyLambda &&Body = {}) {
230 return ForLoopBuilder(Bounds, *this, std::move(Body));
231 }
232
233 template <typename... LoopBuilders>
234 auto buildLoopNest(LoopBuilders &&...Loops) {
235 return LoopNestBuilder(*this, std::forward<LoopBuilders>(Loops)...);
236 }
237
238 void ret(std::optional<std::reference_wrapper<Var>> OptRet = std::nullopt);
239
240 StringRef getName() const { return Name; }
241
243 Name = NewName.str();
244 Function *F = getFunction();
245 F->setName(Name);
246 }
247
248 // Convert the given Var's value to type T and return a new Var holding
249 // the converted value.
250 template <typename T>
251 std::enable_if_t<std::is_arithmetic_v<T>, Var &> convert(Var &V) {
252 auto &Ctx = getFunction()->getContext();
253 auto &IRBRef = getIRBuilder();
254 Type *TargetTy = TypeMap<T>::get(Ctx);
255 Var &Res = declVarInternal("convert.", TargetTy);
256 Value *Converted = proteus::convert(IRBRef, V.getValue(), TargetTy);
257 Res.storeValue(Converted);
258 return Res;
259 }
260};
261
262template <typename RetT, typename... ArgT> class Func final : public FuncBase {
263private:
264 Dispatcher &Dispatch;
265 RetT (*CompiledFunc)(ArgT...) = nullptr;
266
267private:
268 template <std::size_t... Is> auto getArgsImpl(std::index_sequence<Is...>) {
269 return std::tie(getArg(Is)...);
270 }
271
272public:
274 : FuncBase(J, FC), Dispatch(Dispatch) {}
275
276 RetT operator()(ArgT... Args);
277
278 auto getArgs() { return getArgsImpl(std::index_sequence_for<ArgT...>{}); }
279
280 auto getCompiledFunc() const { return CompiledFunc; }
281
283 CompiledFunc = CompiledFuncIn;
284 }
285};
286
287} // namespace proteus
288
289#endif // PROTEUS_FRONTEND_FUNC_HPP
char int void ** Args
Definition CompilerInterfaceHost.cpp:21
#define PROTEUS_FATAL_ERROR(x)
Definition Error.h:7
Definition Dispatcher.hpp:54
Definition LoopNest.hpp:22
Definition Func.hpp:39
void declArgs()
Definition Func.hpp:147
Var & atomicMax(Var &Addr, Var &Val)
Definition Func.cpp:285
std::vector< Scope > Scopes
Definition Func.hpp:63
void setName(StringRef NewName)
Definition Func.hpp:242
Var & declVarInternal(StringRef Name, Type *Ty, Type *PointerElemType=nullptr)
Definition Func.cpp:28
void beginIf(Var &CondVar, const char *File=__builtin_FILE(), int Line=__builtin_LINE())
Definition Func.cpp:146
std::deque< std::unique_ptr< Var > > Variables
Definition Func.hpp:47
std::enable_if_t<!std::is_void_v< typename FnSig< Sig >::RetT >, Var & > call(StringRef Name)
Definition JitFrontend.hpp:271
decltype(auto) callBuiltin(BuiltinFuncT &&BuiltinFunc)
Definition Func.hpp:214
auto forLoop(const LoopBoundInfo &Bounds, BodyLambda &&Body={})
Definition Func.hpp:229
Var & defVar(T Val, StringRef Name="var")
Definition Func.hpp:120
std::deque< std::unique_ptr< Var > > RuntimeConstants
Definition Func.hpp:48
auto defRuntimeConsts(ArgT &&...Args)
Definition Func.hpp:143
std::deque< std::unique_ptr< Var > > Arguments
Definition Func.hpp:46
void endFunction()
Definition Func.cpp:65
Var & emitAtomic(AtomicRMWInst::BinOp Op, Var &Addr, Var &Val)
Definition Func.cpp:254
void beginFunction(const char *File=__builtin_FILE(), int Line=__builtin_LINE())
Definition Func.cpp:42
Var & atomicSub(Var &Addr, Var &Val)
Definition Func.cpp:278
IRBuilder IRB
Definition Func.hpp:43
Var & defRuntimeConst(T Val, StringRef Name="run.const.var")
Definition Func.hpp:136
Function * getFunction()
Definition Func.cpp:78
ScopeKind
Definition Func.hpp:52
void endIf()
Definition Func.cpp:180
Var & getArg(unsigned int ArgNo)
Definition Func.cpp:85
std::string Name
Definition Func.hpp:50
FunctionCallee FC
Definition Func.hpp:42
Var & defVar(const Var &Val, StringRef Name="var")
Definition Func.hpp:127
auto buildLoopNest(LoopBuilders &&...Loops)
Definition Func.hpp:234
IRBuilderBase::InsertPoint IP
Definition Func.hpp:44
JitModule & J
Definition Func.hpp:41
void ret(std::optional< std::reference_wrapper< Var > > OptRet=std::nullopt)
Definition Func.cpp:129
Value * emitArrayCreate(Type *Ty, AddressSpace AT, StringRef Name)
Definition Func.cpp:99
StringRef getName() const
Definition Func.hpp:240
Var & atomicMin(Var &Addr, Var &Val)
Definition Func.cpp:292
void beginFor(Var &IterVar, Var &InitVar, Var &UpperBound, Var &IncVar, const char *File=__builtin_FILE(), int Line=__builtin_LINE())
Definition Func.cpp:196
Var & declVar(size_t NElem, AddressSpace AS=AddressSpace::DEFAULT, StringRef Name="array_var")
Definition Func.hpp:106
Var & atomicAdd(Var &Addr, Var &Val)
Definition Func.cpp:271
void endFor()
Definition Func.cpp:299
std::string toString(ScopeKind Kind)
Definition Func.hpp:65
std::enable_if_t< std::is_arithmetic_v< T >, Var & > convert(Var &V)
Definition Func.hpp:251
IRBuilderBase & getIRBuilder()
Definition Func.cpp:22
AllocaInst * emitAlloca(Type *Ty, StringRef Name, AddressSpace AS=AddressSpace::DEFAULT)
Definition Func.cpp:87
Var & declVar(StringRef Name="var")
Definition Func.hpp:98
Definition Func.hpp:262
Func(JitModule &J, FunctionCallee FC, Dispatcher &Dispatch)
Definition Func.hpp:273
auto getCompiledFunc() const
Definition Func.hpp:280
auto getArgs()
Definition Func.hpp:278
void setCompiledFunc(RetT(*CompiledFuncIn)(ArgT...))
Definition Func.hpp:282
RetT operator()(ArgT... Args)
Definition JitFrontend.hpp:323
Definition JitFrontend.hpp:28
Definition LoopNest.hpp:12
Definition LoopNest.hpp:43
Definition Helpers.h:138
Definition BuiltinsCUDA.cpp:4
AddressSpace
Definition AddressSpace.hpp:6
T getRuntimeConstantValue(void *Arg)
Definition CompilerInterfaceRuntimeConstantInfo.h:114
Value * convert(IRBuilderBase IRB, Value *V, Type *TargetType)
Definition Var.cpp:476
Definition Func.hpp:26
Definition Func.hpp:35
void operator()() const
Definition Func.hpp:36
RetT_ RetT
Definition Func.hpp:30
Definition Func.hpp:27
Definition Func.hpp:53
Scope(const char *File, int Line, ScopeKind Kind, IRBuilderBase::InsertPoint ContIP)
Definition Func.hpp:59
int Line
Definition Func.hpp:55
ScopeKind Kind
Definition Func.hpp:56
std::string File
Definition Func.hpp:54
IRBuilderBase::InsertPoint ContIP
Definition Func.hpp:57
Definition TypeMap.hpp:13
Definition Var.hpp:17