Enzyme main
Loading...
Searching...
No Matches
CallDerivatives.cpp
Go to the documentation of this file.
1//===- CallDerivatives.cpp - Implementation of known call derivatives --===//
2//
3// Enzyme Project
4//
5// Part of the Enzyme Project, under the Apache License v2.0 with LLVM
6// Exceptions. See https://llvm.org/LICENSE.txt for license information.
7// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
8//
9// If using this code in an academic setting, please cite the following:
10// @incollection{enzymeNeurips,
11// title = {Instead of Rewriting Foreign Code for Machine Learning,
12// Automatically Synthesize Fast Gradients},
13// author = {Moses, William S. and Churavy, Valentin},
14// booktitle = {Advances in Neural Information Processing Systems 33},
15// year = {2020},
16// note = {To appear in},
17// }
18//
19//===----------------------------------------------------------------------===//
20//
21// This file contains the implementation of functions in instruction visitor
22// AdjointGenerator that generate corresponding augmented forward pass code,
23// and adjoints for certain known functions.
24//
25//===----------------------------------------------------------------------===//
26
27#include "AdjointGenerator.h"
28
29using namespace llvm;
30
31extern "C" {
32void (*EnzymeShadowAllocRewrite)(LLVMValueRef, void *, LLVMValueRef, uint64_t,
33 LLVMValueRef, uint8_t) = nullptr;
34}
35
36/// Whether \p Begin has exactly one llvm.julia.gc_preserve_end outside the
37/// abort-only blocks the reverse pass drops, as reversing the region requires;
38/// CanonicalizeGCPreserveEnds merges them where it can.
39static bool hasSingleGCPreserveEnd(CallInst *Begin, GradientUtils *gutils) {
40 unsigned Ends = 0;
41 for (auto U : Begin->users())
42 if (auto CI = dyn_cast<CallInst>(U))
43 if (getFuncNameFromCall(CI) == "llvm.julia.gc_preserve_end" &&
44 !gutils->notForAnalysis.count(CI->getParent()))
45 ++Ends;
46 return Ends == 1;
47}
48
49/// Forward mode of an MPI call whose derivative is the same call with the
50/// buffers \p ShadowArgs replaced by their shadows (for the reductions, of a
51/// sum). At vector width > 1 the call is replayed once per lane. All other
52/// arguments, including the trailing `ierr` of the Fortran ABI, are passed as
53/// in the original call. An inactive constant buffer, such as MPI_IN_PLACE or
54/// MPI_STATUS_IGNORE, is passed unchanged.
55static void createMPIForwardCall(CallInst &call, ArrayRef<unsigned> ShadowArgs,
56 GradientUtils *gutils, IRBuilder<> &Builder2) {
57 std::vector<ValueType> BundleTypes(call.arg_size(), ValueType::Primal);
58 SmallVector<Value *, 2> shadows;
59 for (unsigned i : ShadowArgs) {
60 Value *arg = call.getArgOperand(i);
61 if (isa<Constant>(arg) && gutils->isConstantValue(arg)) {
62 shadows.push_back(nullptr);
63 continue;
64 }
65 BundleTypes[i] = ValueType::Shadow;
66 shadows.push_back(gutils->invertPointerM(arg, Builder2));
67 }
68 auto Defs = gutils->getInvertedBundles(&call, BundleTypes, Builder2,
69 /*lookup*/ false);
70
71 for (unsigned l = 0; l < gutils->getWidth(); l++) {
72 SmallVector<Value *, 8> args;
73 for (auto &op : call.args())
74 args.push_back(gutils->getNewFromOriginal(op));
75 for (auto en : llvm::enumerate(ShadowArgs)) {
76 Value *sh = shadows[en.index()];
77 if (!sh)
78 continue;
79 if (gutils->getWidth() > 1)
80 sh = gutils->extractMeta(Builder2, sh, l);
81 Type *T = args[en.value()]->getType();
82 if (sh->getType()->isIntegerTy())
83 sh = Builder2.CreateIntToPtr(sh, T);
84 else if (sh->getType() != T)
85 sh = Builder2.CreatePointerCast(sh, T);
86 args[en.value()] = sh;
87 }
88 auto dcall = Builder2.CreateCall(call.getFunctionType(),
89 call.getCalledOperand(), args, Defs);
90 dcall->setCallingConv(call.getCallingConv());
91 dcall->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
92 }
93}
94
95void AdjointGenerator::handleMPI(llvm::CallInst &call, llvm::Function *called,
96 llvm::StringRef funcName) {
97 using namespace llvm;
98
99 assert(called);
100
101 IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&call));
102 BuilderZ.setFastMathFlags(getFast());
103
104 // In forward mode, the derivative of an MPI call is the call on the shadows,
105 // which is replayed per lane at vector width > 1 (createMPIForwardCall).
106 if (gutils->getWidth() > 1 && Mode != DerivativeMode::ForwardMode &&
107 Mode != DerivativeMode::ForwardModeError) {
108 std::string s;
109 raw_string_ostream ss(s);
110 ss << funcName << " is not supported in vector reverse mode: " << call;
111 EmitNoDerivativeError(ss.str(), call, gutils, BuilderZ);
112 return;
113 }
114
115 // MPI send / recv can only send float/integers
116 if (funcName == "PMPI_Isend" || funcName == "MPI_Isend" ||
117 funcName == "PMPI_Irecv" || funcName == "MPI_Irecv") {
118 if (!gutils->isConstantInstruction(&call)) {
119 if (Mode == DerivativeMode::ReverseModePrimal ||
120 Mode == DerivativeMode::ReverseModeCombined) {
121 assert(!gutils->isConstantValue(call.getOperand(0)));
122 assert(!gutils->isConstantValue(call.getOperand(6)));
123 Value *d_req = gutils->invertPointerM(call.getOperand(6), BuilderZ);
124 if (d_req->getType()->isIntegerTy()) {
125 d_req = BuilderZ.CreateIntToPtr(
126 d_req, getUnqual(getInt8PtrTy(call.getContext())));
127 }
128
129 auto i64 = Type::getInt64Ty(call.getContext());
130 auto impi = getMPIHelper(call.getContext());
131
132 Value *impialloc =
133 CreateAllocation(BuilderZ, impi, ConstantInt::get(i64, 1));
134 BuilderZ.SetInsertPoint(gutils->getNewFromOriginal(&call));
135
136 d_req = BuilderZ.CreateBitCast(d_req, getUnqual(impialloc->getType()));
137 Value *d_req_prev = BuilderZ.CreateLoad(impialloc->getType(), d_req);
138 BuilderZ.CreateStore(
139 BuilderZ.CreatePointerCast(d_req_prev,
140 getInt8PtrTy(call.getContext())),
141 getMPIMemberPtr<MPI_Elem::Old>(BuilderZ, impialloc, impi));
142 BuilderZ.CreateStore(impialloc, d_req);
143
144 if (funcName == "MPI_Isend" || funcName == "PMPI_Isend") {
145 Value *tysize =
146 MPI_TYPE_SIZE(gutils->getNewFromOriginal(call.getOperand(2)),
147 BuilderZ, call.getType(), called);
148
149 auto len_arg = BuilderZ.CreateZExtOrTrunc(
150 gutils->getNewFromOriginal(call.getOperand(1)),
151 Type::getInt64Ty(call.getContext()));
152 len_arg = BuilderZ.CreateMul(
153 len_arg,
154 BuilderZ.CreateZExtOrTrunc(tysize,
155 Type::getInt64Ty(call.getContext())),
156 "", true, true);
157
158 Value *firstallocation =
159 CreateAllocation(BuilderZ, Type::getInt8Ty(call.getContext()),
160 len_arg, "mpirecv_malloccache");
161 BuilderZ.CreateStore(firstallocation, getMPIMemberPtr<MPI_Elem::Buf>(
162 BuilderZ, impialloc, impi));
163 BuilderZ.SetInsertPoint(gutils->getNewFromOriginal(&call));
164 } else {
165 Value *ibuf = gutils->invertPointerM(call.getOperand(0), BuilderZ);
166 if (ibuf->getType()->isIntegerTy())
167 ibuf =
168 BuilderZ.CreateIntToPtr(ibuf, getInt8PtrTy(call.getContext()));
169 BuilderZ.CreateStore(
170 ibuf, getMPIMemberPtr<MPI_Elem::Buf>(BuilderZ, impialloc, impi));
171 }
172
173 BuilderZ.CreateStore(
174 BuilderZ.CreateZExtOrTrunc(
175 gutils->getNewFromOriginal(call.getOperand(1)), i64),
176 getMPIMemberPtr<MPI_Elem::Count>(BuilderZ, impialloc, impi));
177
178 Value *dataType = gutils->getNewFromOriginal(call.getOperand(2));
179 if (dataType->getType()->isIntegerTy())
180 dataType = BuilderZ.CreateIntToPtr(
181 dataType, getInt8PtrTy(dataType->getContext()));
182 BuilderZ.CreateStore(
183 BuilderZ.CreatePointerCast(dataType,
184 getInt8PtrTy(call.getContext())),
185 getMPIMemberPtr<MPI_Elem::DataType>(BuilderZ, impialloc, impi));
186
187 BuilderZ.CreateStore(
188 BuilderZ.CreateZExtOrTrunc(
189 gutils->getNewFromOriginal(call.getOperand(3)), i64),
190 getMPIMemberPtr<MPI_Elem::Src>(BuilderZ, impialloc, impi));
191
192 BuilderZ.CreateStore(
193 BuilderZ.CreateZExtOrTrunc(
194 gutils->getNewFromOriginal(call.getOperand(4)), i64),
195 getMPIMemberPtr<MPI_Elem::Tag>(BuilderZ, impialloc, impi));
196
197 Value *comm = gutils->getNewFromOriginal(call.getOperand(5));
198 if (comm->getType()->isIntegerTy())
199 comm = BuilderZ.CreateIntToPtr(comm,
200 getInt8PtrTy(dataType->getContext()));
201 BuilderZ.CreateStore(
202 BuilderZ.CreatePointerCast(comm, getInt8PtrTy(call.getContext())),
203 getMPIMemberPtr<MPI_Elem::Comm>(BuilderZ, impialloc, impi));
204
205 BuilderZ.CreateStore(
206 ConstantInt::get(
207 Type::getInt8Ty(impialloc->getContext()),
208 (funcName == "MPI_Isend" || funcName == "PMPI_Isend")
209 ? (int)MPI_CallType::ISEND
210 : (int)MPI_CallType::IRECV),
211 getMPIMemberPtr<MPI_Elem::Call>(BuilderZ, impialloc, impi));
212 // TODO old
213 }
214 if (Mode == DerivativeMode::ReverseModeGradient ||
215 Mode == DerivativeMode::ReverseModeCombined) {
216 IRBuilder<> Builder2(&call);
217 getReverseBuilder(Builder2);
218
219 Type *statusType = nullptr;
220#if LLVM_VERSION_MAJOR < 17
221 if (Function *recvfn = called->getParent()->getFunction(
222 getRenamedPerCallingConv(called->getName(), "MPI_Wait"))) {
223 auto statusArg = recvfn->arg_end();
224 statusArg--;
225 if (auto PT = dyn_cast<PointerType>(statusArg->getType()))
226 statusType = PT->getPointerElementType();
227 }
228#endif
229 if (statusType == nullptr) {
230 statusType = ArrayType::get(Type::getInt8Ty(call.getContext()), 24);
231 llvm::errs() << " warning could not automatically determine mpi "
232 "status type, assuming [24 x i8]\n";
233 }
234 Value *req =
235 lookup(gutils->getNewFromOriginal(call.getOperand(6)), Builder2);
236 Value *d_req = lookup(
237 gutils->invertPointerM(call.getOperand(6), Builder2), Builder2);
238 if (d_req->getType()->isIntegerTy()) {
239 d_req =
240 Builder2.CreateIntToPtr(d_req, getInt8PtrTy(call.getContext()));
241 }
242 auto impi = getMPIHelper(call.getContext());
243 Type *helperTy = getUnqual(impi);
244 Value *helper = Builder2.CreatePointerCast(d_req, getUnqual(helperTy));
245 helper = Builder2.CreateLoad(helperTy, helper);
246
247 auto i64 = Type::getInt64Ty(call.getContext());
248
249 Value *firstallocation;
250 firstallocation = Builder2.CreateLoad(
251 getInt8PtrTy(call.getContext()),
252 getMPIMemberPtr<MPI_Elem::Buf>(Builder2, helper, impi));
253 Value *len_arg = nullptr;
254 if (auto C = dyn_cast<Constant>(
255 gutils->getNewFromOriginal(call.getOperand(1)))) {
256 len_arg = Builder2.CreateZExtOrTrunc(C, i64);
257 } else {
258 len_arg = Builder2.CreateLoad(
259 i64, getMPIMemberPtr<MPI_Elem::Count>(Builder2, helper, impi));
260 }
261 Value *tysize = nullptr;
262 if (auto C = dyn_cast<Constant>(
263 gutils->getNewFromOriginal(call.getOperand(2)))) {
264 tysize = C;
265 } else {
266 tysize = Builder2.CreateLoad(
267 getInt8PtrTy(call.getContext()),
268 getMPIMemberPtr<MPI_Elem::DataType>(Builder2, helper, impi));
269 }
270
271 Value *prev;
272 prev = Builder2.CreateLoad(
273 getInt8PtrTy(call.getContext()),
274 getMPIMemberPtr<MPI_Elem::Old>(Builder2, helper, impi));
275
276 Builder2.CreateStore(prev, Builder2.CreatePointerCast(
277 d_req, getUnqual(prev->getType())));
278
279 assert(shouldFree());
280
281 assert(tysize);
282 tysize = MPI_TYPE_SIZE(tysize, Builder2, call.getType(), called);
283
284 Value *args[] = {/*req*/ req,
285 /*status*/ IRBuilder<>(gutils->inversionAllocs)
286 .CreateAlloca(statusType)};
287 FunctionCallee waitFunc = nullptr;
288 for (auto name : {
289 "MPI_Wait",
290 })
291 if (Function *recvfn = called->getParent()->getFunction(
292 getRenamedPerCallingConv(called->getName(), name))) {
293 auto statusArg = recvfn->arg_end();
294 statusArg--;
295 if (statusArg->getType()->isIntegerTy())
296 args[1] = Builder2.CreatePtrToInt(args[1], statusArg->getType());
297 else
298 args[1] = Builder2.CreateBitCast(args[1], statusArg->getType());
299 waitFunc = recvfn;
300 break;
301 }
302 if (!waitFunc) {
303 Type *types[sizeof(args) / sizeof(*args)];
304 for (size_t i = 0; i < sizeof(args) / sizeof(*args); i++)
305 types[i] = args[i]->getType();
306 FunctionType *FT = FunctionType::get(call.getType(), types, false);
307 waitFunc = called->getParent()->getOrInsertFunction(
308 getRenamedPerCallingConv(called->getName(), "MPI_Wait"), FT);
309 }
310 assert(waitFunc);
311
312 // Need to preserve the shadow Request (operand 6 in isend/irecv),
313 // which becomes operand 0 for iwait.
314 auto ReqDefs = gutils->getInvertedBundles(
315 &call,
316 {ValueType::None, ValueType::None, ValueType::None, ValueType::None,
317 ValueType::None, ValueType::None, ValueType::Shadow},
318 Builder2, /*lookup*/ true);
319
320 auto BufferDefs = gutils->getInvertedBundles(
321 &call,
322 {ValueType::Shadow, ValueType::None, ValueType::None,
323 ValueType::None, ValueType::None, ValueType::None,
324 ValueType::None},
325 Builder2, /*lookup*/ true);
326
327 auto fcall = Builder2.CreateCall(waitFunc, args, ReqDefs);
328 fcall->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
329 if (auto F = dyn_cast<Function>(waitFunc.getCallee()))
330 fcall->setCallingConv(F->getCallingConv());
331 len_arg = Builder2.CreateMul(
332 len_arg,
333 Builder2.CreateZExtOrTrunc(tysize,
334 Type::getInt64Ty(Builder2.getContext())),
335 "", true, true);
336 if (funcName == "MPI_Irecv" || funcName == "PMPI_Irecv") {
337 auto val_arg =
338 ConstantInt::get(Type::getInt8Ty(Builder2.getContext()), 0);
339 auto volatile_arg = ConstantInt::getFalse(Builder2.getContext());
340 assert(!gutils->isConstantValue(call.getOperand(0)));
341 auto dbuf = firstallocation;
342 Value *nargs[] = {dbuf, val_arg, len_arg, volatile_arg};
343 Type *tys[] = {dbuf->getType(), len_arg->getType()};
344
345 auto memset = cast<CallInst>(Builder2.CreateCall(
346 getIntrinsicDeclaration(called->getParent(), Intrinsic::memset,
347 tys),
348 nargs, BufferDefs));
349 memset->addParamAttr(0, Attribute::NonNull);
350 } else if (funcName == "MPI_Isend" || funcName == "PMPI_Isend") {
351 assert(!gutils->isConstantValue(call.getOperand(0)));
352 Value *shadow = lookup(
353 gutils->invertPointerM(call.getOperand(0), Builder2), Builder2);
354
355 // TODO add operand bundle (unless force inlined?)
356 DifferentiableMemCopyFloats(call, call.getOperand(0), firstallocation,
357 shadow, len_arg, Builder2, BufferDefs);
358
359 if (shouldFree()) {
360 CreateDealloc(Builder2, firstallocation);
361 }
362 } else
363 assert(0 && "illegal mpi");
364
365 CreateDealloc(Builder2, helper);
366 }
367 if (Mode == DerivativeMode::ForwardMode ||
368 Mode == DerivativeMode::ForwardModeError) {
369 IRBuilder<> Builder2(&call);
370 getForwardBuilder(Builder2);
371
372 assert(!gutils->isConstantValue(call.getOperand(0)));
373 assert(!gutils->isConstantValue(call.getOperand(6)));
374
375 createMPIForwardCall(call, {/*buf*/ 0, /*request*/ 6}, gutils,
376 Builder2);
377 return;
378 }
379 }
380 if (Mode == DerivativeMode::ReverseModeGradient)
381 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
382 return;
383 }
384
385 if (funcName == "MPI_Wait" || funcName == "PMPI_Wait") {
386 Value *d_reqp = nullptr;
387 auto impi = getMPIHelper(call.getContext());
388 if (Mode == DerivativeMode::ReverseModePrimal ||
389 Mode == DerivativeMode::ReverseModeCombined) {
390 Value *req = gutils->getNewFromOriginal(call.getOperand(0));
391 Value *d_req = gutils->invertPointerM(call.getOperand(0), BuilderZ);
392
393 if (req->getType()->isIntegerTy()) {
394 req = BuilderZ.CreateIntToPtr(
395 req, getUnqual(getInt8PtrTy(call.getContext())));
396 }
397
398 Value *isNull = nullptr;
399 if (auto GV = gutils->newFunc->getParent()->getNamedValue(
400 "ompi_request_null")) {
401 Value *reql = BuilderZ.CreatePointerCast(req, getUnqual(GV->getType()));
402 reql = BuilderZ.CreateLoad(GV->getType(), reql);
403 isNull = BuilderZ.CreateICmpEQ(reql, GV);
404 }
405
406 if (d_req->getType()->isIntegerTy()) {
407 d_req = BuilderZ.CreateIntToPtr(
408 d_req, getUnqual(getInt8PtrTy(call.getContext())));
409 }
410
411 d_reqp = BuilderZ.CreateLoad(
412 getUnqual(impi),
413 BuilderZ.CreatePointerCast(d_req, getUnqual(getUnqual(impi))));
414 if (isNull)
415 d_reqp =
416 CreateSelect(BuilderZ, isNull,
417 Constant::getNullValue(d_reqp->getType()), d_reqp);
418 if (auto I = dyn_cast<Instruction>(d_reqp))
419 gutils->TapesToPreventRecomputation.insert(I);
420 d_reqp = gutils->cacheForReverse(
421 BuilderZ, d_reqp, getIndex(&call, CacheType::Tape, BuilderZ));
422 }
423 if (Mode == DerivativeMode::ReverseModeGradient ||
424 Mode == DerivativeMode::ReverseModeCombined) {
425 IRBuilder<> Builder2(&call);
426 getReverseBuilder(Builder2);
427
428 assert(!gutils->isConstantValue(call.getOperand(0)));
429 Value *req =
430 lookup(gutils->getNewFromOriginal(call.getOperand(0)), Builder2);
431
432 if (Mode != DerivativeMode::ReverseModeCombined) {
433 d_reqp = BuilderZ.CreatePHI(getUnqual(impi), 0);
434 d_reqp = gutils->cacheForReverse(
435 BuilderZ, d_reqp, getIndex(&call, CacheType::Tape, BuilderZ));
436 } else
437 assert(d_reqp);
438 d_reqp = lookup(d_reqp, Builder2);
439
440 Value *isNull = Builder2.CreateICmpEQ(
441 d_reqp, Constant::getNullValue(d_reqp->getType()));
442
443 BasicBlock *currentBlock = Builder2.GetInsertBlock();
444 BasicBlock *nonnullBlock = gutils->addReverseBlock(
445 currentBlock, currentBlock->getName() + "_nonnull");
446 BasicBlock *endBlock = gutils->addReverseBlock(
447 nonnullBlock, currentBlock->getName() + "_end",
448 /*fork*/ true, /*push*/ false);
449
450 Builder2.CreateCondBr(isNull, endBlock, nonnullBlock);
451 Builder2.SetInsertPoint(nonnullBlock);
452
453 Value *cache = Builder2.CreateLoad(impi, d_reqp);
454
455 Value *args[] = {
456 getMPIMemberPtr<MPI_Elem::Buf, false>(Builder2, cache, impi),
457 getMPIMemberPtr<MPI_Elem::Count, false>(Builder2, cache, impi),
458 getMPIMemberPtr<MPI_Elem::DataType, false>(Builder2, cache, impi),
459 getMPIMemberPtr<MPI_Elem::Src, false>(Builder2, cache, impi),
460 getMPIMemberPtr<MPI_Elem::Tag, false>(Builder2, cache, impi),
461 getMPIMemberPtr<MPI_Elem::Comm, false>(Builder2, cache, impi),
462 getMPIMemberPtr<MPI_Elem::Call, false>(Builder2, cache, impi),
463 req};
464 Type *types[sizeof(args) / sizeof(*args) - 1];
465 for (size_t i = 0; i < sizeof(args) / sizeof(*args) - 1; i++)
466 types[i] = args[i]->getType();
467 Function *dwait = getOrInsertDifferentialMPI_Wait(
468 *called->getParent(), types, call.getOperand(0)->getType(),
469 called->getName());
470
471 // Need to preserve the shadow Request (operand 0 in wait).
472 // However, this doesn't end up preserving
473 // the underlying buffers for the adjoint. To rememdy, force inline.
474 auto cal =
475 Builder2.CreateCall(dwait, args,
476 gutils->getInvertedBundles(
477 &call, {ValueType::Shadow, ValueType::None},
478 Builder2, /*lookup*/ true));
479 cal->setCallingConv(dwait->getCallingConv());
480 cal->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
481 cal->addFnAttr(Attribute::AlwaysInline);
482 Builder2.CreateBr(endBlock);
483 {
484 auto found = gutils->reverseBlockToPrimal.find(endBlock);
485 assert(found != gutils->reverseBlockToPrimal.end());
486 SmallVector<BasicBlock *, 4> &vec =
487 gutils->reverseBlocks[found->second];
488 assert(vec.size());
489 vec.push_back(endBlock);
490 }
491 Builder2.SetInsertPoint(endBlock);
492 } else if (Mode == DerivativeMode::ForwardMode ||
493 Mode == DerivativeMode::ForwardModeError) {
494 IRBuilder<> Builder2(&call);
495 getForwardBuilder(Builder2);
496
497 assert(!gutils->isConstantValue(call.getOperand(0)));
498
499 createMPIForwardCall(call, {/*request*/ 0, /*status*/ 1}, gutils,
500 Builder2);
501 return;
502 }
503 if (Mode == DerivativeMode::ReverseModeGradient)
504 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
505 return;
506 }
507
508 if (funcName == "MPI_Waitall" || funcName == "PMPI_Waitall") {
509 Value *d_reqp = nullptr;
510 auto impi = getMPIHelper(call.getContext());
511 PointerType *reqType = getUnqual(impi);
512 if (Mode == DerivativeMode::ReverseModePrimal ||
513 Mode == DerivativeMode::ReverseModeCombined) {
514 Value *count = gutils->getNewFromOriginal(call.getOperand(0));
515 Value *req = gutils->getNewFromOriginal(call.getOperand(1));
516 Value *d_req = gutils->invertPointerM(call.getOperand(1), BuilderZ);
517
518 if (req->getType()->isIntegerTy()) {
519 req = BuilderZ.CreateIntToPtr(
520 req, getUnqual(getInt8PtrTy(call.getContext())));
521 }
522
523 if (d_req->getType()->isIntegerTy()) {
524 d_req = BuilderZ.CreateIntToPtr(
525 d_req, getUnqual(getInt8PtrTy(call.getContext())));
526 }
527
528 Function *dsave = getOrInsertDifferentialWaitallSave(
529 *gutils->oldFunc->getParent(),
530 {count->getType(), req->getType(), d_req->getType()}, reqType);
531
532 d_reqp = BuilderZ.CreateCall(dsave, {count, req, d_req});
533 cast<CallInst>(d_reqp)->setCallingConv(dsave->getCallingConv());
534 cast<CallInst>(d_reqp)->setDebugLoc(
535 gutils->getNewFromOriginal(call.getDebugLoc()));
536 d_reqp = gutils->cacheForReverse(
537 BuilderZ, d_reqp, getIndex(&call, CacheType::Tape, BuilderZ));
538 }
539 if (Mode == DerivativeMode::ReverseModeGradient ||
540 Mode == DerivativeMode::ReverseModeCombined) {
541 IRBuilder<> Builder2(&call);
542 getReverseBuilder(Builder2);
543
544 assert(!gutils->isConstantValue(call.getOperand(1)));
545 Value *count =
546 lookup(gutils->getNewFromOriginal(call.getOperand(0)), Builder2);
547 Value *req_orig =
548 lookup(gutils->getNewFromOriginal(call.getOperand(1)), Builder2);
549
550 if (Mode != DerivativeMode::ReverseModeCombined) {
551 d_reqp = BuilderZ.CreatePHI(getUnqual(reqType), 0);
552 d_reqp = gutils->cacheForReverse(
553 BuilderZ, d_reqp, getIndex(&call, CacheType::Tape, BuilderZ));
554 }
555
556 d_reqp = lookup(d_reqp, Builder2);
557
558 BasicBlock *currentBlock = Builder2.GetInsertBlock();
559 BasicBlock *loopBlock = gutils->addReverseBlock(
560 currentBlock, currentBlock->getName() + "_loop");
561 BasicBlock *nonnullBlock = gutils->addReverseBlock(
562 loopBlock, currentBlock->getName() + "_nonnull");
563 BasicBlock *eloopBlock = gutils->addReverseBlock(
564 nonnullBlock, currentBlock->getName() + "_eloop");
565 BasicBlock *endBlock =
566 gutils->addReverseBlock(eloopBlock, currentBlock->getName() + "_end",
567 /*fork*/ true, /*push*/ false);
568
569 Builder2.CreateCondBr(
570 Builder2.CreateICmpNE(count,
571 ConstantInt::get(count->getType(), 0, false)),
572 loopBlock, endBlock);
573
574 Builder2.SetInsertPoint(loopBlock);
575 auto idx = Builder2.CreatePHI(count->getType(), 2);
576 idx->addIncoming(ConstantInt::get(count->getType(), 0, false),
577 currentBlock);
578 Value *inc = Builder2.CreateAdd(
579 idx, ConstantInt::get(count->getType(), 1, false), "", true, true);
580 idx->addIncoming(inc, eloopBlock);
581
582 Value *idxs[] = {idx};
583 Value *req = Builder2.CreateInBoundsGEP(reqType, req_orig, idxs);
584 Value *d_req = Builder2.CreateInBoundsGEP(reqType, d_reqp, idxs);
585
586 d_req = Builder2.CreateLoad(
587 getUnqual(impi),
588 Builder2.CreatePointerCast(d_req, getUnqual(getUnqual(impi))));
589
590 Value *isNull = Builder2.CreateICmpEQ(
591 d_req, Constant::getNullValue(d_req->getType()));
592
593 Builder2.CreateCondBr(isNull, eloopBlock, nonnullBlock);
594 Builder2.SetInsertPoint(nonnullBlock);
595
596 Value *cache = Builder2.CreateLoad(impi, d_req);
597
598 Value *args[] = {
599 getMPIMemberPtr<MPI_Elem::Buf, false>(Builder2, cache, impi),
600 getMPIMemberPtr<MPI_Elem::Count, false>(Builder2, cache, impi),
601 getMPIMemberPtr<MPI_Elem::DataType, false>(Builder2, cache, impi),
602 getMPIMemberPtr<MPI_Elem::Src, false>(Builder2, cache, impi),
603 getMPIMemberPtr<MPI_Elem::Tag, false>(Builder2, cache, impi),
604 getMPIMemberPtr<MPI_Elem::Comm, false>(Builder2, cache, impi),
605 getMPIMemberPtr<MPI_Elem::Call, false>(Builder2, cache, impi),
606 req};
607 Type *types[sizeof(args) / sizeof(*args) - 1];
608 for (size_t i = 0; i < sizeof(args) / sizeof(*args) - 1; i++)
609 types[i] = args[i]->getType();
610 Function *dwait = getOrInsertDifferentialMPI_Wait(
611 *called->getParent(), types, req->getType(), called->getName());
612 // Need to preserve the shadow Request (operand 6 in isend/irecv), which
613 // becomes operand 0 for iwait. However, this doesn't end up preserving
614 // the underlying buffers for the adjoint. To remedy, force inline the
615 // function.
616 auto cal = Builder2.CreateCall(
617 dwait, args,
618 gutils->getInvertedBundles(&call,
619 {ValueType::None, ValueType::None,
620 ValueType::None, ValueType::None,
621 ValueType::None, ValueType::None,
622 ValueType::Shadow},
623 Builder2, /*lookup*/ true));
624 cal->setCallingConv(dwait->getCallingConv());
625 cal->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
626 cal->addFnAttr(Attribute::AlwaysInline);
627 Builder2.CreateBr(eloopBlock);
628
629 Builder2.SetInsertPoint(eloopBlock);
630 Builder2.CreateCondBr(Builder2.CreateICmpEQ(inc, count), endBlock,
631 loopBlock);
632 {
633 auto found = gutils->reverseBlockToPrimal.find(endBlock);
634 assert(found != gutils->reverseBlockToPrimal.end());
635 SmallVector<BasicBlock *, 4> &vec =
636 gutils->reverseBlocks[found->second];
637 assert(vec.size());
638 vec.push_back(endBlock);
639 }
640 Builder2.SetInsertPoint(endBlock);
641 if (shouldFree()) {
642 CreateDealloc(Builder2, d_reqp);
643 }
644 } else if (Mode == DerivativeMode::ForwardMode ||
645 Mode == DerivativeMode::ForwardModeError) {
646 IRBuilder<> Builder2(&call);
647 getForwardBuilder(Builder2);
648
649 assert(!gutils->isConstantValue(call.getOperand(1)));
650
651 createMPIForwardCall(call, {/*array_of_requests*/ 1}, gutils, Builder2);
652 return;
653 }
654 if (Mode == DerivativeMode::ReverseModeGradient)
655 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
656 return;
657 }
658
659 if (funcName == "MPI_Send" || funcName == "MPI_Ssend" ||
660 funcName == "PMPI_Send" || funcName == "PMPI_Ssend") {
661 if (Mode == DerivativeMode::ReverseModeGradient ||
662 Mode == DerivativeMode::ReverseModeCombined ||
663 Mode == DerivativeMode::ForwardMode ||
664 Mode == DerivativeMode::ForwardModeError) {
665 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
666 Mode == DerivativeMode::ForwardModeError;
667
668 IRBuilder<> Builder2 =
669 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
670 if (forwardMode) {
671 getForwardBuilder(Builder2);
672 } else {
673 getReverseBuilder(Builder2);
674 }
675
676 Value *shadow = gutils->invertPointerM(call.getOperand(0), Builder2);
677 if (!forwardMode)
678 shadow = lookup(shadow, Builder2);
679 if (shadow->getType()->isIntegerTy())
680 shadow =
681 Builder2.CreateIntToPtr(shadow, getInt8PtrTy(call.getContext()));
682
683 Type *statusType = nullptr;
684#if LLVM_VERSION_MAJOR < 17
685 if (called->getContext().supportsTypedPointers()) {
686 if (Function *recvfn = called->getParent()->getFunction(
687 getRenamedPerCallingConv(called->getName(), "MPI_Recv"))) {
688 auto statusArg = recvfn->arg_end();
689 statusArg--;
690 if (auto PT = dyn_cast<PointerType>(statusArg->getType()))
691 statusType = PT->getPointerElementType();
692 }
693 }
694#endif
695 if (statusType == nullptr) {
696 statusType = ArrayType::get(Type::getInt8Ty(call.getContext()), 24);
697 llvm::errs() << " warning could not automatically determine mpi "
698 "status type, assuming [24 x i8]\n";
699 }
700
701 Value *count = gutils->getNewFromOriginal(call.getOperand(1));
702 if (!forwardMode)
703 count = lookup(count, Builder2);
704
705 Value *datatype = gutils->getNewFromOriginal(call.getOperand(2));
706 if (!forwardMode)
707 datatype = lookup(datatype, Builder2);
708
709 Value *src = gutils->getNewFromOriginal(call.getOperand(3));
710 if (!forwardMode)
711 src = lookup(src, Builder2);
712
713 Value *tag = gutils->getNewFromOriginal(call.getOperand(4));
714 if (!forwardMode)
715 tag = lookup(tag, Builder2);
716
717 Value *comm = gutils->getNewFromOriginal(call.getOperand(5));
718 if (!forwardMode)
719 comm = lookup(comm, Builder2);
720
721 if (forwardMode) {
722 createMPIForwardCall(call, {/*buf*/ 0}, gutils, Builder2);
723 return;
724 }
725
726 Value *args[] = {
727 /*buf*/ NULL,
728 /*count*/ count,
729 /*datatype*/ datatype,
730 /*src*/ src,
731 /*tag*/ tag,
732 /*comm*/ comm,
733 /*status*/
734 IRBuilder<>(gutils->inversionAllocs).CreateAlloca(statusType)};
735
736 Value *tysize = MPI_TYPE_SIZE(datatype, Builder2, call.getType(), called);
737
738 auto len_arg = Builder2.CreateZExtOrTrunc(
739 args[1], Type::getInt64Ty(call.getContext()));
740 len_arg =
741 Builder2.CreateMul(len_arg,
742 Builder2.CreateZExtOrTrunc(
743 tysize, Type::getInt64Ty(call.getContext())),
744 "", true, true);
745
746 Value *firstallocation =
747 CreateAllocation(Builder2, Type::getInt8Ty(call.getContext()),
748 len_arg, "mpirecv_malloccache");
749 args[0] = firstallocation;
750
751 Type *types[sizeof(args) / sizeof(*args)];
752 for (size_t i = 0; i < sizeof(args) / sizeof(*args); i++)
753 types[i] = args[i]->getType();
754 FunctionType *FT = FunctionType::get(call.getType(), types, false);
755
756 Builder2.SetInsertPoint(Builder2.GetInsertBlock());
757
758 auto BufferDefs = gutils->getInvertedBundles(
759 &call,
760 {ValueType::Shadow, ValueType::None, ValueType::None, ValueType::None,
761 ValueType::None, ValueType::None, ValueType::None},
762 Builder2, /*lookup*/ true);
763
764 auto fcall = Builder2.CreateCall(
765 called->getParent()->getOrInsertFunction(
766 getRenamedPerCallingConv(called->getName(), "MPI_Recv"), FT),
767 args);
768 fcall->setCallingConv(call.getCallingConv());
769
770 DifferentiableMemCopyFloats(call, call.getOperand(0), firstallocation,
771 shadow, len_arg, Builder2, BufferDefs);
772
773 if (shouldFree()) {
774 CreateDealloc(Builder2, firstallocation);
775 }
776 }
777 if (Mode == DerivativeMode::ReverseModeGradient)
778 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
779 return;
780 }
781
782 if (funcName == "MPI_Recv" || funcName == "PMPI_Recv") {
783 if (Mode == DerivativeMode::ReverseModeGradient ||
784 Mode == DerivativeMode::ReverseModeCombined ||
785 Mode == DerivativeMode::ForwardMode ||
786 Mode == DerivativeMode::ForwardModeError) {
787 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
788 Mode == DerivativeMode::ForwardModeError;
789
790 IRBuilder<> Builder2 =
791 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
792 if (forwardMode) {
793 getForwardBuilder(Builder2);
794 } else {
795 getReverseBuilder(Builder2);
796 }
797
798 Value *shadow = gutils->invertPointerM(call.getOperand(0), Builder2);
799 if (!forwardMode)
800 shadow = lookup(shadow, Builder2);
801
802 Value *count = gutils->getNewFromOriginal(call.getOperand(1));
803 if (!forwardMode)
804 count = lookup(count, Builder2);
805
806 Value *datatype = gutils->getNewFromOriginal(call.getOperand(2));
807 if (!forwardMode)
808 datatype = lookup(datatype, Builder2);
809
810 Value *source = gutils->getNewFromOriginal(call.getOperand(3));
811 if (!forwardMode)
812 source = lookup(source, Builder2);
813
814 Value *tag = gutils->getNewFromOriginal(call.getOperand(4));
815 if (!forwardMode)
816 tag = lookup(tag, Builder2);
817
818 Value *comm = gutils->getNewFromOriginal(call.getOperand(5));
819 if (!forwardMode)
820 comm = lookup(comm, Builder2);
821
822 if (forwardMode) {
823 createMPIForwardCall(call, {/*buf*/ 0}, gutils, Builder2);
824 return;
825 }
826
827 Value *args[] = {shadow, count, datatype, source, tag, comm};
828
829 auto Defs = gutils->getInvertedBundles(
830 &call,
831 {ValueType::Shadow, ValueType::Primal, ValueType::Primal,
832 ValueType::Primal, ValueType::Primal, ValueType::Primal,
833 ValueType::None},
834 Builder2, /*lookup*/ !forwardMode);
835
836 Type *types[sizeof(args) / sizeof(*args)];
837 for (size_t i = 0; i < sizeof(args) / sizeof(*args); i++)
838 types[i] = args[i]->getType();
839 FunctionType *FT = FunctionType::get(call.getType(), types, false);
840
841 auto fcall = Builder2.CreateCall(
842 called->getParent()->getOrInsertFunction(
843 getRenamedPerCallingConv(called->getName(), "MPI_Send"), FT),
844 args, Defs);
845 fcall->setCallingConv(call.getCallingConv());
846
847 auto dst_arg =
848 Builder2.CreateBitCast(args[0], getInt8PtrTy(call.getContext()));
849 auto val_arg = ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
850 auto len_arg = Builder2.CreateZExtOrTrunc(
851 args[1], Type::getInt64Ty(call.getContext()));
852 auto tysize = MPI_TYPE_SIZE(datatype, Builder2, call.getType(), called);
853 len_arg =
854 Builder2.CreateMul(len_arg,
855 Builder2.CreateZExtOrTrunc(
856 tysize, Type::getInt64Ty(call.getContext())),
857 "", true, true);
858 auto volatile_arg = ConstantInt::getFalse(call.getContext());
859
860 Value *nargs[] = {dst_arg, val_arg, len_arg, volatile_arg};
861 Type *tys[] = {dst_arg->getType(), len_arg->getType()};
862
863 auto MemsetDefs = gutils->getInvertedBundles(
864 &call,
865 {ValueType::Shadow, ValueType::None, ValueType::None, ValueType::None,
866 ValueType::None, ValueType::None, ValueType::None},
867 Builder2, /*lookup*/ true);
868 auto memset = cast<CallInst>(Builder2.CreateCall(
869 getIntrinsicDeclaration(gutils->newFunc->getParent(),
870 Intrinsic::memset, tys),
871 nargs));
872 memset->addParamAttr(0, Attribute::NonNull);
873 }
874 if (Mode == DerivativeMode::ReverseModeGradient)
875 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
876 return;
877 }
878
879 // int MPI_Bcast( void *buffer, int count, MPI_Datatype datatype, int root,
880 // MPI_Comm comm )
881 // 1. if root, malloc intermediate buffer
882 // 2. reduce sum diff(buffer) into intermediate
883 // 3. if root, set shadow(buffer) = intermediate [memcpy] then free
884 // 3-e. else, set shadow(buffer) = 0 [memset]
885 if (funcName == "MPI_Bcast" || funcName == "PMPI_Bcast") {
886 if (Mode == DerivativeMode::ReverseModeGradient ||
887 Mode == DerivativeMode::ReverseModeCombined ||
888 Mode == DerivativeMode::ForwardMode ||
889 Mode == DerivativeMode::ForwardModeError) {
890 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
891 Mode == DerivativeMode::ForwardModeError;
892
893 IRBuilder<> Builder2 =
894 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
895 if (forwardMode) {
896 getForwardBuilder(Builder2);
897 } else {
898 getReverseBuilder(Builder2);
899 }
900
901 Value *shadow = gutils->invertPointerM(call.getOperand(0), Builder2);
902 if (!forwardMode)
903 shadow = lookup(shadow, Builder2);
904 if (shadow->getType()->isIntegerTy())
905 shadow =
906 Builder2.CreateIntToPtr(shadow, getInt8PtrTy(call.getContext()));
907
908 ConcreteType CT =
909 TR.firstPointer(1, call.getOperand(0), &call, gutils, &Builder2);
910 auto MPI_OP_type = getInt8PtrTy(call.getContext());
911 Type *MPI_OP_Ptr_type = getUnqual(MPI_OP_type);
912
913 Value *count = gutils->getNewFromOriginal(call.getOperand(1));
914 if (!forwardMode)
915 count = lookup(count, Builder2);
916 Value *datatype = gutils->getNewFromOriginal(call.getOperand(2));
917 if (!forwardMode)
918 datatype = lookup(datatype, Builder2);
919 Value *root = gutils->getNewFromOriginal(call.getOperand(3));
920 if (!forwardMode)
921 root = lookup(root, Builder2);
922
923 Value *comm = gutils->getNewFromOriginal(call.getOperand(4));
924 if (!forwardMode)
925 comm = lookup(comm, Builder2);
926
927 if (forwardMode) {
928 createMPIForwardCall(call, {/*buffer*/ 0}, gutils, Builder2);
929 return;
930 }
931
932 Value *rank = MPI_COMM_RANK(comm, Builder2, root->getType(), called);
933 Value *tysize = MPI_TYPE_SIZE(datatype, Builder2, call.getType(), called);
934
935 auto len_arg = Builder2.CreateZExtOrTrunc(
936 count, Type::getInt64Ty(call.getContext()));
937 len_arg =
938 Builder2.CreateMul(len_arg,
939 Builder2.CreateZExtOrTrunc(
940 tysize, Type::getInt64Ty(call.getContext())),
941 "", true, true);
942
943 // 1. if root, malloc intermediate buffer, else undef
944 PHINode *buf;
945
946 {
947 BasicBlock *currentBlock = Builder2.GetInsertBlock();
948 BasicBlock *rootBlock = gutils->addReverseBlock(
949 currentBlock, currentBlock->getName() + "_root", gutils->newFunc);
950 BasicBlock *mergeBlock = gutils->addReverseBlock(
951 rootBlock, currentBlock->getName() + "_post", gutils->newFunc);
952
953 Builder2.CreateCondBr(Builder2.CreateICmpEQ(rank, root), rootBlock,
954 mergeBlock);
955
956 Builder2.SetInsertPoint(rootBlock);
957
958 Value *rootbuf =
959 CreateAllocation(Builder2, Type::getInt8Ty(call.getContext()),
960 len_arg, "mpireduce_malloccache");
961 Builder2.CreateBr(mergeBlock);
962
963 Builder2.SetInsertPoint(mergeBlock);
964
965 buf = Builder2.CreatePHI(rootbuf->getType(), 2);
966 buf->addIncoming(rootbuf, rootBlock);
967 buf->addIncoming(UndefValue::get(buf->getType()), currentBlock);
968 }
969
970 // Need to preserve the shadow buffer.
971 auto BufferDefs = gutils->getInvertedBundles(
972 &call,
973 {ValueType::Shadow, ValueType::Primal, ValueType::Primal,
974 ValueType::Primal, ValueType::Primal},
975 Builder2, /*lookup*/ true);
976
977 // 2. reduce sum diff(buffer) into intermediate
978 {
979 // int MPI_Reduce(const void *sendbuf, void *recvbuf, int count,
980 // MPI_Datatype datatype,
981 // MPI_Op op, int root, MPI_Comm comm)
982 Value *args[] = {
983 /*sendbuf*/ shadow,
984 /*recvbuf*/ buf,
985 /*count*/ count,
986 /*datatype*/ datatype,
987 /*op (MPI_SUM)*/
988 getOrInsertOpFloatSum(*gutils->newFunc->getParent(), called,
989 MPI_OP_Ptr_type, MPI_OP_type, CT,
990 root->getType(), Builder2),
991 /*int root*/ root,
992 /*comm*/ comm,
993 };
994 Type *types[sizeof(args) / sizeof(*args)];
995 for (size_t i = 0; i < sizeof(args) / sizeof(*args); i++)
996 types[i] = args[i]->getType();
997
998 FunctionType *FT = FunctionType::get(call.getType(), types, false);
999
1000 Builder2.CreateCall(
1001 called->getParent()->getOrInsertFunction(
1002 getRenamedPerCallingConv(called->getName(), "MPI_Reduce"), FT),
1003 args, BufferDefs);
1004 }
1005
1006 // 3. if root, set shadow(buffer) = intermediate [memcpy]
1007 BasicBlock *currentBlock = Builder2.GetInsertBlock();
1008 BasicBlock *rootBlock = gutils->addReverseBlock(
1009 currentBlock, currentBlock->getName() + "_root", gutils->newFunc);
1010 BasicBlock *nonrootBlock = gutils->addReverseBlock(
1011 rootBlock, currentBlock->getName() + "_nonroot", gutils->newFunc);
1012 BasicBlock *mergeBlock = gutils->addReverseBlock(
1013 nonrootBlock, currentBlock->getName() + "_post", gutils->newFunc);
1014
1015 Builder2.CreateCondBr(Builder2.CreateICmpEQ(rank, root), rootBlock,
1016 nonrootBlock);
1017
1018 Builder2.SetInsertPoint(rootBlock);
1019
1020 {
1021 auto volatile_arg = ConstantInt::getFalse(call.getContext());
1022 Value *nargs[] = {shadow, buf, len_arg, volatile_arg};
1023
1024 Type *tys[] = {shadow->getType(), buf->getType(), len_arg->getType()};
1025
1026 auto memcpyF = getIntrinsicDeclaration(gutils->newFunc->getParent(),
1027 Intrinsic::memcpy, tys);
1028
1029 auto mem =
1030 cast<CallInst>(Builder2.CreateCall(memcpyF, nargs, BufferDefs));
1031 mem->setCallingConv(memcpyF->getCallingConv());
1032
1033 // Free up the memory of the buffer
1034 if (shouldFree()) {
1035 CreateDealloc(Builder2, buf);
1036 }
1037 }
1038
1039 Builder2.CreateBr(mergeBlock);
1040
1041 Builder2.SetInsertPoint(nonrootBlock);
1042
1043 // 3-e. else, set shadow(buffer) = 0 [memset]
1044 auto val_arg = ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
1045 auto volatile_arg = ConstantInt::getFalse(call.getContext());
1046 Value *args[] = {shadow, val_arg, len_arg, volatile_arg};
1047 Type *tys[] = {args[0]->getType(), args[2]->getType()};
1048 auto memset = cast<CallInst>(Builder2.CreateCall(
1049 getIntrinsicDeclaration(gutils->newFunc->getParent(),
1050 Intrinsic::memset, tys),
1051 args, BufferDefs));
1052 memset->addParamAttr(0, Attribute::NonNull);
1053 Builder2.CreateBr(mergeBlock);
1054
1055 Builder2.SetInsertPoint(mergeBlock);
1056 }
1057 if (Mode == DerivativeMode::ReverseModeGradient)
1058 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
1059 return;
1060 }
1061
1062 // Approximate algo (for sum): -> if statement yet to be
1063 // 1. malloc intermediate buffer
1064 // 1.5 if root, set intermediate = diff(recvbuffer)
1065 // 2. MPI_Bcast intermediate to all
1066 // 3. if root, Zero diff(recvbuffer) [memset to 0]
1067 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
1068 // 5. free intermediate buffer
1069
1070 // int MPI_Reduce(const void *sendbuf, void *recvbuf, int count,
1071 // MPI_Datatype datatype,
1072 // MPI_Op op, int root, MPI_Comm comm)
1073
1074 llvm::StringRef canonMPIName = canonicalizeMPIName(funcName);
1075
1076 if (canonMPIName == "MPI_Reduce") {
1077 if (Mode == DerivativeMode::ReverseModeGradient ||
1078 Mode == DerivativeMode::ReverseModeCombined ||
1079 Mode == DerivativeMode::ForwardMode ||
1080 Mode == DerivativeMode::ForwardModeError) {
1081 // TODO insert a check for sum
1082
1083 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
1084 Mode == DerivativeMode::ForwardModeError;
1085
1086 IRBuilder<> Builder2 =
1087 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
1088 if (forwardMode) {
1089 getForwardBuilder(Builder2);
1090 } else {
1091 getReverseBuilder(Builder2);
1092 }
1093
1094 // Get the operations from MPI_Reduce
1095 Value *orig_sendbuf = call.getOperand(0);
1096 Value *orig_recvbuf = call.getOperand(1);
1097 Value *orig_count = call.getOperand(2);
1098 Value *orig_datatype = call.getOperand(3);
1099 Value *orig_op = call.getOperand(4);
1100 Value *orig_root = call.getOperand(5);
1101 Value *orig_comm = call.getOperand(6);
1102
1103 // The Fortran MPI ABI ("mpi_reduce_", "mpi_reduce__", ...) passes all
1104 // arguments by reference and appends an `ierr` argument.
1105 bool fortranABI = isFortranMPICall(called->getName());
1106
1107 bool isSum = false;
1108 if (Constant *C = dyn_cast<Constant>(orig_op)) {
1109 while (ConstantExpr *CE = dyn_cast<ConstantExpr>(C)) {
1110 C = CE->getOperand(0);
1111 }
1112 if (auto GV = dyn_cast<GlobalVariable>(C)) {
1113 if (GV->getName() == "ompi_mpi_op_sum") {
1114 isSum = true;
1115 } else if (fortranABI && GV->isConstant() &&
1116 GV->hasDefinitiveInitializer()) {
1117 // The Fortran ABI passes the operator as a reference to an
1118 // integer handle
1119 if (auto *CI = dyn_cast<ConstantInt>(GV->getInitializer())) {
1120 // MPICH native ABI (also covers the MPICH ABI Compatibility
1121 // Initiative: Intel MPI, MVAPICH, Cray MPICH)
1122 if (CI->getValue() == 1476395011) {
1123 isSum = true;
1124 }
1125 // MPI 5.0 standard ABI (Chapter 20), where predefined op
1126 // handles are fixed compile-time constants, identical in C
1127 // and Fortran: MPI_SUM == 33.
1128 if (CI->getValue() == 33) {
1129 isSum = true;
1130 }
1131 }
1132 }
1133 }
1134 // MPICH native ABI
1135 if (ConstantInt *CI = dyn_cast<ConstantInt>(C)) {
1136 if (CI->getValue() == 1476395011) {
1137 isSum = true;
1138 }
1139 }
1140 // MPI 5.0 standard ABI: predefined op handles are small-integer
1141 // pointer constants (inttoptr), with MPI_SUM == 33.
1142 if (ConstantInt *CI = dyn_cast<ConstantInt>(C)) {
1143 if (CI->getValue() == 33) {
1144 isSum = true;
1145 }
1146 }
1147 }
1148 if (!isSum) {
1149 if (fortranABI) {
1150 // Integer MPI operator handles of a native Fortran ABI are
1151 // implementation-defined and cannot be mapped portably at compile
1152 // time; warn and assume the common case of MPI_SUM. This remains
1153 // true with MPI 5.0: although its standard ABI (Chapter 20) does
1154 // fix handle values as portable compile-time constants (MPI_SUM
1155 // == 33, recognized above), supporting that ABI is optional and
1156 // neither Open MPI nor MPICH use it by default, so code compiled
1157 // against the default/native ABIs still carries
1158 // implementation-defined handle values.
1159 llvm::errs() << "warning: cannot determine MPI op used in `" << call
1160 << "`, assuming MPI_SUM\n";
1161 } else {
1162 std::string s;
1163 llvm::raw_string_ostream ss(s);
1164 ss << " call: " << call << "\n";
1165 ss << " unhandled mpi_reduce op: " << *orig_op << "\n";
1166 EmitNoDerivativeError(ss.str(), call, gutils, BuilderZ);
1167 return;
1168 }
1169 }
1170
1171 Value *shadow_recvbuf = gutils->invertPointerM(orig_recvbuf, Builder2);
1172 if (!forwardMode)
1173 shadow_recvbuf = lookup(shadow_recvbuf, Builder2);
1174 if (shadow_recvbuf->getType()->isIntegerTy())
1175 shadow_recvbuf = Builder2.CreateIntToPtr(
1176 shadow_recvbuf, getInt8PtrTy(call.getContext()));
1177
1178 Value *shadow_sendbuf = gutils->invertPointerM(orig_sendbuf, Builder2);
1179 if (!forwardMode)
1180 shadow_sendbuf = lookup(shadow_sendbuf, Builder2);
1181 if (shadow_sendbuf->getType()->isIntegerTy())
1182 shadow_sendbuf = Builder2.CreateIntToPtr(
1183 shadow_sendbuf, getInt8PtrTy(call.getContext()));
1184
1185 // Need to preserve the shadow send/recv buffers. The Fortran ABI call
1186 // has an extra `ierr` argument, so build the bundle to match the actual
1187 // call arity: shadow send/recv buffers, primal everything else.
1188 std::vector<ValueType> BufferBundleTypes(call.arg_size(),
1189 ValueType::Primal);
1190 BufferBundleTypes[0] = ValueType::Shadow;
1191 BufferBundleTypes[1] = ValueType::Shadow;
1192 auto BufferDefs =
1193 gutils->getInvertedBundles(&call, BufferBundleTypes, Builder2,
1194 /*lookup*/ !forwardMode);
1195
1196 Value *count = gutils->getNewFromOriginal(orig_count);
1197 if (!forwardMode)
1198 count = lookup(count, Builder2);
1199
1200 Value *datatype = gutils->getNewFromOriginal(orig_datatype);
1201 if (!forwardMode)
1202 datatype = lookup(datatype, Builder2);
1203
1204 Value *op = gutils->getNewFromOriginal(orig_op);
1205 if (!forwardMode)
1206 op = lookup(op, Builder2);
1207
1208 Value *root = gutils->getNewFromOriginal(orig_root);
1209 if (!forwardMode)
1210 root = lookup(root, Builder2);
1211
1212 Value *comm = gutils->getNewFromOriginal(orig_comm);
1213 if (!forwardMode)
1214 comm = lookup(comm, Builder2);
1215
1216 // The Fortran ABI passes the count and root arguments by reference;
1217 // load them where their value is needed.
1218 Type *i32Ty = Type::getInt32Ty(call.getContext());
1219 Value *countVal = fortranABI ? Builder2.CreateLoad(i32Ty, count) : count;
1220 Value *rootVal = fortranABI ? Builder2.CreateLoad(i32Ty, root) : root;
1221
1222 Value *rank = MPI_COMM_RANK(comm, Builder2, rootVal->getType(), called);
1223
1224 if (forwardMode) {
1225 createMPIForwardCall(call, {/*sendbuf*/ 0, /*recvbuf*/ 1}, gutils,
1226 Builder2);
1227 return;
1228 }
1229
1230 Value *tysize = MPI_TYPE_SIZE(datatype, Builder2, call.getType(), called);
1231
1232 // Get the length for the allocation of the intermediate buffer
1233 auto len_arg = Builder2.CreateZExtOrTrunc(
1234 countVal, Type::getInt64Ty(call.getContext()));
1235 len_arg =
1236 Builder2.CreateMul(len_arg,
1237 Builder2.CreateZExtOrTrunc(
1238 tysize, Type::getInt64Ty(call.getContext())),
1239 "", true, true);
1240
1241 // 1. Alloc intermediate buffer
1242 Value *buf =
1243 CreateAllocation(Builder2, Type::getInt8Ty(call.getContext()),
1244 len_arg, "mpireduce_malloccache");
1245
1246 // 1.5 if root, set intermediate = diff(recvbuffer)
1247 {
1248
1249 BasicBlock *currentBlock = Builder2.GetInsertBlock();
1250 BasicBlock *rootBlock = gutils->addReverseBlock(
1251 currentBlock, currentBlock->getName() + "_root", gutils->newFunc);
1252 BasicBlock *mergeBlock = gutils->addReverseBlock(
1253 rootBlock, currentBlock->getName() + "_post", gutils->newFunc);
1254
1255 Builder2.CreateCondBr(Builder2.CreateICmpEQ(rank, rootVal), rootBlock,
1256 mergeBlock);
1257
1258 Builder2.SetInsertPoint(rootBlock);
1259
1260 {
1261 auto volatile_arg = ConstantInt::getFalse(call.getContext());
1262 Value *nargs[] = {buf, shadow_recvbuf, len_arg, volatile_arg};
1263
1264 Type *tys[] = {nargs[0]->getType(), nargs[1]->getType(),
1265 len_arg->getType()};
1266
1267 auto memcpyF = getIntrinsicDeclaration(gutils->newFunc->getParent(),
1268 Intrinsic::memcpy, tys);
1269
1270 auto mem =
1271 cast<CallInst>(Builder2.CreateCall(memcpyF, nargs, BufferDefs));
1272 mem->setCallingConv(memcpyF->getCallingConv());
1273 }
1274
1275 Builder2.CreateBr(mergeBlock);
1276 Builder2.SetInsertPoint(mergeBlock);
1277 }
1278
1279 // 2. MPI_Bcast intermediate to all
1280 {
1281 // int MPI_Bcast( void *buffer, int count, MPI_Datatype datatype, int
1282 // root,
1283 // MPI_Comm comm )
1284 // The Fortran ABI passes count/datatype/root/comm by reference (as
1285 // held here) and appends an `ierr` argument.
1286 SmallVector<Value *, 8> args = {
1287 /*buf*/ buf,
1288 /*count*/ count,
1289 /*datatype*/ datatype,
1290 /*int root*/ root,
1291 /*comm*/ comm,
1292 };
1293 if (fortranABI) {
1294 args.push_back(IRBuilder<>(gutils->inversionAllocs)
1295 .CreateAlloca(i32Ty, nullptr, "enzyme_mpi_ierr"));
1296 }
1297 SmallVector<Type *, 8> types;
1298 for (auto *arg : args) {
1299 types.push_back(arg->getType());
1300 }
1301
1302 FunctionType *FT = FunctionType::get(
1303 fortranABI ? Type::getVoidTy(call.getContext()) : call.getType(),
1304 types, false);
1305 Builder2.CreateCall(
1306 called->getParent()->getOrInsertFunction(
1307 getRenamedPerCallingConv(called->getName(), "MPI_Bcast"), FT),
1308 args, BufferDefs);
1309 }
1310
1311 // 3. if root, Zero diff(recvbuffer) [memset to 0]
1312 {
1313 BasicBlock *currentBlock = Builder2.GetInsertBlock();
1314 BasicBlock *rootBlock = gutils->addReverseBlock(
1315 currentBlock, currentBlock->getName() + "_root", gutils->newFunc);
1316 BasicBlock *mergeBlock = gutils->addReverseBlock(
1317 rootBlock, currentBlock->getName() + "_post", gutils->newFunc);
1318
1319 Builder2.CreateCondBr(Builder2.CreateICmpEQ(rank, rootVal), rootBlock,
1320 mergeBlock);
1321
1322 Builder2.SetInsertPoint(rootBlock);
1323
1324 auto val_arg = ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
1325 auto volatile_arg = ConstantInt::getFalse(call.getContext());
1326 Value *args[] = {shadow_recvbuf, val_arg, len_arg, volatile_arg};
1327 Type *tys[] = {args[0]->getType(), args[2]->getType()};
1328 auto memset = cast<CallInst>(Builder2.CreateCall(
1329 getIntrinsicDeclaration(gutils->newFunc->getParent(),
1330 Intrinsic::memset, tys),
1331 args, BufferDefs));
1332 memset->addParamAttr(0, Attribute::NonNull);
1333
1334 Builder2.CreateBr(mergeBlock);
1335 Builder2.SetInsertPoint(mergeBlock);
1336 }
1337
1338 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
1339 DifferentiableMemCopyFloats(call, orig_sendbuf, buf, shadow_sendbuf,
1340 len_arg, Builder2, BufferDefs);
1341
1342 // Free up intermediate buffer
1343 if (shouldFree()) {
1344 CreateDealloc(Builder2, buf);
1345 }
1346 }
1347 if (Mode == DerivativeMode::ReverseModeGradient)
1348 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
1349 return;
1350 }
1351
1352 // Approximate algo (for sum): -> if statement yet to be
1353 // 1. malloc intermediate buffers
1354 // 2. MPI_Allreduce (sum) of diff(recvbuffer) to intermediate
1355 // 3. Zero diff(recvbuffer) [memset to 0]
1356 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
1357 // 5. free intermediate buffer
1358
1359 // int MPI_Allreduce(const void *sendbuf, void *recvbuf, int count,
1360 // MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)
1361
1362 if (canonMPIName == "MPI_Allreduce") {
1363 if (Mode == DerivativeMode::ReverseModeGradient ||
1364 Mode == DerivativeMode::ReverseModeCombined ||
1365 Mode == DerivativeMode::ForwardMode ||
1366 Mode == DerivativeMode::ForwardModeError) {
1367 // TODO insert a check for sum
1368
1369 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
1370 Mode == DerivativeMode::ForwardModeError;
1371
1372 IRBuilder<> Builder2 =
1373 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
1374 if (forwardMode) {
1375 getForwardBuilder(Builder2);
1376 } else {
1377 getReverseBuilder(Builder2);
1378 }
1379
1380 // Get the operations from MPI_Allreduce
1381 Value *orig_sendbuf = call.getOperand(0);
1382 Value *orig_recvbuf = call.getOperand(1);
1383 Value *orig_count = call.getOperand(2);
1384 Value *orig_datatype = call.getOperand(3);
1385 Value *orig_op = call.getOperand(4);
1386 Value *orig_comm = call.getOperand(5);
1387
1388 // The Fortran MPI ABI ("mpi_allreduce_", "mpi_allreduce__", ...) passes
1389 // all arguments by reference and appends an `ierr` argument.
1390 bool fortranABI = isFortranMPICall(called->getName());
1391
1392 bool isSum = false;
1393 if (Constant *C = dyn_cast<Constant>(orig_op)) {
1394 while (ConstantExpr *CE = dyn_cast<ConstantExpr>(C)) {
1395 C = CE->getOperand(0);
1396 }
1397 if (auto GV = dyn_cast<GlobalVariable>(C)) {
1398 if (GV->getName() == "ompi_mpi_op_sum") {
1399 isSum = true;
1400 } else if (fortranABI && GV->isConstant() &&
1401 GV->hasDefinitiveInitializer()) {
1402 // The Fortran ABI passes the operator as a reference to an
1403 // integer handle
1404 if (auto *CI = dyn_cast<ConstantInt>(GV->getInitializer())) {
1405 // MPICH native ABI (also covers the MPICH ABI Compatibility
1406 // Initiative: Intel MPI, MVAPICH, Cray MPICH)
1407 if (CI->getValue() == 1476395011) {
1408 isSum = true;
1409 }
1410 // MPI 5.0 standard ABI (Chapter 20), where predefined op
1411 // handles are fixed compile-time constants, identical in C
1412 // and Fortran: MPI_SUM == 33.
1413 if (CI->getValue() == 33) {
1414 isSum = true;
1415 }
1416 }
1417 }
1418 }
1419 // MPICH native ABI
1420 if (ConstantInt *CI = dyn_cast<ConstantInt>(C)) {
1421 if (CI->getValue() == 1476395011) {
1422 isSum = true;
1423 }
1424 }
1425 // MPI 5.0 standard ABI: predefined op handles are small-integer
1426 // pointer constants (inttoptr), with MPI_SUM == 33.
1427 if (ConstantInt *CI = dyn_cast<ConstantInt>(C)) {
1428 if (CI->getValue() == 33) {
1429 isSum = true;
1430 }
1431 }
1432 }
1433 if (!isSum) {
1434 if (fortranABI) {
1435 // Integer MPI operator handles of a native Fortran ABI are
1436 // implementation-defined and cannot be mapped portably at compile
1437 // time; warn and assume the common case of MPI_SUM. This remains
1438 // true with MPI 5.0: although its standard ABI (Chapter 20) does
1439 // fix handle values as portable compile-time constants (MPI_SUM
1440 // == 33, recognized above), supporting that ABI is optional and
1441 // neither Open MPI nor MPICH use it by default, so code compiled
1442 // against the default/native ABIs still carries
1443 // implementation-defined handle values.
1444 llvm::errs() << "warning: cannot determine MPI op used in `" << call
1445 << "`, assuming MPI_SUM\n";
1446 } else {
1447 std::string s;
1448 llvm::raw_string_ostream ss(s);
1449 ss << " call: " << call << "\n";
1450 ss << " unhandled mpi_allreduce op: " << *orig_op << "\n";
1451 EmitNoDerivativeError(ss.str(), call, gutils, BuilderZ);
1452 return;
1453 }
1454 }
1455
1456 Value *shadow_recvbuf = gutils->invertPointerM(orig_recvbuf, Builder2);
1457 if (!forwardMode)
1458 shadow_recvbuf = lookup(shadow_recvbuf, Builder2);
1459 if (shadow_recvbuf->getType()->isIntegerTy())
1460 shadow_recvbuf = Builder2.CreateIntToPtr(
1461 shadow_recvbuf, getInt8PtrTy(call.getContext()));
1462
1463 Value *shadow_sendbuf = gutils->invertPointerM(orig_sendbuf, Builder2);
1464 if (!forwardMode)
1465 shadow_sendbuf = lookup(shadow_sendbuf, Builder2);
1466 if (shadow_sendbuf->getType()->isIntegerTy())
1467 shadow_sendbuf = Builder2.CreateIntToPtr(
1468 shadow_sendbuf, getInt8PtrTy(call.getContext()));
1469
1470 // Need to preserve the shadow send/recv buffers. The Fortran ABI call
1471 // has an extra `ierr` argument, so build the bundle to match the actual
1472 // call arity: shadow send/recv buffers, primal everything else.
1473 std::vector<ValueType> BufferBundleTypes(call.arg_size(),
1474 ValueType::Primal);
1475 BufferBundleTypes[0] = ValueType::Shadow;
1476 BufferBundleTypes[1] = ValueType::Shadow;
1477 auto BufferDefs =
1478 gutils->getInvertedBundles(&call, BufferBundleTypes, Builder2,
1479 /*lookup*/ !forwardMode);
1480
1481 Value *count = gutils->getNewFromOriginal(orig_count);
1482 if (!forwardMode)
1483 count = lookup(count, Builder2);
1484
1485 Value *datatype = gutils->getNewFromOriginal(orig_datatype);
1486 if (!forwardMode)
1487 datatype = lookup(datatype, Builder2);
1488
1489 Value *comm = gutils->getNewFromOriginal(orig_comm);
1490 if (!forwardMode)
1491 comm = lookup(comm, Builder2);
1492
1493 Value *op = gutils->getNewFromOriginal(orig_op);
1494 if (!forwardMode)
1495 op = lookup(op, Builder2);
1496
1497 // The Fortran ABI passes the count argument by reference; load it
1498 // where its value is needed.
1499 Type *i32Ty = Type::getInt32Ty(call.getContext());
1500 Value *countVal = fortranABI ? Builder2.CreateLoad(i32Ty, count) : count;
1501
1502 if (forwardMode) {
1503 createMPIForwardCall(call, {/*sendbuf*/ 0, /*recvbuf*/ 1}, gutils,
1504 Builder2);
1505 return;
1506 }
1507
1508 Value *tysize = MPI_TYPE_SIZE(datatype, Builder2, call.getType(), called);
1509
1510 // Get the length for the allocation of the intermediate buffer
1511 auto len_arg = Builder2.CreateZExtOrTrunc(
1512 countVal, Type::getInt64Ty(call.getContext()));
1513 len_arg =
1514 Builder2.CreateMul(len_arg,
1515 Builder2.CreateZExtOrTrunc(
1516 tysize, Type::getInt64Ty(call.getContext())),
1517 "", true, true);
1518
1519 // 1. Alloc intermediate buffer
1520 Value *buf =
1521 CreateAllocation(Builder2, Type::getInt8Ty(call.getContext()),
1522 len_arg, "mpireduce_malloccache");
1523
1524 // 2. MPI_Allreduce (sum) of diff(recvbuffer) to intermediate
1525 {
1526 // int MPI_Allreduce(const void *sendbuf, void *recvbuf, int count,
1527 // MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)
1528 // The Fortran ABI passes count/datatype/op/comm by reference (as
1529 // held here) and appends an `ierr` argument.
1530 SmallVector<Value *, 8> args = {
1531 /*sendbuf*/ shadow_recvbuf,
1532 /*recvbuf*/ buf,
1533 /*count*/ count,
1534 /*datatype*/ datatype,
1535 /*op*/ op,
1536 /*comm*/ comm,
1537 };
1538 if (fortranABI) {
1539 args.push_back(IRBuilder<>(gutils->inversionAllocs)
1540 .CreateAlloca(i32Ty, nullptr, "enzyme_mpi_ierr"));
1541 }
1542 SmallVector<Type *, 8> types;
1543 for (auto *arg : args)
1544 types.push_back(arg->getType());
1545
1546 FunctionType *FT = FunctionType::get(
1547 fortranABI ? Type::getVoidTy(call.getContext()) : call.getType(),
1548 types, false);
1549 Builder2.CreateCall(
1550 called->getParent()->getOrInsertFunction(
1551 getRenamedPerCallingConv(called->getName(), "MPI_Allreduce"),
1552 FT),
1553 args, BufferDefs);
1554 }
1555
1556 // 3. Zero diff(recvbuffer) [memset to 0]
1557 auto val_arg = ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
1558 auto volatile_arg = ConstantInt::getFalse(call.getContext());
1559 Value *args[] = {shadow_recvbuf, val_arg, len_arg, volatile_arg};
1560 Type *tys[] = {args[0]->getType(), args[2]->getType()};
1561 auto memset = cast<CallInst>(Builder2.CreateCall(
1562 getIntrinsicDeclaration(gutils->newFunc->getParent(),
1563 Intrinsic::memset, tys),
1564 args, BufferDefs));
1565 memset->addParamAttr(0, Attribute::NonNull);
1566
1567 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
1568 DifferentiableMemCopyFloats(call, orig_sendbuf, buf, shadow_sendbuf,
1569 len_arg, Builder2, BufferDefs);
1570
1571 // Free up intermediate buffer
1572 if (shouldFree()) {
1573 CreateDealloc(Builder2, buf);
1574 }
1575 }
1576 if (Mode == DerivativeMode::ReverseModeGradient)
1577 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
1578 return;
1579 }
1580
1581 // Approximate algo (for sum): -> if statement yet to be
1582 // 1. malloc intermediate buffer
1583 // 2. Scatter diff(recvbuffer) to intermediate buffer
1584 // 3. if root, Zero diff(recvbuffer) [memset to 0]
1585 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
1586 // 5. free intermediate buffer
1587
1588 // int MPI_Gather(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
1589 // void *recvbuf, int recvcount, MPI_Datatype recvtype,
1590 // int root, MPI_Comm comm)
1591
1592 if (funcName == "MPI_Gather" || funcName == "PMPI_Gather") {
1593 if (Mode == DerivativeMode::ReverseModeGradient ||
1594 Mode == DerivativeMode::ReverseModeCombined ||
1595 Mode == DerivativeMode::ForwardMode ||
1596 Mode == DerivativeMode::ForwardModeError) {
1597 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
1598 Mode == DerivativeMode::ForwardModeError;
1599
1600 IRBuilder<> Builder2 =
1601 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
1602 if (forwardMode) {
1603 getForwardBuilder(Builder2);
1604 } else {
1605 getReverseBuilder(Builder2);
1606 }
1607
1608 Value *orig_sendbuf = call.getOperand(0);
1609 Value *orig_sendcount = call.getOperand(1);
1610 Value *orig_sendtype = call.getOperand(2);
1611 Value *orig_recvbuf = call.getOperand(3);
1612 Value *orig_recvcount = call.getOperand(4);
1613 Value *orig_recvtype = call.getOperand(5);
1614 Value *orig_root = call.getOperand(6);
1615 Value *orig_comm = call.getOperand(7);
1616
1617 Value *shadow_recvbuf = gutils->invertPointerM(orig_recvbuf, Builder2);
1618 if (!forwardMode)
1619 shadow_recvbuf = lookup(shadow_recvbuf, Builder2);
1620 if (shadow_recvbuf->getType()->isIntegerTy())
1621 shadow_recvbuf = Builder2.CreateIntToPtr(
1622 shadow_recvbuf, getInt8PtrTy(call.getContext()));
1623
1624 Value *shadow_sendbuf = gutils->invertPointerM(orig_sendbuf, Builder2);
1625 if (!forwardMode)
1626 shadow_sendbuf = lookup(shadow_sendbuf, Builder2);
1627 if (shadow_sendbuf->getType()->isIntegerTy())
1628 shadow_sendbuf = Builder2.CreateIntToPtr(
1629 shadow_sendbuf, getInt8PtrTy(call.getContext()));
1630
1631 Value *recvcount = gutils->getNewFromOriginal(orig_recvcount);
1632 if (!forwardMode)
1633 recvcount = lookup(recvcount, Builder2);
1634
1635 Value *recvtype = gutils->getNewFromOriginal(orig_recvtype);
1636 if (!forwardMode)
1637 recvtype = lookup(recvtype, Builder2);
1638
1639 Value *sendcount = gutils->getNewFromOriginal(orig_sendcount);
1640 if (!sendcount)
1641 sendcount = lookup(sendcount, Builder2);
1642
1643 Value *sendtype = gutils->getNewFromOriginal(orig_sendtype);
1644 if (!forwardMode)
1645 sendtype = lookup(sendtype, Builder2);
1646
1647 bool fortranABI = isFortranMPICall(called->getName());
1648
1649 Value *root = gutils->getNewFromOriginal(orig_root);
1650 if (!forwardMode)
1651 root = lookup(root, Builder2);
1652
1653 Value *comm = gutils->getNewFromOriginal(orig_comm);
1654 if (!forwardMode)
1655 comm = lookup(comm, Builder2);
1656
1657 Type *i32Ty = Type::getInt32Ty(call.getContext());
1658 Value *sendcountVal =
1659 fortranABI ? Builder2.CreateLoad(i32Ty, sendcount) : sendcount;
1660 Value *recvcountVal =
1661 fortranABI ? Builder2.CreateLoad(i32Ty, recvcount) : recvcount;
1662 Value *rootVal = fortranABI ? Builder2.CreateLoad(i32Ty, root) : root;
1663
1664 Value *rank = MPI_COMM_RANK(comm, Builder2, rootVal->getType(), called);
1665 Value *tysize = MPI_TYPE_SIZE(sendtype, Builder2, call.getType(), called);
1666
1667 if (forwardMode) {
1668 createMPIForwardCall(call, {/*sendbuf*/ 0, /*recvbuf*/ 3}, gutils,
1669 Builder2);
1670 return;
1671 }
1672
1673 // Get the length for the allocation of the intermediate buffer
1674 auto sendlen_arg = Builder2.CreateZExtOrTrunc(
1675 sendcountVal, Type::getInt64Ty(call.getContext()));
1676 sendlen_arg =
1677 Builder2.CreateMul(sendlen_arg,
1678 Builder2.CreateZExtOrTrunc(
1679 tysize, Type::getInt64Ty(call.getContext())),
1680 "", true, true);
1681
1682 // Need to preserve the shadow send/recv buffers. The Fortran ABI call
1683 // has an extra `ierr` argument, so size the bundle to match the actual
1684 // call arity: shadow send/recv buffers, primal everything else.
1685 std::vector<ValueType> BufferBundleTypes(call.arg_size(),
1686 ValueType::Primal);
1687 BufferBundleTypes[0] = ValueType::Shadow;
1688 BufferBundleTypes[3] = ValueType::Shadow;
1689 auto BufferDefs = gutils->getInvertedBundles(&call, BufferBundleTypes,
1690 Builder2, /*lookup*/ true);
1691
1692 // 1. Alloc intermediate buffer
1693 Value *buf =
1694 CreateAllocation(Builder2, Type::getInt8Ty(call.getContext()),
1695 sendlen_arg, "mpireduce_malloccache");
1696
1697 // 2. Scatter diff(recvbuffer) to intermediate buffer
1698 {
1699 // int MPI_Scatter(const void *sendbuf, int sendcount, MPI_Datatype
1700 // sendtype,
1701 // void *recvbuf, int recvcount, MPI_Datatype recvtype, int root,
1702 // MPI_Comm comm)
1703 //
1704 // The Fortran MPI ABI passes all arguments by reference and appends
1705 // an `ierr` argument, so the generated call must match the convention
1706 // of the caller.
1707 SmallVector<Value *, 10> args;
1708 if (fortranABI) {
1709 args = {shadow_recvbuf, recvcount, recvtype, buf,
1710 sendcount, sendtype, root, comm};
1711 for (size_t i = 8, e = call.arg_size(); i < e; i++)
1712 args.push_back(lookup(
1713 gutils->getNewFromOriginal(call.getArgOperand(i)), Builder2));
1714 } else {
1715 args = {shadow_recvbuf, recvcountVal, recvtype, buf,
1716 sendcountVal, sendtype, rootVal, comm};
1717 }
1718 SmallVector<Type *, 10> types;
1719 for (auto *arg : args)
1720 types.push_back(arg->getType());
1721
1722 FunctionType *FT = FunctionType::get(call.getType(), types, false);
1723 Builder2.CreateCall(
1724 called->getParent()->getOrInsertFunction(
1725 getRenamedPerCallingConv(called->getName(), "MPI_Scatter"), FT),
1726 args, BufferDefs);
1727 }
1728
1729 // 3. if root, Zero diff(recvbuffer) [memset to 0]
1730 {
1731
1732 BasicBlock *currentBlock = Builder2.GetInsertBlock();
1733 BasicBlock *rootBlock = gutils->addReverseBlock(
1734 currentBlock, currentBlock->getName() + "_root", gutils->newFunc);
1735 BasicBlock *mergeBlock = gutils->addReverseBlock(
1736 rootBlock, currentBlock->getName() + "_post", gutils->newFunc);
1737
1738 Builder2.CreateCondBr(Builder2.CreateICmpEQ(rank, rootVal), rootBlock,
1739 mergeBlock);
1740
1741 Builder2.SetInsertPoint(rootBlock);
1742 auto recvlen_arg = Builder2.CreateZExtOrTrunc(
1743 recvcountVal, Type::getInt64Ty(call.getContext()));
1744 recvlen_arg =
1745 Builder2.CreateMul(recvlen_arg,
1746 Builder2.CreateZExtOrTrunc(
1747 tysize, Type::getInt64Ty(call.getContext())),
1748 "", true, true);
1749 recvlen_arg = Builder2.CreateMul(
1750 recvlen_arg,
1751 Builder2.CreateZExtOrTrunc(
1752 MPI_COMM_SIZE(comm, Builder2, rootVal->getType(), called),
1753 Type::getInt64Ty(call.getContext())),
1754 "", true, true);
1755
1756 auto val_arg = ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
1757 auto volatile_arg = ConstantInt::getFalse(call.getContext());
1758 Value *args[] = {shadow_recvbuf, val_arg, recvlen_arg, volatile_arg};
1759 Type *tys[] = {args[0]->getType(), args[2]->getType()};
1760 auto memset = cast<CallInst>(Builder2.CreateCall(
1761 getIntrinsicDeclaration(gutils->newFunc->getParent(),
1762 Intrinsic::memset, tys),
1763 args, BufferDefs));
1764 memset->addParamAttr(0, Attribute::NonNull);
1765
1766 Builder2.CreateBr(mergeBlock);
1767 Builder2.SetInsertPoint(mergeBlock);
1768 }
1769
1770 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
1771 DifferentiableMemCopyFloats(call, orig_sendbuf, buf, shadow_sendbuf,
1772 sendlen_arg, Builder2, BufferDefs);
1773
1774 // Free up intermediate buffer
1775 if (shouldFree()) {
1776 CreateDealloc(Builder2, buf);
1777 }
1778 }
1779 if (Mode == DerivativeMode::ReverseModeGradient)
1780 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
1781 return;
1782 }
1783
1784 // Approximate algo (for sum): -> if statement yet to be
1785 // 1. if root, malloc intermediate buffer, else undef
1786 // 2. Gather diff(recvbuffer) to intermediate buffer
1787 // 3. Zero diff(recvbuffer) [memset to 0]
1788 // 4. if root, diff(sendbuffer) += intermediate buffer (diffmemcopy)
1789 // 5. if root, free intermediate buffer
1790
1791 // int MPI_Scatter(const void *sendbuf, int sendcount, MPI_Datatype
1792 // sendtype,
1793 // void *recvbuf, int recvcount, MPI_Datatype recvtype, int root,
1794 // MPI_Comm comm)
1795 if (canonMPIName == "MPI_Scatter") {
1796 if (Mode == DerivativeMode::ReverseModeGradient ||
1797 Mode == DerivativeMode::ReverseModeCombined ||
1798 Mode == DerivativeMode::ForwardMode ||
1799 Mode == DerivativeMode::ForwardModeError) {
1800 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
1801 Mode == DerivativeMode::ForwardModeError;
1802
1803 IRBuilder<> Builder2 =
1804 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
1805 if (forwardMode) {
1806 getForwardBuilder(Builder2);
1807 } else {
1808 getReverseBuilder(Builder2);
1809 }
1810
1811 Value *orig_sendbuf = call.getOperand(0);
1812 Value *orig_sendcount = call.getOperand(1);
1813 Value *orig_sendtype = call.getOperand(2);
1814 Value *orig_recvbuf = call.getOperand(3);
1815 Value *orig_recvcount = call.getOperand(4);
1816 Value *orig_recvtype = call.getOperand(5);
1817 Value *orig_root = call.getOperand(6);
1818 Value *orig_comm = call.getOperand(7);
1819
1820 Value *shadow_recvbuf = gutils->invertPointerM(orig_recvbuf, Builder2);
1821 if (!forwardMode)
1822 shadow_recvbuf = lookup(shadow_recvbuf, Builder2);
1823 if (shadow_recvbuf->getType()->isIntegerTy())
1824 shadow_recvbuf = Builder2.CreateIntToPtr(
1825 shadow_recvbuf, getInt8PtrTy(call.getContext()));
1826
1827 Value *shadow_sendbuf = gutils->invertPointerM(orig_sendbuf, Builder2);
1828 if (!forwardMode)
1829 shadow_sendbuf = lookup(shadow_sendbuf, Builder2);
1830 if (shadow_sendbuf->getType()->isIntegerTy())
1831 shadow_sendbuf = Builder2.CreateIntToPtr(
1832 shadow_sendbuf, getInt8PtrTy(call.getContext()));
1833
1834 Value *recvcount = gutils->getNewFromOriginal(orig_recvcount);
1835 if (!forwardMode)
1836 recvcount = lookup(recvcount, Builder2);
1837
1838 Value *recvtype = gutils->getNewFromOriginal(orig_recvtype);
1839 if (!forwardMode)
1840 recvtype = lookup(recvtype, Builder2);
1841
1842 Value *sendcount = gutils->getNewFromOriginal(orig_sendcount);
1843 if (!forwardMode)
1844 sendcount = lookup(sendcount, Builder2);
1845
1846 Value *sendtype = gutils->getNewFromOriginal(orig_sendtype);
1847 if (!forwardMode)
1848 sendtype = lookup(sendtype, Builder2);
1849
1850 Value *root = gutils->getNewFromOriginal(orig_root);
1851 if (!forwardMode)
1852 root = lookup(root, Builder2);
1853
1854 Value *comm = gutils->getNewFromOriginal(orig_comm);
1855 if (!forwardMode)
1856 comm = lookup(comm, Builder2);
1857
1858 bool fortranABI = isFortranMPICall(called->getName());
1859
1860 Type *i32Ty = Type::getInt32Ty(call.getContext());
1861 Value *sendcountVal =
1862 fortranABI ? Builder2.CreateLoad(i32Ty, sendcount) : sendcount;
1863 Value *recvcountVal =
1864 fortranABI ? Builder2.CreateLoad(i32Ty, recvcount) : recvcount;
1865 Value *rootVal = fortranABI ? Builder2.CreateLoad(i32Ty, root) : root;
1866
1867 Value *rank = MPI_COMM_RANK(comm, Builder2, rootVal->getType(), called);
1868 Value *tysize = MPI_TYPE_SIZE(sendtype, Builder2, call.getType(), called);
1869
1870 if (forwardMode) {
1871 createMPIForwardCall(call, {/*sendbuf*/ 0, /*recvbuf*/ 3}, gutils,
1872 Builder2);
1873 return;
1874 }
1875 // Get the length for the allocation of the intermediate buffer
1876 auto recvlen_arg = Builder2.CreateZExtOrTrunc(
1877 recvcountVal, Type::getInt64Ty(call.getContext()));
1878 recvlen_arg =
1879 Builder2.CreateMul(recvlen_arg,
1880 Builder2.CreateZExtOrTrunc(
1881 tysize, Type::getInt64Ty(call.getContext())),
1882 "", true, true);
1883
1884 // Need to preserve the shadow send/recv buffers. The Fortran ABI call
1885 // has an extra `ierr` argument, so size the bundle to match the actual
1886 // call arity: shadow send/recv buffers, primal everything else.
1887 std::vector<ValueType> BufferBundleTypes(call.arg_size(),
1888 ValueType::Primal);
1889 BufferBundleTypes[0] = ValueType::Shadow;
1890 BufferBundleTypes[3] = ValueType::Shadow;
1891 auto BufferDefs = gutils->getInvertedBundles(&call, BufferBundleTypes,
1892 Builder2, /*lookup*/ true);
1893
1894 // 1. if root, malloc intermediate buffer, else undef
1895 PHINode *buf;
1896 PHINode *sendlen_phi;
1897
1898 {
1899 BasicBlock *currentBlock = Builder2.GetInsertBlock();
1900 BasicBlock *rootBlock = gutils->addReverseBlock(
1901 currentBlock, currentBlock->getName() + "_root", gutils->newFunc);
1902 BasicBlock *mergeBlock = gutils->addReverseBlock(
1903 rootBlock, currentBlock->getName() + "_post", gutils->newFunc);
1904
1905 Builder2.CreateCondBr(Builder2.CreateICmpEQ(rank, rootVal), rootBlock,
1906 mergeBlock);
1907
1908 Builder2.SetInsertPoint(rootBlock);
1909
1910 auto sendlen_arg = Builder2.CreateZExtOrTrunc(
1911 sendcountVal, Type::getInt64Ty(call.getContext()));
1912 sendlen_arg =
1913 Builder2.CreateMul(sendlen_arg,
1914 Builder2.CreateZExtOrTrunc(
1915 tysize, Type::getInt64Ty(call.getContext())),
1916 "", true, true);
1917 sendlen_arg = Builder2.CreateMul(
1918 sendlen_arg,
1919 Builder2.CreateZExtOrTrunc(
1920 MPI_COMM_SIZE(comm, Builder2, rootVal->getType(), called),
1921 Type::getInt64Ty(call.getContext())),
1922 "", true, true);
1923
1924 Value *rootbuf =
1925 CreateAllocation(Builder2, Type::getInt8Ty(call.getContext()),
1926 sendlen_arg, "mpireduce_malloccache");
1927
1928 Builder2.CreateBr(mergeBlock);
1929
1930 Builder2.SetInsertPoint(mergeBlock);
1931
1932 buf = Builder2.CreatePHI(rootbuf->getType(), 2);
1933 buf->addIncoming(rootbuf, rootBlock);
1934 buf->addIncoming(UndefValue::get(buf->getType()), currentBlock);
1935
1936 sendlen_phi = Builder2.CreatePHI(sendlen_arg->getType(), 2);
1937 sendlen_phi->addIncoming(sendlen_arg, rootBlock);
1938 sendlen_phi->addIncoming(UndefValue::get(sendlen_arg->getType()),
1939 currentBlock);
1940 }
1941
1942 // 2. Gather diff(recvbuffer) to intermediate buffer
1943 {
1944 // int MPI_Gather(const void *sendbuf, int sendcount, MPI_Datatype
1945 // sendtype,
1946 // void *recvbuf, int recvcount, MPI_Datatype recvtype,
1947 // int root, MPI_Comm comm)
1948 //
1949 // The Fortran MPI ABI passes all arguments by reference and appends
1950 // an `ierr` argument, so the generated call must match the convention
1951 // of the caller.
1952 SmallVector<Value *, 10> args;
1953 if (fortranABI) {
1954 args = {shadow_recvbuf, recvcount, recvtype, buf,
1955 sendcount, sendtype, root, comm};
1956 for (size_t i = 8, e = call.arg_size(); i < e; i++)
1957 args.push_back(lookup(
1958 gutils->getNewFromOriginal(call.getArgOperand(i)), Builder2));
1959 } else {
1960 args = {shadow_recvbuf, recvcountVal, recvtype, buf,
1961 sendcountVal, sendtype, rootVal, comm};
1962 }
1963 SmallVector<Type *, 10> types;
1964 for (auto *arg : args)
1965 types.push_back(arg->getType());
1966
1967 FunctionType *FT = FunctionType::get(call.getType(), types, false);
1968 Builder2.CreateCall(
1969 called->getParent()->getOrInsertFunction(
1970 getRenamedPerCallingConv(called->getName(), "MPI_Gather"), FT),
1971 args, BufferDefs);
1972 }
1973
1974 // 3. Zero diff(recvbuffer) [memset to 0]
1975 {
1976 auto val_arg = ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
1977 auto volatile_arg = ConstantInt::getFalse(call.getContext());
1978 Value *args[] = {shadow_recvbuf, val_arg, recvlen_arg, volatile_arg};
1979 Type *tys[] = {args[0]->getType(), args[2]->getType()};
1980 auto memset = cast<CallInst>(Builder2.CreateCall(
1981 getIntrinsicDeclaration(gutils->newFunc->getParent(),
1982 Intrinsic::memset, tys),
1983 args, BufferDefs));
1984 memset->addParamAttr(0, Attribute::NonNull);
1985 }
1986
1987 // 4. if root, diff(sendbuffer) += intermediate buffer (diffmemcopy)
1988 // 5. if root, free intermediate buffer
1989
1990 {
1991 BasicBlock *currentBlock = Builder2.GetInsertBlock();
1992 BasicBlock *rootBlock = gutils->addReverseBlock(
1993 currentBlock, currentBlock->getName() + "_root", gutils->newFunc);
1994 BasicBlock *mergeBlock = gutils->addReverseBlock(
1995 rootBlock, currentBlock->getName() + "_post", gutils->newFunc);
1996
1997 Builder2.CreateCondBr(Builder2.CreateICmpEQ(rank, rootVal), rootBlock,
1998 mergeBlock);
1999
2000 Builder2.SetInsertPoint(rootBlock);
2001
2002 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
2003 DifferentiableMemCopyFloats(call, orig_sendbuf, buf, shadow_sendbuf,
2004 sendlen_phi, Builder2, BufferDefs);
2005
2006 // Free up intermediate buffer
2007 if (shouldFree()) {
2008 CreateDealloc(Builder2, buf);
2009 }
2010
2011 Builder2.CreateBr(mergeBlock);
2012 Builder2.SetInsertPoint(mergeBlock);
2013 }
2014 }
2015 if (Mode == DerivativeMode::ReverseModeGradient)
2016 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2017 return;
2018 }
2019
2020 // Approximate algo (for sum): -> if statement yet to be
2021 // 1. malloc intermediate buffer
2022 // 2. reduce diff(recvbuffer) then scatter to corresponding input node's
2023 // intermediate buffer
2024 // 3. Zero diff(recvbuffer) [memset to 0]
2025 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
2026 // 5. free intermediate buffer
2027
2028 // int MPI_Allgather(const void *sendbuf, int sendcount, MPI_Datatype
2029 // sendtype,
2030 // void *recvbuf, int recvcount, MPI_Datatype recvtype,
2031 // MPI_Comm comm)
2032
2033 if (funcName == "MPI_Allgather" || funcName == "PMPI_Allgather") {
2034 if (Mode == DerivativeMode::ReverseModeGradient ||
2035 Mode == DerivativeMode::ReverseModeCombined ||
2036 Mode == DerivativeMode::ForwardMode ||
2037 Mode == DerivativeMode::ForwardModeError) {
2038 bool forwardMode = Mode == DerivativeMode::ForwardMode ||
2039 Mode == DerivativeMode::ForwardModeError;
2040
2041 IRBuilder<> Builder2 =
2042 forwardMode ? IRBuilder<>(&call) : IRBuilder<>(call.getParent());
2043 if (forwardMode) {
2044 getForwardBuilder(Builder2);
2045 } else {
2046 getReverseBuilder(Builder2);
2047 }
2048
2049 Value *orig_sendbuf = call.getOperand(0);
2050 Value *orig_sendcount = call.getOperand(1);
2051 Value *orig_sendtype = call.getOperand(2);
2052 Value *orig_recvbuf = call.getOperand(3);
2053 Value *orig_recvcount = call.getOperand(4);
2054 Value *orig_recvtype = call.getOperand(5);
2055 Value *orig_comm = call.getOperand(6);
2056
2057 Value *shadow_recvbuf = gutils->invertPointerM(orig_recvbuf, Builder2);
2058 if (!forwardMode)
2059 shadow_recvbuf = lookup(shadow_recvbuf, Builder2);
2060
2061 if (shadow_recvbuf->getType()->isIntegerTy())
2062 shadow_recvbuf = Builder2.CreateIntToPtr(
2063 shadow_recvbuf, getInt8PtrTy(call.getContext()));
2064
2065 Value *shadow_sendbuf = gutils->invertPointerM(orig_sendbuf, Builder2);
2066 if (!forwardMode)
2067 shadow_sendbuf = lookup(shadow_sendbuf, Builder2);
2068
2069 if (shadow_sendbuf->getType()->isIntegerTy())
2070 shadow_sendbuf = Builder2.CreateIntToPtr(
2071 shadow_sendbuf, getInt8PtrTy(call.getContext()));
2072
2073 Value *recvcount = gutils->getNewFromOriginal(orig_recvcount);
2074 if (!forwardMode)
2075 recvcount = lookup(recvcount, Builder2);
2076
2077 Value *recvtype = gutils->getNewFromOriginal(orig_recvtype);
2078 if (!forwardMode)
2079 recvtype = lookup(recvtype, Builder2);
2080
2081 Value *sendcount = gutils->getNewFromOriginal(orig_sendcount);
2082 if (!forwardMode)
2083 sendcount = lookup(sendcount, Builder2);
2084
2085 Value *sendtype = gutils->getNewFromOriginal(orig_sendtype);
2086 if (!forwardMode)
2087 sendtype = lookup(sendtype, Builder2);
2088
2089 Value *comm = gutils->getNewFromOriginal(orig_comm);
2090 if (!forwardMode)
2091 comm = lookup(comm, Builder2);
2092
2093 Value *tysize = MPI_TYPE_SIZE(sendtype, Builder2, call.getType(), called);
2094
2095 if (forwardMode) {
2096 createMPIForwardCall(call, {/*sendbuf*/ 0, /*recvbuf*/ 3}, gutils,
2097 Builder2);
2098 return;
2099 }
2100 // Get the length for the allocation of the intermediate buffer
2101 auto sendlen_arg = Builder2.CreateZExtOrTrunc(
2102 sendcount, Type::getInt64Ty(call.getContext()));
2103 sendlen_arg =
2104 Builder2.CreateMul(sendlen_arg,
2105 Builder2.CreateZExtOrTrunc(
2106 tysize, Type::getInt64Ty(call.getContext())),
2107 "", true, true);
2108
2109 // Need to preserve the shadow send/recv buffers.
2110 auto BufferDefs = gutils->getInvertedBundles(
2111 &call,
2112 {ValueType::Shadow, ValueType::Primal, ValueType::Primal,
2113 ValueType::Shadow, ValueType::Primal, ValueType::Primal,
2114 ValueType::Primal},
2115 Builder2, /*lookup*/ true);
2116
2117 // 1. Alloc intermediate buffer
2118 Value *buf =
2119 CreateAllocation(Builder2, Type::getInt8Ty(call.getContext()),
2120 sendlen_arg, "mpireduce_malloccache");
2121
2122 ConcreteType CT =
2123 TR.firstPointer(1, orig_sendbuf, &call, gutils, &Builder2);
2124 auto MPI_OP_type = getInt8PtrTy(call.getContext());
2125 Type *MPI_OP_Ptr_type = getUnqual(MPI_OP_type);
2126
2127 // 2. reduce diff(recvbuffer) then scatter to corresponding input node's
2128 // intermediate buffer
2129 {
2130 // int MPI_Reduce_scatter_block(const void* send_buffer,
2131 // void* receive_buffer,
2132 // int count,
2133 // MPI_Datatype datatype,
2134 // MPI_Op operation,
2135 // MPI_Comm communicator);
2136 Value *args[] = {
2137 /*sendbuf*/ shadow_recvbuf,
2138 /*recvbuf*/ buf,
2139 /*recvcount*/ sendcount,
2140 /*recvtype*/ sendtype,
2141 /*op (MPI_SUM)*/
2142 getOrInsertOpFloatSum(*gutils->newFunc->getParent(), called,
2143 MPI_OP_Ptr_type, MPI_OP_type, CT,
2144 call.getType(), Builder2),
2145 /*comm*/ comm,
2146 };
2147 Type *types[sizeof(args) / sizeof(*args)];
2148 for (size_t i = 0; i < sizeof(args) / sizeof(*args); i++)
2149 types[i] = args[i]->getType();
2150
2151 FunctionType *FT = FunctionType::get(call.getType(), types, false);
2152 Builder2.CreateCall(
2153 called->getParent()->getOrInsertFunction(
2154 getRenamedPerCallingConv(called->getName(),
2155 "MPI_Reduce_scatter_block"),
2156 FT),
2157 args, BufferDefs);
2158 }
2159
2160 // 3. zero diff(recvbuffer) [memset to 0]
2161 {
2162 auto recvlen_arg = Builder2.CreateZExtOrTrunc(
2163 recvcount, Type::getInt64Ty(call.getContext()));
2164 recvlen_arg =
2165 Builder2.CreateMul(recvlen_arg,
2166 Builder2.CreateZExtOrTrunc(
2167 tysize, Type::getInt64Ty(call.getContext())),
2168 "", true, true);
2169 recvlen_arg = Builder2.CreateMul(
2170 recvlen_arg,
2171 Builder2.CreateZExtOrTrunc(
2172 MPI_COMM_SIZE(comm, Builder2, call.getType(), called),
2173 Type::getInt64Ty(call.getContext())),
2174 "", true, true);
2175 auto val_arg = ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
2176 auto volatile_arg = ConstantInt::getFalse(call.getContext());
2177 Value *args[] = {shadow_recvbuf, val_arg, recvlen_arg, volatile_arg};
2178 Type *tys[] = {args[0]->getType(), args[2]->getType()};
2179 auto memset = cast<CallInst>(Builder2.CreateCall(
2180 getIntrinsicDeclaration(gutils->newFunc->getParent(),
2181 Intrinsic::memset, tys),
2182 args, BufferDefs));
2183 memset->addParamAttr(0, Attribute::NonNull);
2184 }
2185
2186 // 4. diff(sendbuffer) += intermediate buffer (diffmemcopy)
2187 DifferentiableMemCopyFloats(call, orig_sendbuf, buf, shadow_sendbuf,
2188 sendlen_arg, Builder2, BufferDefs);
2189
2190 // Free up intermediate buffer
2191 if (shouldFree()) {
2192 CreateDealloc(Builder2, buf);
2193 }
2194 }
2195 if (Mode == DerivativeMode::ReverseModeGradient)
2196 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2197 return;
2198 }
2199
2200 // Adjoint of barrier is to place a barrier at the corresponding
2201 // location in the reverse.
2202 if (funcName == "MPI_Barrier" || funcName == "PMPI_Barrier") {
2203 if (Mode == DerivativeMode::ReverseModeGradient ||
2204 Mode == DerivativeMode::ReverseModeCombined) {
2205 IRBuilder<> Builder2(&call);
2206 getReverseBuilder(Builder2);
2207 auto callval = call.getCalledOperand();
2208 // Copy all arguments to match the call's arity: Fortran ABI manglings
2209 // of MPI_Barrier (e.g. "mpi_barrier_") take an extra `ierr` argument.
2210 SmallVector<Value *, 4> args;
2211 for (unsigned i = 0; i < call.arg_size(); i++)
2212 args.push_back(lookup(gutils->getNewFromOriginal(call.getArgOperand(i)),
2213 Builder2));
2214 Builder2.CreateCall(call.getFunctionType(), callval, args);
2215 }
2216 if (Mode == DerivativeMode::ReverseModeGradient)
2217 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2218 return;
2219 }
2220
2221 // Remove free's in forward pass so the comm can be used in the reverse
2222 // pass
2223 if (funcName == "MPI_Comm_free" || funcName == "MPI_Comm_disconnect") {
2224 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2225 return;
2226 }
2227
2228 // Adjoint of MPI_Comm_split / MPI_Graph_create (which allocates a comm in a
2229 // pointer) is to free the created comm at the corresponding place in the
2230 // reverse pass
2231 auto commFound = MPIInactiveCommAllocators.find(funcName);
2232 if (commFound != MPIInactiveCommAllocators.end()) {
2233 if (Mode == DerivativeMode::ReverseModeGradient ||
2234 Mode == DerivativeMode::ReverseModeCombined) {
2235 IRBuilder<> Builder2(&call);
2236 getReverseBuilder(Builder2);
2237
2238 Value *args[] = {lookup(call.getOperand(commFound->second), Builder2)};
2239 Type *types[] = {args[0]->getType()};
2240
2241 FunctionType *FT = FunctionType::get(call.getType(), types, false);
2242 Builder2.CreateCall(
2243 called->getParent()->getOrInsertFunction(
2244 getRenamedPerCallingConv(called->getName(), "MPI_Comm_free"), FT),
2245 args);
2246 }
2247 if (Mode == DerivativeMode::ReverseModeGradient)
2248 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2249 return;
2250 }
2251
2252 // MPI_Init / MPI_Finalize don't participate in the computation being
2253 // differentiated - just duplicate them as-is in both passes.
2254 if (funcName == "MPI_Init" || funcName == "PMPI_Init" ||
2255 funcName == "MPI_Init_thread" || funcName == "PMPI_Init_thread" ||
2256 funcName == "MPI_Finalize" || funcName == "PMPI_Finalize" ||
2257 funcName == "MPI_Test" || funcName == "PMPI_Test" ||
2258 funcName == "MPI_Probe" || funcName == "PMPI_Probe") {
2259 if (Mode == DerivativeMode::ReverseModeGradient ||
2260 Mode == DerivativeMode::ReverseModeCombined ||
2261 Mode == DerivativeMode::ReverseModePrimal) {
2262 IRBuilder<> Builder2(&call);
2263 getReverseBuilder(Builder2);
2264 SmallVector<Value *, 8> args;
2265 for (unsigned i = 0; i < call.arg_size(); ++i) {
2266 args.push_back(lookup(gutils->getNewFromOriginal(call.getArgOperand(i)),
2267 Builder2));
2268 }
2269 Builder2.CreateCall(call.getFunctionType(), call.getCalledOperand(),
2270 args);
2271 }
2272 if (Mode == DerivativeMode::ForwardMode ||
2273 Mode == DerivativeMode::ForwardModeError ||
2274 Mode == DerivativeMode::ForwardModeSplit) {
2275 IRBuilder<> Builder2(&call);
2276 getForwardBuilder(Builder2);
2277 SmallVector<Value *, 8> args;
2278 for (unsigned i = 0; i < call.arg_size(); ++i) {
2279 args.push_back(gutils->getNewFromOriginal(call.getArgOperand(i)));
2280 }
2281 Builder2.CreateCall(call.getFunctionType(), call.getCalledOperand(),
2282 args);
2283 }
2284 if (Mode == DerivativeMode::ReverseModeGradient)
2285 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2286 return;
2287 }
2288
2289 llvm::errs() << *gutils->oldFunc->getParent() << "\n";
2290 llvm::errs() << *gutils->oldFunc << "\n";
2291 llvm::errs() << call << "\n";
2292 llvm::errs() << called << "\n";
2293 llvm_unreachable("Unhandled MPI FUNCTION");
2294}
2295
2296// Classify the op callee of a julia.atomicmodify pseudo-intrinsic call
2297// (op has signature elty (elty oldval, args...), with elty an integer type
2298// holding the bits of the modified value) as an atomicrmw-style operation
2299// whose derivative Enzyme knows. On success sets opArgNo to the index of the
2300// op parameter providing the update value and FT to the floating point type
2301// the update is performed in (nullptr for Xchg).
2302static llvm::AtomicRMWInst::BinOp
2303classifyAtomicModifyOp(llvm::Function *op, unsigned &opArgNo, llvm::Type *&FT) {
2304 using namespace llvm;
2305 FT = nullptr;
2306 if (!op || op->empty() || op->isVarArg() ||
2307 op->getFunctionType()->getNumParams() == 0)
2308 return AtomicRMWInst::BAD_BINOP;
2309 ReturnInst *Ret = nullptr;
2310 for (auto &BB : *op)
2311 if (auto R = dyn_cast<ReturnInst>(BB.getTerminator())) {
2312 if (Ret)
2313 return AtomicRMWInst::BAD_BINOP;
2314 Ret = R;
2315 }
2316 if (!Ret || !Ret->getReturnValue())
2317 return AtomicRMWInst::BAD_BINOP;
2318 auto peel = [](Value *V) {
2319 while (auto BC = dyn_cast<BitCastInst>(V))
2320 V = BC->getOperand(0);
2321 return V;
2322 };
2323 Value *RV = peel(Ret->getReturnValue());
2324 Value *Old = op->getArg(0);
2325 if (auto A = dyn_cast<Argument>(RV)) {
2326 if (A == Old)
2327 return AtomicRMWInst::BAD_BINOP;
2328 opArgNo = A->getArgNo();
2329 return AtomicRMWInst::Xchg;
2330 }
2331 if (auto BO = dyn_cast<BinaryOperator>(RV)) {
2332 Value *L = peel(BO->getOperand(0));
2333 Value *R = peel(BO->getOperand(1));
2334 FT = BO->getType();
2335 if (!FT->isFPOrFPVectorTy())
2336 return AtomicRMWInst::BAD_BINOP;
2337 if (BO->getOpcode() == Instruction::FAdd) {
2338 if (L == Old && isa<Argument>(R) && R != Old) {
2339 opArgNo = cast<Argument>(R)->getArgNo();
2340 return AtomicRMWInst::FAdd;
2341 }
2342 if (R == Old && isa<Argument>(L) && L != Old) {
2343 opArgNo = cast<Argument>(L)->getArgNo();
2344 return AtomicRMWInst::FAdd;
2345 }
2346 } else if (BO->getOpcode() == Instruction::FSub) {
2347 if (L == Old && isa<Argument>(R) && R != Old) {
2348 opArgNo = cast<Argument>(R)->getArgNo();
2349 return AtomicRMWInst::FSub;
2350 }
2351 }
2352 }
2353 return AtomicRMWInst::BAD_BINOP;
2354}
2355
2356bool AdjointGenerator::handleKnownCallDerivatives(
2357 CallInst &call, Function *called, StringRef funcName,
2358 bool subsequent_calls_may_write, const std::vector<bool> &overwritten_args,
2359 CallInst *const newCall) {
2360 bool subretused = false;
2361 bool shadowReturnUsed = false;
2362 DIFFE_TYPE subretType =
2363 gutils->getReturnDiffeType(&call, &subretused, &shadowReturnUsed);
2364
2365 IRBuilder<> BuilderZ(newCall);
2366 BuilderZ.setFastMathFlags(getFast());
2367
2368 // Julia's atomic modify pseudo-intrinsic (introduced in Julia 1.13):
2369 // {old, new} = julia.atomicmodify.iN.pAS(ptr, op, ordering, syncscope,
2370 // args...)
2371 // which atomically performs old = *ptr; new = op(old, args...);
2372 // *ptr = new, where op's parameter i (i >= 1) is forwarded from call
2373 // operand i + 3. The pseudo-intrinsic must be kept intact (including in
2374 // generated derivative code); it is only expanded to atomicrmw/cmpxchg by
2375 // Julia's ExpandAtomicModify pass after GC lowering.
2376 if (startsWith(funcName, "julia.atomicmodify.")) {
2377 // Retire the primal call replayed in the pass being generated. Like the
2378 // generic call path, tape its result in the augmented primal and read it
2379 // back from the tape in the derivative pass if the latter needs the
2380 // primal value: the modification must never be repeated by recomputing
2381 // the call there.
2382 auto finishPrimal = [&]() {
2383 bool primalNeededInReverse = false;
2384 {
2385 auto found = gutils->knownRecomputeHeuristic.find(&call);
2386 if (found != gutils->knownRecomputeHeuristic.end())
2387 primalNeededInReverse = !found->second;
2388 }
2389 if (!primalNeededInReverse &&
2390 Mode != DerivativeMode::ReverseModeCombined &&
2391 Mode != DerivativeMode::ForwardMode &&
2392 Mode != DerivativeMode::ForwardModeError && subretused &&
2393 !gutils->unnecessaryIntermediates.count(&call)) {
2394 std::map<UsageKey, bool> Seen =
2395 gutils->populateSeenFromKnownRecompute();
2396 auto minCutMode = (Mode == DerivativeMode::ReverseModePrimal)
2397 ? DerivativeMode::ReverseModeGradient
2398 : Mode;
2399 primalNeededInReverse =
2400 DifferentialUseAnalysis::is_value_needed_in_reverse<
2401 QueryType::Primal>(gutils, &call, minCutMode, Seen,
2402 oldUnreachable);
2403 }
2404 if (primalNeededInReverse) {
2405 gutils->cacheForReverse(BuilderZ, newCall,
2406 getIndex(&call, CacheType::Self, BuilderZ));
2407 eraseIfUnused(call);
2408 } else if (Mode == DerivativeMode::ReverseModeGradient ||
2409 Mode == DerivativeMode::ForwardModeSplit) {
2410 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2411 } else {
2412 eraseIfUnused(call);
2413 }
2414 };
2415
2416 // The pseudo-intrinsic called with the operands of the primal call,
2417 // except for the location and (if non-null) the operand at valIdx.
2418 auto emitCall = [&](Value *ptr, unsigned valIdx, Value *val) -> Value * {
2419 SmallVector<Value *, 6> args;
2420 for (size_t i = 0; i < call.arg_size(); i++) {
2421 if (i == 0)
2422 args.push_back(ptr);
2423 else if (i == valIdx && val)
2424 args.push_back(val);
2425 else
2426 args.push_back(gutils->getNewFromOriginal(call.getArgOperand(i)));
2427 }
2428 auto CI = BuilderZ.CreateCall(call.getFunctionType(), called, args);
2429 CI->setCallingConv(call.getCallingConv());
2430 CI->setAttributes(call.getAttributes());
2431 CI->copyMetadata(*newCall);
2432 return CI;
2433 };
2434
2435 // Fully inactive calls only need the primal replayed. This must be
2436 // handled here rather than falling through to the generic call path,
2437 // as the latter may decide against the constant fallback (e.g. for a
2438 // nocapture, non-readonly pointer argument) and then reject the
2439 // variadic call with `Number of arg operands != function parameters`.
2440 if (gutils->isConstantInstruction(&call) &&
2441 gutils->isConstantValue(&call)) {
2442 finishPrimal();
2443 return true;
2444 }
2445
2446 Type *elty = cast<StructType>(call.getType())->getElementType(0);
2447 unsigned opArgNo = 0;
2448 Type *FT = nullptr;
2449 auto opKind = classifyAtomicModifyOp(
2450 dyn_cast<Function>(call.getArgOperand(1)), opArgNo, FT);
2451 if (opKind != AtomicRMWInst::BAD_BINOP &&
2452 call.arg_size() - 3 != cast<Function>(call.getArgOperand(1))
2453 ->getFunctionType()
2454 ->getNumParams())
2455 opKind = AtomicRMWInst::BAD_BINOP;
2456 // Call operand forwarded to op's update value parameter.
2457 unsigned valIdx = opArgNo + 3;
2458 Value *valOp = opKind != AtomicRMWInst::BAD_BINOP
2459 ? call.getArgOperand(valIdx)
2460 : nullptr;
2461
2462 bool constval = gutils->isConstantValue(&call);
2463 bool constptr = gutils->isConstantValue(call.getArgOperand(0));
2464
2465 // No shadow memory is involved; replaying the primal suffices.
2466 if (constval && constptr) {
2467 finishPrimal();
2468 return true;
2469 }
2470
2471 auto &DL = gutils->newFunc->getParent()->getDataLayout();
2472 auto storeSize = (DL.getTypeSizeInBits(elty) + 7) / 8;
2473 auto vd = TR.firstPointer(storeSize, call.getArgOperand(0), &call, gutils,
2474 /*errifnotfound*/ nullptr,
2475 /*pointerIntSame*/ true);
2476
2477 bool constargs = true;
2478 for (size_t i = 4; i < call.arg_size(); i++)
2479 if (!gutils->isConstantValue(call.getArgOperand(i))) {
2480 constargs = false;
2481 break;
2482 }
2483
2484 // Non-differentiable (integer/pointer) data modified within duplicated
2485 // memory: replicate the modification on the shadow location with the
2486 // primal arguments (like inactive stores into active memory), keeping
2487 // e.g. lock states and counters of shadow objects consistent. As for an
2488 // inactive store, only the modified data has to be inactive, not the
2489 // operands op derives it from; when the type of the data is unknown, an
2490 // active operand however suggests floating point data, which is left
2491 // to the rules for recognized ops below.
2492 if (constval &&
2493 (vd.isKnown() ? !vd.isFloat() : (looseTypeAnalysis && constargs))) {
2494 // Which pass emits the shadow modification is decided exactly as for
2495 // an inactive store into duplicated memory (visitCommonStore); in
2496 // particular the augmented primal and the split derivative pass must
2497 // not both replay it, and a shadow that is only materialized in the
2498 // reverse pass must be updated there.
2499 bool forwardsShadow, backwardsShadow;
2500 shadowStoreSchedule(call, call.getArgOperand(0), forwardsShadow,
2501 backwardsShadow);
2502
2503 if ((Mode == DerivativeMode::ReverseModePrimal && forwardsShadow) ||
2504 (Mode == DerivativeMode::ReverseModeGradient && backwardsShadow) ||
2505 (Mode == DerivativeMode::ForwardModeSplit && backwardsShadow) ||
2506 (Mode == DerivativeMode::ReverseModeCombined &&
2507 (forwardsShadow || backwardsShadow)) ||
2508 Mode == DerivativeMode::ForwardMode ||
2509 Mode == DerivativeMode::ForwardModeError) {
2510 Value *dptr = gutils->invertPointerM(call.getArgOperand(0), BuilderZ);
2511 applyChainRule(
2512 BuilderZ, [&](Value *dptr) { emitCall(dptr, 0, nullptr); }, dptr);
2513 }
2514 finishPrimal();
2515 return true;
2516 }
2517
2518 if (Mode == DerivativeMode::ForwardMode ||
2519 Mode == DerivativeMode::ForwardModeError ||
2520 Mode == DerivativeMode::ForwardModeSplit) {
2521 if (opKind != AtomicRMWInst::BAD_BINOP) {
2522 // new = op(old, v) is linear in (old, v) for FAdd/FSub/Xchg, so the
2523 // same op applied to the shadow location and the shadow of v
2524 // computes the tangent {dold, dnew}.
2525 //
2526 // The shadow of v is its tangent (zero if inactive), except for an
2527 // inactive value exchanged into duplicated memory: like an inactive
2528 // store, that writes the primal value for pointer/integer data (and
2529 // zero for floating point data) into the shadow location.
2530 Value *dval = nullptr;
2531 if (!gutils->isConstantValue(valOp)) {
2532 dval = gutils->invertPointerM(valOp, BuilderZ);
2533 } else if (opKind == AtomicRMWInst::Xchg && !constptr) {
2534 if (!gutils->runtimeActivity && vd == BaseType::Pointer &&
2535 !isa<UndefValue>(valOp) && !isa<ConstantPointerNull>(valOp)) {
2536 std::string str;
2537 raw_string_ostream ss(str);
2538 ss << "Mismatched activity for: " << call
2539 << " const val: " << *valOp;
2540 if (CustomErrorHandler) {
2541 dval = unwrap(CustomErrorHandler(
2542 str.c_str(), wrap(&call), ErrorType::MixedActivityError,
2543 gutils, wrap(valOp), wrap(&BuilderZ)));
2544 } else
2545 EmitWarningAlways("MixedActivityError", call, ss.str(),
2546 MixedActivityHint);
2547 }
2548 if (!dval)
2549 dval = gutils->invertPointerM(valOp, BuilderZ,
2550 TypeTree(vd).Only(-1, nullptr));
2551 }
2552 if (!dval)
2553 dval =
2554 Constant::getNullValue(gutils->getShadowType(valOp->getType()));
2555
2556 auto rule = [&](Value *dptr, Value *dval) -> Value * {
2557 if (!dptr) {
2558 // Inactive location: dold = 0 and dnew = dv (-dv for FSub), in
2559 // the bits of the modified value.
2560 Value *dnew = dval;
2561 if (opKind == AtomicRMWInst::FSub) {
2562 dnew = BuilderZ.CreateBitCast(dnew, FT);
2563 dnew = BuilderZ.CreateFNeg(dnew);
2564 }
2565 dnew = BuilderZ.CreateBitCast(dnew, elty);
2566 return BuilderZ.CreateInsertValue(
2567 Constant::getNullValue(call.getType()), dnew, {1});
2568 }
2569 return emitCall(dptr, valIdx, dval);
2570 };
2571 Value *diff = applyChainRule(
2572 call.getType(), BuilderZ, rule,
2573 constptr ? nullptr
2574 : gutils->invertPointerM(call.getArgOperand(0), BuilderZ),
2575 dval);
2576 if (!constval)
2577 setDiffe(&call, diff, BuilderZ);
2578 finishPrimal();
2579 return true;
2580 }
2581 } else if (Mode == DerivativeMode::ReverseModePrimal) {
2582 // Nothing to do in the primal pass beyond replaying the call. Besides
2583 // the inactive case this also covers the augmented primal paired with
2584 // a ForwardModeSplit derivative, which handles any recognized op.
2585 if (constval || opKind != AtomicRMWInst::BAD_BINOP) {
2586 if (!constval)
2587 resolveShadowPlaceholder(call);
2588 finishPrimal();
2589 return true;
2590 }
2591 } else if ((Mode == DerivativeMode::ReverseModeCombined ||
2592 Mode == DerivativeMode::ReverseModeGradient) &&
2593 constval &&
2594 (opKind == AtomicRMWInst::FAdd ||
2595 opKind == AtomicRMWInst::FSub)) {
2596 if (!gutils->isConstantValue(valOp) && !constptr) {
2597 auto order = static_cast<AtomicOrdering>(
2598 cast<ConstantInt>(call.getArgOperand(2))->getZExtValue());
2599 auto ssid = static_cast<SyncScope::ID>(
2600 cast<ConstantInt>(call.getArgOperand(3))->getZExtValue());
2601 auto align = call.getParamAlign(0).value_or(DL.getABITypeAlign(elty));
2602 addAtomicRMWValueAdjoint(call, call.getArgOperand(0), valOp, elty, FT,
2603 /*negate*/ opKind == AtomicRMWInst::FSub,
2604 align, order, ssid, /*isVolatile*/ false);
2605 }
2606 finishPrimal();
2607 return true;
2608 }
2609
2610 // Remaining cases (active result in reverse mode, or an op that is not
2611 // a recognized linear modification) are not supported.
2612 std::string s;
2613 llvm::raw_string_ostream ss(s);
2614 ss << *gutils->oldFunc << "\n" << call << "\n";
2615 ss << " Active atomic modify not yet handled";
2616 Value *rval = EmitNoDerivativeError(ss.str(), call, gutils, BuilderZ);
2617 if (!constval) {
2618 if (Mode == DerivativeMode::ForwardMode ||
2619 Mode == DerivativeMode::ForwardModeError ||
2620 Mode == DerivativeMode::ForwardModeSplit) {
2621 if (!rval)
2622 rval = Constant::getNullValue(gutils->getShadowType(call.getType()));
2623 setDiffe(&call, rval, BuilderZ);
2624 } else {
2625 resolveShadowPlaceholder(call, rval);
2626 }
2627 }
2628 if (!call.getType()->isVoidTy()) {
2629 for (auto &U :
2630 make_early_inc_range(gutils->getNewFromOriginal(&call)->uses())) {
2631 U.set(UndefValue::get(call.getType()));
2632 }
2633 }
2634 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2635 return true;
2636 }
2637
2638 if (Mode != DerivativeMode::ReverseModePrimal && called) {
2639 if (funcName == "__kmpc_for_static_init_4" ||
2640 funcName == "__kmpc_for_static_init_4u" ||
2641 funcName == "__kmpc_for_static_init_8" ||
2642 funcName == "__kmpc_for_static_init_8u") {
2643 IRBuilder<> Builder2(&call);
2644 getReverseBuilder(Builder2);
2645 auto fini = called->getParent()->getFunction("__kmpc_for_static_fini");
2646 assert(fini);
2647 Value *args[] = {
2648 lookup(gutils->getNewFromOriginal(call.getArgOperand(0)), Builder2),
2649 lookup(gutils->getNewFromOriginal(call.getArgOperand(1)), Builder2)};
2650 auto fcall = Builder2.CreateCall(fini->getFunctionType(), fini, args);
2651 fcall->setCallingConv(fini->getCallingConv());
2652 return true;
2653 }
2654 }
2655
2656 // Canonicalize MPI routine names across calling conventions: the C
2657 // convention ("MPI_Recv", "PMPI_Recv") as well as Fortran ABI manglings
2658 // ("mpi_recv_", "mpi_comm_rank__", ...) all map to the canonical C name
2659 // (without profiling prefix) used throughout handleMPI.
2660 llvm::StringRef canonMPIName = canonicalizeMPIName(funcName);
2661 if (!canonMPIName.empty() &&
2662 (!gutils->isConstantInstruction(&call) || canonMPIName == "MPI_Barrier" ||
2663 canonMPIName == "MPI_Comm_free" ||
2664 canonMPIName == "MPI_Comm_disconnect" || canonMPIName == "MPI_Init" ||
2665 canonMPIName == "MPI_Init_thread" || canonMPIName == "MPI_Finalize" ||
2666 canonMPIName == "MPI_Test" || canonMPIName == "MPI_Probe" ||
2667 MPIInactiveCommAllocators.find(canonMPIName) !=
2668 MPIInactiveCommAllocators.end())) {
2669 handleMPI(call, called, canonMPIName);
2670 return true;
2671 }
2672
2673 if (auto blas = extractBLAS(funcName)) {
2674 if (handleBLAS(call, called, *blas, overwritten_args))
2675 return true;
2676 }
2677
2678 if (funcName == "printf" || funcName == "puts" ||
2679 startsWith(funcName, "_ZN3std2io5stdio6_print") ||
2680 startsWith(funcName, "_ZN4core3fmt")) {
2681 if (Mode == DerivativeMode::ReverseModeGradient) {
2682 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2683 }
2684 return true;
2685 }
2686 if (called && (called->getName().contains("__enzyme_float") ||
2687 called->getName().contains("__enzyme_double") ||
2688 called->getName().contains("__enzyme_integer") ||
2689 called->getName().contains("__enzyme_pointer"))) {
2690 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2691 return true;
2692 }
2693
2694 // Handle lgamma, safe to recompute so no store/change to forward
2695 if (called) {
2696 if (funcName == "__kmpc_for_static_init_4" ||
2697 funcName == "__kmpc_for_static_init_4u" ||
2698 funcName == "__kmpc_for_static_init_8" ||
2699 funcName == "__kmpc_for_static_init_8u") {
2700 if (Mode != DerivativeMode::ReverseModePrimal) {
2701 IRBuilder<> Builder2(&call);
2702 getReverseBuilder(Builder2);
2703 auto fini = called->getParent()->getFunction("__kmpc_for_static_fini");
2704 assert(fini);
2705 Value *args[] = {
2706 lookup(gutils->getNewFromOriginal(call.getArgOperand(0)), Builder2),
2707 lookup(gutils->getNewFromOriginal(call.getArgOperand(1)),
2708 Builder2)};
2709 auto fcall = Builder2.CreateCall(fini->getFunctionType(), fini, args);
2710 fcall->setCallingConv(fini->getCallingConv());
2711 }
2712 return true;
2713 }
2714 if (funcName == "__kmpc_for_static_fini") {
2715 if (Mode != DerivativeMode::ReverseModePrimal) {
2716 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
2717 }
2718 return true;
2719 }
2720 // TODO check
2721 // Adjoint of barrier is to place a barrier at the corresponding
2722 // location in the reverse.
2723 if (funcName == "__kmpc_barrier") {
2724 if (Mode == DerivativeMode::ReverseModeGradient ||
2725 Mode == DerivativeMode::ReverseModeCombined) {
2726 IRBuilder<> Builder2(&call);
2727 getReverseBuilder(Builder2);
2728 auto callval = call.getCalledOperand();
2729 Value *args[] = {
2730 lookup(gutils->getNewFromOriginal(call.getOperand(0)), Builder2),
2731 lookup(gutils->getNewFromOriginal(call.getOperand(1)), Builder2)};
2732 Builder2.CreateCall(call.getFunctionType(), callval, args);
2733 }
2734 return true;
2735 }
2736 if (funcName == "__kmpc_critical") {
2737 if (Mode != DerivativeMode::ReverseModePrimal) {
2738 IRBuilder<> Builder2(&call);
2739 getReverseBuilder(Builder2);
2740 auto crit2 = called->getParent()->getFunction("__kmpc_end_critical");
2741 assert(crit2);
2742 Value *args[] = {
2743 lookup(gutils->getNewFromOriginal(call.getArgOperand(0)), Builder2),
2744 lookup(gutils->getNewFromOriginal(call.getArgOperand(1)), Builder2),
2745 lookup(gutils->getNewFromOriginal(call.getArgOperand(2)),
2746 Builder2)};
2747 auto fcall = Builder2.CreateCall(crit2->getFunctionType(), crit2, args);
2748 fcall->setCallingConv(crit2->getCallingConv());
2749 }
2750 return true;
2751 }
2752 if (funcName == "__kmpc_end_critical") {
2753 if (Mode != DerivativeMode::ReverseModePrimal) {
2754 IRBuilder<> Builder2(&call);
2755 getReverseBuilder(Builder2);
2756 auto crit2 = called->getParent()->getFunction("__kmpc_critical");
2757 assert(crit2);
2758 Value *args[] = {
2759 lookup(gutils->getNewFromOriginal(call.getArgOperand(0)), Builder2),
2760 lookup(gutils->getNewFromOriginal(call.getArgOperand(1)), Builder2),
2761 lookup(gutils->getNewFromOriginal(call.getArgOperand(2)),
2762 Builder2)};
2763 auto fcall = Builder2.CreateCall(crit2->getFunctionType(), crit2, args);
2764 fcall->setCallingConv(crit2->getCallingConv());
2765 }
2766 return true;
2767 }
2768
2769 if (startsWith(funcName, "__kmpc") &&
2770 funcName != "__kmpc_global_thread_num") {
2771 std::string s;
2772 llvm::raw_string_ostream ss(s);
2773 ss << " unhandled openmp function: " << call << "\n";
2774 EmitNoDerivativeError(ss.str(), call, gutils, BuilderZ);
2775 return true;
2776 }
2777
2778 auto mod = call.getParent()->getParent()->getParent();
2779#include "CallDerivatives.inc"
2780
2781 if (funcName == "llvm.julia.gc_preserve_end") {
2782 if (Mode == DerivativeMode::ReverseModeGradient ||
2783 Mode == DerivativeMode::ReverseModeCombined) {
2784
2785 auto begin_call = cast<CallInst>(call.getOperand(0));
2786
2787 IRBuilder<> Builder2(&call);
2788 getReverseBuilder(Builder2);
2789
2790 if (!hasSingleGCPreserveEnd(begin_call, gutils)) {
2791 std::string s;
2792 llvm::raw_string_ostream ss(s);
2793 ss << "cannot reverse gc preserve region with several ends: "
2794 << *begin_call << "\n";
2795 EmitNoDerivativeError(ss.str(), call, gutils, Builder2);
2796 return true;
2797 }
2798 SmallVector<Value *, 1> args;
2799 for (auto &arg : begin_call->args()) {
2800 bool primalUsed = false;
2801 bool shadowUsed = false;
2802 gutils->getReturnDiffeType(arg, &primalUsed, &shadowUsed);
2803
2804 if (primalUsed)
2805 args.push_back(
2806 gutils->lookupM(gutils->getNewFromOriginal(arg), Builder2));
2807
2808 if (!gutils->isConstantValue(arg) && shadowUsed) {
2809 Value *ptrshadow = gutils->lookupM(
2810 gutils->invertPointerM(arg, BuilderZ), Builder2);
2811 if (gutils->getWidth() == 1)
2812 args.push_back(ptrshadow);
2813 else
2814 for (size_t i = 0; i < gutils->getWidth(); ++i)
2815 args.push_back(gutils->extractMeta(Builder2, ptrshadow, i));
2816 }
2817 }
2818
2819 auto newp = Builder2.CreateCall(
2820 called->getParent()->getOrInsertFunction(
2821 "llvm.julia.gc_preserve_begin",
2822 FunctionType::get(Type::getTokenTy(call.getContext()),
2823 ArrayRef<Type *>(), true)),
2824 args);
2825 auto ifound = gutils->invertedPointers.find(begin_call);
2826 assert(ifound != gutils->invertedPointers.end());
2827 auto placeholder = cast<CallInst>(&*ifound->second);
2828 gutils->invertedPointers.erase(ifound);
2829 gutils->invertedPointers.insert(std::make_pair(
2830 (const Value *)begin_call, InvertedPointerVH(gutils, newp)));
2831
2832 gutils->replaceAWithB(placeholder, newp);
2833 gutils->erase(placeholder);
2834 }
2835 return true;
2836 }
2837 if (funcName == "llvm.julia.gc_preserve_begin") {
2838 SmallVector<Value *, 1> args;
2839 for (auto &arg : call.args()) {
2840 bool primalUsed = false;
2841 bool shadowUsed = false;
2842 gutils->getReturnDiffeType(arg, &primalUsed, &shadowUsed);
2843
2844 if (primalUsed)
2845 args.push_back(gutils->getNewFromOriginal(arg));
2846
2847 if (!gutils->isConstantValue(arg) && shadowUsed) {
2848 Value *ptrshadow = gutils->invertPointerM(arg, BuilderZ);
2849 if (gutils->getWidth() == 1)
2850 args.push_back(ptrshadow);
2851 else
2852 for (size_t i = 0; i < gutils->getWidth(); ++i)
2853 args.push_back(gutils->extractMeta(BuilderZ, ptrshadow, i));
2854 }
2855 }
2856
2857 auto newp = BuilderZ.CreateCall(called, args);
2858 auto oldp = gutils->getNewFromOriginal(&call);
2859 gutils->replaceAWithB(oldp, newp);
2860 gutils->erase(oldp);
2861
2862 if (Mode == DerivativeMode::ReverseModeGradient ||
2863 Mode == DerivativeMode::ReverseModeCombined) {
2864 IRBuilder<> Builder2(&call);
2865 getReverseBuilder(Builder2);
2866
2867 auto ifound = gutils->invertedPointers.find(&call);
2868 assert(ifound != gutils->invertedPointers.end());
2869 auto placeholder = cast<CallInst>(&*ifound->second);
2870 if (!hasSingleGCPreserveEnd(&call, gutils)) {
2871 gutils->invertedPointers.erase(ifound);
2872 gutils->erase(placeholder);
2873 return true;
2874 }
2875 Builder2.CreateCall(
2876 called->getParent()->getOrInsertFunction(
2877 "llvm.julia.gc_preserve_end",
2878 FunctionType::get(Builder2.getVoidTy(), call.getType(), false)),
2879 placeholder);
2880 }
2881 return true;
2882 }
2883
2884 // void _FortranAAssign(Descriptor &to, const Descriptor &from,
2885 // const char *sourceFile, int sourceLine)
2886 // is LLVM flang's assignment between descriptors; its derivative is the
2887 // same assignment applied to the shadow descriptors.
2888 if (funcName == "_FortranAAssign" && call.arg_size() == 4) {
2889 if (Mode == DerivativeMode::ForwardMode ||
2890 Mode == DerivativeMode::ForwardModeError) {
2891 if (!gutils->isConstantInstruction(&call)) {
2892 IRBuilder<> Builder2(&call);
2893 getForwardBuilder(Builder2);
2894
2895 Value *file = gutils->getNewFromOriginal(call.getArgOperand(2));
2896 Value *line = gutils->getNewFromOriginal(call.getArgOperand(3));
2897 Value *shadowTo =
2898 gutils->invertPointerM(call.getArgOperand(0), Builder2);
2899 Value *shadowFrom =
2900 gutils->invertPointerM(call.getArgOperand(1), Builder2);
2901
2902 auto rule = [&](Value *sTo, Value *sFrom) {
2903 auto dcall = Builder2.CreateCall(called->getFunctionType(), called,
2904 {sTo, sFrom, file, line});
2905 dcall->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
2906 };
2907 applyChainRule(Builder2, rule, shadowTo, shadowFrom);
2908
2909 eraseIfUnused(call);
2910 return true;
2911 }
2912 }
2913 }
2914
2915 // void _FortranAEtime(const Descriptor *values, const Descriptor *time,
2916 // const char *sourceFile, int line)
2917 // is LLVM flang's ETIME. Its results are inactive, but it overwrites the
2918 // real(4) elements values(1:2) and time, whose shadow must then be zero:
2919 // after the call in forward mode, and in the reverse pass.
2920 if (funcName == "_FortranAEtime" && call.arg_size() == 4 &&
2921 (!gutils->isConstantValue(call.getArgOperand(0)) ||
2922 !gutils->isConstantValue(call.getArgOperand(1))) &&
2923 (Mode == DerivativeMode::ForwardMode ||
2924 Mode == DerivativeMode::ReverseModePrimal ||
2925 Mode == DerivativeMode::ReverseModeGradient ||
2926 Mode == DerivativeMode::ReverseModeCombined)) {
2927 auto &DL = gutils->newFunc->getParent()->getDataLayout();
2928 Type *I8PtrTy = getInt8PtrTy(call.getContext());
2929 Type *IdxTy = DL.getIntPtrType(call.getContext());
2930 unsigned P = DL.getPointerSize();
2931 // Offsets in a CFI descriptor: base_addr, and extent and sm of dim[0],
2932 // which follows elem_len, version, rank, type, attribute and extra.
2933 unsigned DimOff = 2 * P + 8;
2934
2935 // The shadow elements to zero, or null where the shadow is the primal
2936 // or the element does not exist.
2937 auto shadowElements = [&](IRBuilder<> &B) {
2938 auto field = [&](Value *desc, Type *T, unsigned off) {
2939 Value *p = B.CreatePointerCast(desc, I8PtrTy);
2940 p = B.CreateConstInBoundsGEP1_64(B.getInt8Ty(), p, off);
2941 return B.CreateLoad(T, B.CreatePointerCast(p, getUnqual(T)));
2942 };
2943 SmallVector<Value *, 6> elems;
2944 for (unsigned i = 0; i < 2; ++i) {
2945 Value *orig = call.getArgOperand(i);
2946 if (gutils->isConstantValue(orig))
2947 continue;
2948 Value *desc = gutils->getNewFromOriginal(orig);
2949 Value *sdescs = gutils->invertPointerM(orig, B);
2950 Value *base = field(desc, I8PtrTy, 0);
2951 for (unsigned w = 0; w < gutils->getWidth(); ++w) {
2952 Value *sdesc = gutils->getWidth() == 1
2953 ? sdescs
2954 : gutils->extractMeta(B, sdescs, w);
2955 Value *sbase = field(sdesc, I8PtrTy, 0);
2956 Value *notPrimal = B.CreateICmpNE(sbase, base);
2957 Value *notNull = B.CreateIsNotNull(sbase);
2958 Value *valid = B.CreateAnd(notPrimal, notNull);
2959 auto add = [&](Value *cond, Value *addr) {
2960 elems.push_back(
2961 B.CreateSelect(cond, addr, Constant::getNullValue(I8PtrTy)));
2962 };
2963 if (i == 1) {
2964 add(valid, sbase);
2965 continue;
2966 }
2967 Value *extent = field(desc, IdxTy, DimOff + P);
2968 Value *sm = field(desc, IdxTy, DimOff + 2 * P);
2969 Value *has1 = B.CreateICmpSGE(extent, ConstantInt::get(IdxTy, 1));
2970 add(B.CreateAnd(valid, has1), sbase);
2971 Value *has2 = B.CreateICmpSGE(extent, ConstantInt::get(IdxTy, 2));
2972 Value *valid2 = B.CreateAnd(valid, has2);
2973 add(valid2, B.CreateInBoundsGEP(B.getInt8Ty(), sbase, sm));
2974 }
2975 }
2976 return elems;
2977 };
2978 auto zero = [&](IRBuilder<> &B, Value *addr) {
2979 auto VT = FixedVectorType::get(B.getFloatTy(), 1);
2980 Value *ptr = B.CreatePointerCast(addr, getUnqual(VT));
2981 Value *notNull = B.CreateIsNotNull(addr);
2982 Value *mask = B.CreateVectorSplat(1, notNull);
2983 B.CreateMaskedStore(Constant::getNullValue(VT), ptr, Align(4), mask);
2984 };
2985
2986 if (Mode == DerivativeMode::ForwardMode) {
2987 IRBuilder<> B(newCall->getNextNode());
2988 for (auto addr : shadowElements(B))
2989 zero(B, addr);
2990 return true;
2991 }
2992
2993 unsigned N = 0;
2994 for (unsigned i = 0; i < 2; ++i)
2995 if (!gutils->isConstantValue(call.getArgOperand(i)))
2996 N += (i == 0 ? 2 : 1) * gutils->getWidth();
2997 Type *TapeTy = ArrayType::get(I8PtrTy, N);
2998 Value *tape;
2999 if (Mode == DerivativeMode::ReverseModeGradient) {
3000 tape = BuilderZ.CreatePHI(TapeTy, 0);
3001 } else {
3002 tape = UndefValue::get(TapeTy);
3003 auto elems = shadowElements(BuilderZ);
3004 for (unsigned i = 0; i < N; ++i)
3005 tape = BuilderZ.CreateInsertValue(tape, elems[i], i);
3006 if (auto I = dyn_cast<Instruction>(tape))
3007 gutils->TapesToPreventRecomputation.insert(I);
3008 }
3009 tape = gutils->cacheForReverse(
3010 BuilderZ, tape, getIndex(&call, CacheType::Tape, BuilderZ));
3011
3012 if (Mode == DerivativeMode::ReverseModeGradient ||
3013 Mode == DerivativeMode::ReverseModeCombined) {
3014 IRBuilder<> Builder2(&call);
3015 getReverseBuilder(Builder2);
3016 tape = lookup(tape, Builder2);
3017 for (unsigned i = 0; i < N; ++i)
3018 zero(Builder2, Builder2.CreateExtractValue(tape, i));
3019 }
3020 if (Mode == DerivativeMode::ReverseModeGradient)
3021 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
3022 return true;
3023 }
3024
3025 /*
3026 * int gsl_sf_legendre_array_e(const gsl_sf_legendre_t norm,
3027 const size_t lmax,
3028 const double x,
3029 const double csphase,
3030 double result_array[]);
3031 */
3032 // d L(n, x) / dx = L(n,x) * x * (n-1) + 1
3033 if (funcName == "gsl_sf_legendre_array_e") {
3034 if (gutils->isConstantValue(call.getArgOperand(4))) {
3035 eraseIfUnused(call);
3036 return true;
3037 }
3038 if (Mode == DerivativeMode::ReverseModePrimal) {
3039 eraseIfUnused(call);
3040 return true;
3041 }
3042 if (Mode == DerivativeMode::ReverseModeCombined ||
3043 Mode == DerivativeMode::ReverseModeGradient) {
3044 IRBuilder<> Builder2(&call);
3045 getReverseBuilder(Builder2);
3046 ValueType BundleTypes[5] = {ValueType::None, ValueType::None,
3047 ValueType::None, ValueType::None,
3048 ValueType::Shadow};
3049 auto Defs = gutils->getInvertedBundles(&call, BundleTypes, Builder2,
3050 /*lookup*/ true);
3051
3052 Type *types[6] = {
3053 call.getOperand(0)->getType(), call.getOperand(1)->getType(),
3054 call.getOperand(2)->getType(), call.getOperand(3)->getType(),
3055 call.getOperand(4)->getType(), call.getOperand(4)->getType(),
3056 };
3057 FunctionType *FT = FunctionType::get(call.getType(), types, false);
3058 auto F = getOrInsertPerCallingConv(*called->getParent(), called,
3059 "gsl_sf_legendre_deriv_array_e", FT);
3060
3061 llvm::Value *args[6] = {
3062 gutils->lookupM(gutils->getNewFromOriginal(call.getOperand(0)),
3063 Builder2),
3064 gutils->lookupM(gutils->getNewFromOriginal(call.getOperand(1)),
3065 Builder2),
3066 gutils->lookupM(gutils->getNewFromOriginal(call.getOperand(2)),
3067 Builder2),
3068 gutils->lookupM(gutils->getNewFromOriginal(call.getOperand(3)),
3069 Builder2),
3070 nullptr,
3071 nullptr};
3072
3073 Type *typesS[] = {args[1]->getType()};
3074 FunctionType *FTS =
3075 FunctionType::get(args[1]->getType(), typesS, false);
3076 auto FS = getOrInsertPerCallingConv(*called->getParent(), called,
3077 "gsl_sf_legendre_array_n", FTS);
3078 Value *alSize = Builder2.CreateCall(FS, args[1]);
3079 Value *tmp = CreateAllocation(Builder2, types[2], alSize);
3080 Value *dtmp = CreateAllocation(Builder2, types[2], alSize);
3081 Builder2.CreateLifetimeStart(tmp);
3082 Builder2.CreateLifetimeStart(dtmp);
3083
3084 args[4] = Builder2.CreateBitCast(tmp, types[4]);
3085 args[5] = Builder2.CreateBitCast(dtmp, types[5]);
3086
3087 Builder2.CreateCall(F, args, Defs);
3088 Builder2.CreateLifetimeEnd(tmp);
3089 CreateDealloc(Builder2, tmp);
3090
3091 BasicBlock *currentBlock = Builder2.GetInsertBlock();
3092
3093 BasicBlock *loopBlock = gutils->addReverseBlock(
3094 currentBlock, currentBlock->getName() + "_loop");
3095 BasicBlock *endBlock =
3096 gutils->addReverseBlock(loopBlock, currentBlock->getName() + "_end",
3097 /*fork*/ true, /*push*/ false);
3098
3099 Builder2.CreateCondBr(
3100 Builder2.CreateICmpEQ(args[1], Constant::getNullValue(types[1])),
3101 endBlock, loopBlock);
3102 Builder2.SetInsertPoint(loopBlock);
3103
3104 auto idx = Builder2.CreatePHI(types[1], 2);
3105 idx->addIncoming(ConstantInt::get(types[1], 0, false), currentBlock);
3106
3107 auto acc_idx = Builder2.CreatePHI(types[2], 2);
3108
3109 Value *inc = Builder2.CreateAdd(
3110 idx, ConstantInt::get(types[1], 1, false), "", true, true);
3111 idx->addIncoming(inc, loopBlock);
3112 acc_idx->addIncoming(Constant::getNullValue(types[2]), currentBlock);
3113
3114 Value *idxs[] = {idx};
3115 Value *dtmp_idx = Builder2.CreateInBoundsGEP(types[2], dtmp, idxs);
3116 Value *d_req = Builder2.CreateInBoundsGEP(
3117 types[2],
3118 Builder2.CreatePointerCast(
3119 gutils->invertPointerM(call.getOperand(4), Builder2),
3120 getUnqual(types[2])),
3121 idxs);
3122
3123 auto l0 = Builder2.CreateLoad(types[2], dtmp_idx);
3124 auto l1 = Builder2.CreateLoad(types[2], d_req);
3125 auto acc = Builder2.CreateFAdd(acc_idx, Builder2.CreateFMul(l0, l1));
3126 Builder2.CreateStore(Constant::getNullValue(types[2]), d_req);
3127
3128 acc_idx->addIncoming(acc, loopBlock);
3129
3130 Builder2.CreateCondBr(Builder2.CreateICmpEQ(inc, args[1]), endBlock,
3131 loopBlock);
3132
3133 Builder2.SetInsertPoint(endBlock);
3134 {
3135 auto found = gutils->reverseBlockToPrimal.find(endBlock);
3136 assert(found != gutils->reverseBlockToPrimal.end());
3137 SmallVector<BasicBlock *, 4> &vec =
3138 gutils->reverseBlocks[found->second];
3139 assert(vec.size());
3140 vec.push_back(endBlock);
3141 }
3142
3143 auto fin_idx = Builder2.CreatePHI(types[2], 2);
3144 fin_idx->addIncoming(Constant::getNullValue(types[2]), currentBlock);
3145 fin_idx->addIncoming(acc, loopBlock);
3146
3147 Builder2.CreateLifetimeEnd(dtmp);
3148 CreateDealloc(Builder2, dtmp);
3149
3150 ((DiffeGradientUtils *)gutils)
3151 ->addToDiffe(call.getOperand(2), fin_idx, Builder2, types[2]);
3152
3153 return true;
3154 }
3155 }
3156
3157 // Functions that only modify pointers and don't allocate memory,
3158 // needs to be run on shadow in primal
3159 if (funcName == "_ZSt29_Rb_tree_insert_and_rebalancebPSt18_Rb_tree_"
3160 "node_baseS0_RS_") {
3161 if (Mode == DerivativeMode::ReverseModeGradient) {
3162 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
3163 return true;
3164 }
3165 if (gutils->isConstantValue(call.getArgOperand(3)))
3166 return true;
3167 SmallVector<Value *, 2> args;
3168 for (auto &arg : call.args()) {
3169 if (gutils->isConstantValue(arg))
3170 args.push_back(gutils->getNewFromOriginal(arg));
3171 else
3172 args.push_back(gutils->invertPointerM(arg, BuilderZ));
3173 }
3174 BuilderZ.CreateCall(called, args);
3175 return true;
3176 }
3177
3178 // Functions that initialize a shadow data structure (with no
3179 // other arguments) needs to be run on shadow in primal.
3180 if (funcName == "_ZNSt8ios_baseC2Ev" || funcName == "_ZNSt8ios_baseD2Ev" ||
3181 funcName == "_ZNSt6localeC1Ev" || funcName == "_ZNSt6localeD1Ev" ||
3182 funcName == "_ZNKSt5ctypeIcE13_M_widen_initEv") {
3183 if (Mode == DerivativeMode::ReverseModeGradient ||
3184 Mode == DerivativeMode::ForwardModeSplit) {
3185 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
3186 return true;
3187 }
3188 if (gutils->isConstantValue(call.getArgOperand(0)))
3189 return true;
3190 Value *args[] = {gutils->invertPointerM(call.getArgOperand(0), BuilderZ)};
3191 BuilderZ.CreateCall(called, args);
3192 return true;
3193 }
3194
3195 if (funcName == "_ZNSt9basic_iosIcSt11char_traitsIcEE4initEPSt15basic_"
3196 "streambufIcS1_E") {
3197 if (Mode == DerivativeMode::ReverseModeGradient ||
3198 Mode == DerivativeMode::ForwardModeSplit) {
3199 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
3200 return true;
3201 }
3202 if (gutils->isConstantValue(call.getArgOperand(0)))
3203 return true;
3204 Value *args[] = {gutils->invertPointerM(call.getArgOperand(0), BuilderZ),
3205 gutils->invertPointerM(call.getArgOperand(1), BuilderZ)};
3206 BuilderZ.CreateCall(called, args);
3207 return true;
3208 }
3209
3210 // if constant instruction and readonly (thus must be pointer return)
3211 // and shadow return recomputable from shadow arguments.
3212 if (funcName == "__dynamic_cast" ||
3213 funcName == "_ZSt18_Rb_tree_decrementPKSt18_Rb_tree_node_base" ||
3214 funcName == "_ZSt18_Rb_tree_incrementPKSt18_Rb_tree_node_base" ||
3215 funcName == "_ZSt18_Rb_tree_decrementPSt18_Rb_tree_node_base" ||
3216 funcName == "_ZSt18_Rb_tree_incrementPSt18_Rb_tree_node_base" ||
3217 funcName == "jl_ptr_to_array" || funcName == "jl_ptr_to_array_1d") {
3218 bool shouldCache = false;
3219 if (gutils->knownRecomputeHeuristic.find(&call) !=
3220 gutils->knownRecomputeHeuristic.end()) {
3221 if (!gutils->knownRecomputeHeuristic[&call]) {
3222 shouldCache = true;
3223 }
3224 }
3225 ValueToValueMapTy empty;
3226 bool lrc = gutils->legalRecompute(&call, empty, nullptr);
3227
3228 if (!gutils->isConstantValue(&call)) {
3229 auto ifound = gutils->invertedPointers.find(&call);
3230 assert(ifound != gutils->invertedPointers.end());
3231 auto placeholder = cast<PHINode>(&*ifound->second);
3232
3233 if (subretType == DIFFE_TYPE::DUP_ARG) {
3234 Value *shadow = placeholder;
3235 if (lrc || Mode == DerivativeMode::ReverseModePrimal ||
3236 Mode == DerivativeMode::ReverseModeCombined ||
3237 Mode == DerivativeMode::ForwardMode ||
3238 Mode == DerivativeMode::ForwardModeError) {
3239 if (gutils->isConstantValue(call.getArgOperand(0)))
3240 shadow = gutils->getNewFromOriginal(&call);
3241 else {
3242 SmallVector<Value *, 2> args;
3243 size_t i = 0;
3244 for (auto &arg : call.args()) {
3245 if (gutils->isConstantValue(arg) ||
3246 (funcName == "__dynamic_cast" && i > 0) ||
3247 (funcName == "jl_ptr_to_array_1d" && i != 1) ||
3248 (funcName == "jl_ptr_to_array" && i != 1))
3249 args.push_back(gutils->getNewFromOriginal(arg));
3250 else
3251 args.push_back(gutils->invertPointerM(arg, BuilderZ));
3252 i++;
3253 }
3254 shadow = BuilderZ.CreateCall(called, args);
3255 }
3256 }
3257
3258 bool needsReplacement = true;
3259 if (!lrc && (Mode == DerivativeMode::ReverseModePrimal ||
3260 Mode == DerivativeMode::ReverseModeGradient)) {
3261 shadow = gutils->cacheForReverse(
3262 BuilderZ, shadow, getIndex(&call, CacheType::Shadow, BuilderZ));
3263 if (Mode == DerivativeMode::ReverseModeGradient)
3264 needsReplacement = false;
3265 }
3266 gutils->invertedPointers.erase((const Value *)&call);
3267 gutils->invertedPointers.insert(std::make_pair(
3268 (const Value *)&call, InvertedPointerVH(gutils, shadow)));
3269 if (needsReplacement) {
3270 assert(shadow != placeholder);
3271 gutils->replaceAWithB(placeholder, shadow);
3272 gutils->erase(placeholder);
3273 }
3274 } else {
3275 gutils->invertedPointers.erase((const Value *)&call);
3276 gutils->erase(placeholder);
3277 }
3278 }
3279
3280 if (Mode == DerivativeMode::ForwardMode ||
3281 Mode == DerivativeMode::ForwardModeError) {
3282 eraseIfUnused(call);
3283 assert(gutils->isConstantInstruction(&call));
3284 return true;
3285 }
3286
3287 if (!shouldCache && !lrc) {
3288 std::map<UsageKey, bool> Seen =
3289 gutils->populateSeenFromKnownRecompute();
3290 bool primalNeededInReverse =
3291 DifferentialUseAnalysis::is_value_needed_in_reverse<
3292 QueryType::Primal>(gutils, &call, Mode, Seen, oldUnreachable);
3293 {
3294 auto found = gutils->knownRecomputeHeuristic.find(&call);
3295 if (found != gutils->knownRecomputeHeuristic.end()) {
3296 if (!found->second) {
3297 primalNeededInReverse = true;
3298 }
3299 }
3300 }
3301 shouldCache = primalNeededInReverse;
3302 }
3303
3304 if (shouldCache) {
3305 BuilderZ.SetInsertPoint(newCall->getNextNode());
3306 gutils->cacheForReverse(BuilderZ, newCall,
3307 getIndex(&call, CacheType::Self, BuilderZ));
3308 }
3309 eraseIfUnused(call);
3310 assert(gutils->isConstantInstruction(&call));
3311 return true;
3312 }
3313
3314 if (called) {
3315 if (funcName == "julia.write_barrier" ||
3316 funcName == "julia.write_barrier_binding") {
3317 std::map<UsageKey, bool> Seen =
3318 gutils->populateSeenFromKnownRecompute();
3319 bool backwardsShadow = false;
3320 bool forwardsShadow = true;
3321 for (auto pair : gutils->backwardsOnlyShadows) {
3322 if (pair.second.stores.count(&call)) {
3323 backwardsShadow = true;
3324 forwardsShadow = pair.second.primalInitialize;
3325 if (auto inst = dyn_cast<Instruction>(pair.first))
3326 if (!forwardsShadow && pair.second.LI &&
3327 pair.second.LI->contains(inst->getParent()))
3328 backwardsShadow = false;
3329 break;
3330 }
3331 }
3332
3333 if (Mode == DerivativeMode::ForwardMode ||
3334 Mode == DerivativeMode::ForwardModeError ||
3335 (Mode == DerivativeMode::ReverseModeCombined &&
3336 (forwardsShadow || backwardsShadow)) ||
3337 (Mode == DerivativeMode::ReverseModePrimal && forwardsShadow) ||
3338 (Mode == DerivativeMode::ReverseModeGradient && backwardsShadow)) {
3339 IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&call));
3340 for (size_t i = 0; i < gutils->getWidth(); i++) {
3341 SmallVector<Value *, 1> iargs;
3342 bool first = true;
3343 for (auto &arg : call.args()) {
3344 if (!gutils->isConstantValue(arg)) {
3345 Value *ptrshadow = gutils->invertPointerM(arg, BuilderZ);
3346 if (gutils->getWidth() > 1) {
3347 ptrshadow = gutils->extractMeta(BuilderZ, ptrshadow, i);
3348 }
3349 iargs.push_back(ptrshadow);
3350 } else {
3351 if (first)
3352 break;
3353 }
3354 first = false;
3355 }
3356 if (iargs.size()) {
3357 BuilderZ.CreateCall(called, iargs);
3358 }
3359 }
3360 }
3361
3362 bool forceErase = false;
3363 if (Mode == DerivativeMode::ReverseModeGradient) {
3364
3365 // Since we won't redo the store in the reverse pass, do not
3366 // force the write barrier.
3367 forceErase = true;
3368 for (const auto &pair : gutils->rematerializableAllocations) {
3369 if (!pair.second.stores.count(&call))
3370 continue;
3371 if (gutils->allocationsToBeRematerialized.count(pair.first))
3372 // However, if we are rematerailizing the allocation and not
3373 // inside the loop level rematerialization, we do still need the
3374 // reverse passes ``fake primal'' store and therefore write
3375 // barrier
3376 if (!pair.second.LI || !pair.second.LI->contains(&call)) {
3377 forceErase = false;
3378 }
3379 }
3380 }
3381 if (forceErase)
3382 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
3383 else
3384 eraseIfUnused(call);
3385
3386 return true;
3387 }
3388 Intrinsic::ID ID = Intrinsic::not_intrinsic;
3389 if (isMemFreeLibMFunction(funcName, &ID)) {
3390 if (Mode == DerivativeMode::ReverseModePrimal ||
3391 gutils->isConstantInstruction(&call)) {
3392
3393 if (gutils->knownRecomputeHeuristic.find(&call) !=
3394 gutils->knownRecomputeHeuristic.end()) {
3395 if (!gutils->knownRecomputeHeuristic[&call]) {
3396 gutils->cacheForReverse(
3397 BuilderZ, newCall,
3398 getIndex(&call, CacheType::Self, BuilderZ));
3399 }
3400 }
3401 eraseIfUnused(call);
3402 return true;
3403 }
3404
3405 if (ID != Intrinsic::not_intrinsic) {
3406 SmallVector<Value *, 2> orig_ops(call.getNumOperands());
3407 for (unsigned i = 0; i < call.getNumOperands(); ++i) {
3408 orig_ops[i] = call.getOperand(i);
3409 }
3410 bool cached = handleAdjointForIntrinsic(ID, call, orig_ops);
3411 if (!cached) {
3412 if (gutils->knownRecomputeHeuristic.find(&call) !=
3413 gutils->knownRecomputeHeuristic.end()) {
3414 if (!gutils->knownRecomputeHeuristic[&call]) {
3415 gutils->cacheForReverse(
3416 BuilderZ, newCall,
3417 getIndex(&call, CacheType::Self, BuilderZ));
3418 }
3419 }
3420 }
3421 eraseIfUnused(call);
3422 return true;
3423 }
3424 }
3425 }
3426 }
3427 if (auto assembly = dyn_cast<InlineAsm>(call.getCalledOperand())) {
3428 if (assembly->getAsmString() == "maxpd $1, $0") {
3429 if (Mode == DerivativeMode::ReverseModePrimal ||
3430 gutils->isConstantInstruction(&call)) {
3431
3432 if (gutils->knownRecomputeHeuristic.find(&call) !=
3433 gutils->knownRecomputeHeuristic.end()) {
3434 if (!gutils->knownRecomputeHeuristic[&call]) {
3435 gutils->cacheForReverse(BuilderZ, newCall,
3436 getIndex(&call, CacheType::Self, BuilderZ));
3437 }
3438 }
3439 eraseIfUnused(call);
3440 return true;
3441 }
3442
3443 SmallVector<Value *, 2> orig_ops(call.getNumOperands());
3444 for (unsigned i = 0; i < call.getNumOperands(); ++i) {
3445 orig_ops[i] = call.getOperand(i);
3446 }
3447 handleAdjointForIntrinsic(Intrinsic::maxnum, call, orig_ops);
3448 if (gutils->knownRecomputeHeuristic.find(&call) !=
3449 gutils->knownRecomputeHeuristic.end()) {
3450 if (!gutils->knownRecomputeHeuristic[&call]) {
3451 gutils->cacheForReverse(BuilderZ, newCall,
3452 getIndex(&call, CacheType::Self, BuilderZ));
3453 }
3454 }
3455 eraseIfUnused(call);
3456 return true;
3457 }
3458 }
3459
3460 if (funcName == "realloc") {
3461 if (Mode == DerivativeMode::ForwardMode ||
3462 Mode == DerivativeMode::ForwardModeError) {
3463 if (!gutils->isConstantValue(&call)) {
3464 IRBuilder<> Builder2(&call);
3465 getForwardBuilder(Builder2);
3466
3467 auto dbgLoc = gutils->getNewFromOriginal(&call)->getDebugLoc();
3468
3469 auto rule = [&](Value *ip) {
3470 ValueType BundleTypes[2] = {ValueType::Shadow, ValueType::Primal};
3471
3472 auto Defs = gutils->getInvertedBundles(&call, BundleTypes, Builder2,
3473 /*lookup*/ false);
3474
3475 llvm::Value *args[2] = {
3476 ip, gutils->getNewFromOriginal(call.getOperand(1))};
3477 CallInst *CI = Builder2.CreateCall(
3478 call.getFunctionType(), call.getCalledFunction(), args, Defs);
3479 CI->setAttributes(call.getAttributes());
3480 CI->setCallingConv(call.getCallingConv());
3481 CI->setTailCallKind(call.getTailCallKind());
3482 CI->setDebugLoc(dbgLoc);
3483 return CI;
3484 };
3485
3486 Value *CI = applyChainRule(
3487 call.getType(), Builder2, rule,
3488 gutils->invertPointerM(call.getOperand(0), Builder2));
3489
3490 auto found = gutils->invertedPointers.find(&call);
3491 PHINode *placeholder = cast<PHINode>(&*found->second);
3492
3493 gutils->invertedPointers.erase(found);
3494 gutils->replaceAWithB(placeholder, CI);
3495 gutils->erase(placeholder);
3496 gutils->invertedPointers.insert(
3497 std::make_pair(&call, InvertedPointerVH(gutils, CI)));
3498 }
3499 eraseIfUnused(call);
3500 return true;
3501 }
3502 }
3503
3504 if (isAllocationFunction(funcName, gutils->TLI)) {
3505
3506 bool constval = gutils->isConstantValue(&call);
3507
3508 if (!constval) {
3509 auto dbgLoc = gutils->getNewFromOriginal(&call)->getDebugLoc();
3510 auto found = gutils->invertedPointers.find(&call);
3511 PHINode *placeholder = cast<PHINode>(&*found->second);
3512 IRBuilder<> bb(placeholder);
3513
3514 SmallVector<Value *, 8> args;
3515 for (auto &arg : call.args()) {
3516 args.push_back(gutils->getNewFromOriginal(arg));
3517 }
3518
3519 if (Mode == DerivativeMode::ReverseModeCombined ||
3520 Mode == DerivativeMode::ReverseModeGradient ||
3521 Mode == DerivativeMode::ReverseModePrimal ||
3522 Mode == DerivativeMode::ForwardModeSplit) {
3523
3524 Value *anti = placeholder;
3525 // If rematerializable allocations and split mode, we can
3526 // simply elect to build the entire piece in the reverse
3527 // since it should be possible to perform any shadow stores
3528 // of pointers (from rematerializable property) and it does
3529 // not escape the function scope (lest it not be
3530 // rematerializable) so all input derivatives remain zero.
3531 bool backwardsShadow = false;
3532 bool forwardsShadow = true;
3533 bool inLoop = false;
3534 bool isAlloca = isa<AllocaInst>(&call);
3535 {
3536 auto found = gutils->backwardsOnlyShadows.find(&call);
3537 if (found != gutils->backwardsOnlyShadows.end()) {
3538 backwardsShadow = true;
3539 forwardsShadow = found->second.primalInitialize;
3540 // If in a loop context, maintain the same free behavior.
3541 if (found->second.LI &&
3542 found->second.LI->contains(call.getParent()))
3543 inLoop = true;
3544 }
3545 }
3546 {
3547
3548 if (!forwardsShadow) {
3549 if (Mode == DerivativeMode::ReverseModePrimal) {
3550 // Needs a stronger replacement check/assertion.
3551 Value *replacement = getUndefinedValueForType(
3552 *gutils->oldFunc->getParent(), placeholder->getType());
3553 gutils->replaceAWithB(placeholder, replacement);
3554 gutils->invertedPointers.erase(found);
3555 gutils->invertedPointers.insert(std::make_pair(
3556 &call, InvertedPointerVH(gutils, replacement)));
3557 gutils->erase(placeholder);
3558 anti = nullptr;
3559 goto endAnti;
3560 } else if (inLoop) {
3561 gutils->rematerializedPrimalOrShadowAllocations.push_back(
3562 placeholder);
3563 if (hasMetadata(&call, "enzyme_fromstack"))
3564 isAlloca = true;
3565 goto endAnti;
3566 }
3567 }
3568 placeholder->setName("");
3569 if (shadowHandlers.find(funcName) != shadowHandlers.end()) {
3570 bb.SetInsertPoint(placeholder);
3571
3572 if (Mode == DerivativeMode::ReverseModeCombined ||
3573 (Mode == DerivativeMode::ReverseModePrimal && forwardsShadow) ||
3574 (Mode == DerivativeMode::ReverseModeGradient &&
3575 backwardsShadow)) {
3576 anti = applyChainRule(call.getType(), bb, [&]() {
3577 return shadowHandlers[funcName](bb, &call, args, gutils);
3578 });
3579 if (anti->getType() != placeholder->getType()) {
3580 llvm::errs() << "orig: " << call << "\n";
3581 llvm::errs() << "placeholder: " << *placeholder << "\n";
3582 llvm::errs() << "anti: " << *anti << "\n";
3583 }
3584 gutils->invertedPointers.erase(found);
3585 bb.SetInsertPoint(placeholder);
3586
3587 gutils->replaceAWithB(placeholder, anti);
3588 gutils->erase(placeholder);
3589 }
3590
3591 if (auto inst = dyn_cast<Instruction>(anti))
3592 bb.SetInsertPoint(inst);
3593
3594 if (!backwardsShadow)
3595 anti = gutils->cacheForReverse(
3596 bb, anti, getIndex(&call, CacheType::Shadow, BuilderZ));
3597 } else {
3598 bool zeroed = false;
3599 uint64_t idx = 0;
3600 Value *prev = nullptr;
3601 ;
3602 auto rule = [&]() {
3603 Value *anti =
3604 bb.CreateCall(call.getFunctionType(), call.getCalledOperand(),
3605 args, call.getName() + "'mi");
3606 cast<CallInst>(anti)->setAttributes(call.getAttributes());
3607 cast<CallInst>(anti)->setCallingConv(call.getCallingConv());
3608 cast<CallInst>(anti)->setTailCallKind(call.getTailCallKind());
3609 cast<CallInst>(anti)->setDebugLoc(dbgLoc);
3610
3611 if (anti->getType()->isPointerTy()) {
3612 cast<CallInst>(anti)->addAttributeAtIndex(
3613 AttributeList::ReturnIndex, Attribute::NoAlias);
3614 cast<CallInst>(anti)->addAttributeAtIndex(
3615 AttributeList::ReturnIndex, Attribute::NonNull);
3616
3617 if (funcName == "malloc" || funcName == "_Znwm" ||
3618 funcName == "??2@YAPAXI@Z" ||
3619 funcName == "??2@YAPEAX_K@Z") {
3620 if (auto ci = dyn_cast<ConstantInt>(args[0])) {
3621 unsigned derefBytes = ci->getLimitedValue();
3622 CallInst *cal =
3623 cast<CallInst>(gutils->getNewFromOriginal(&call));
3624 cast<CallInst>(anti)->addDereferenceableRetAttr(derefBytes);
3625 cal->addDereferenceableRetAttr(derefBytes);
3626#if !defined(FLANG) && !defined(ROCM)
3627 AttrBuilder B(ci->getContext());
3628#else
3629 AttrBuilder B;
3630#endif
3631 B.addDereferenceableOrNullAttr(derefBytes);
3632 cast<CallInst>(anti)->setAttributes(
3633 cast<CallInst>(anti)->getAttributes().addRetAttributes(
3634 call.getContext(), B));
3635 cal->setAttributes(cal->getAttributes().addRetAttributes(
3636 call.getContext(), B));
3637 cal->addAttributeAtIndex(AttributeList::ReturnIndex,
3638 Attribute::NoAlias);
3639 cal->addAttributeAtIndex(AttributeList::ReturnIndex,
3640 Attribute::NonNull);
3641 }
3642 }
3643 if (funcName == "julia.gc_alloc_obj" ||
3644 funcName == "jl_gc_alloc_typed" ||
3645 funcName == "ijl_gc_alloc_typed") {
3646 if (EnzymeShadowAllocRewrite) {
3647 bool used = unnecessaryInstructions.find(&call) ==
3648 unnecessaryInstructions.end();
3649 EnzymeShadowAllocRewrite(wrap(anti), gutils, wrap(&call),
3650 idx, wrap(prev), used);
3651 }
3652 }
3653 }
3654 if (Mode == DerivativeMode::ReverseModeCombined ||
3655 (Mode == DerivativeMode::ReverseModePrimal &&
3656 forwardsShadow) ||
3657 (Mode == DerivativeMode::ReverseModeGradient &&
3658 backwardsShadow) ||
3659 (Mode == DerivativeMode::ForwardModeSplit &&
3660 backwardsShadow)) {
3661 if (!inLoop) {
3662 zeroKnownAllocation(bb, anti, args, funcName, gutils->TLI,
3663 &call);
3664 zeroed = true;
3665 }
3666 }
3667 idx++;
3668 prev = anti;
3669 return anti;
3670 };
3671
3672 anti = applyChainRule(call.getType(), bb, rule);
3673
3674 gutils->invertedPointers.erase(found);
3675 if (&*bb.GetInsertPoint() == placeholder)
3676 bb.SetInsertPoint(placeholder->getNextNode());
3677 gutils->replaceAWithB(placeholder, anti);
3678 gutils->erase(placeholder);
3679
3680 if (!backwardsShadow)
3681 anti = gutils->cacheForReverse(
3682 bb, anti, getIndex(&call, CacheType::Shadow, BuilderZ));
3683 else {
3684 if (auto MD = hasMetadata(&call, "enzyme_fromstack")) {
3685 isAlloca = true;
3686 bb.SetInsertPoint(cast<Instruction>(anti));
3687 Value *Size;
3688 if (funcName == "malloc")
3689 Size = args[0];
3690 else if (funcName == "julia.gc_alloc_obj" ||
3691 funcName == "jl_gc_alloc_typed" ||
3692 funcName == "ijl_gc_alloc_typed")
3693 Size = args[1];
3694 else
3695 llvm_unreachable("Unknown allocation to upgrade");
3696
3697 Type *elTy = Type::getInt8Ty(call.getContext());
3698 if (MD->getNumOperands() == 2) {
3699 elTy = (Type *)cast<ConstantInt>(
3700 cast<ConstantAsMetadata>(MD->getOperand(1))
3701 ->getValue())
3702 ->getLimitedValue();
3703 Value *tsize = ConstantInt::get(
3704 Size->getType(), (gutils->newFunc->getParent()
3705 ->getDataLayout()
3706 .getTypeAllocSizeInBits(elTy) +
3707 7) /
3708 8);
3709
3710 Size = bb.CreateUDiv(Size, tsize, "", /*exact*/ true);
3711 }
3712 std::string name = "";
3713#if LLVM_VERSION_MAJOR < 17
3714 if (call.getContext().supportsTypedPointers()) {
3715 for (auto U : call.users()) {
3716 if (hasMetadata(cast<Instruction>(U), "enzyme_caststack")) {
3717 if (MD->getNumOperands() == 1) {
3718 elTy = U->getType()->getPointerElementType();
3719 Value *tsize = ConstantInt::get(
3720 Size->getType(),
3721 (gutils->newFunc->getParent()
3722 ->getDataLayout()
3723 .getTypeAllocSizeInBits(elTy) +
3724 7) /
3725 8);
3726
3727 Size = bb.CreateUDiv(Size, tsize, "", /*exact*/ true);
3728 }
3729 name = (U->getName() + "'ai").str();
3730 break;
3731 }
3732 }
3733 }
3734#endif
3735 auto rule = [&](Value *anti) {
3736 bb.SetInsertPoint(cast<Instruction>(anti));
3737 Value *replacement = bb.CreateAlloca(elTy, Size, name);
3738 if (name.size() == 0)
3739 replacement->takeName(anti);
3740 else
3741 anti->setName("");
3742 auto Alignment = cast<ConstantInt>(cast<ConstantAsMetadata>(
3743 MD->getOperand(0))
3744 ->getValue())
3745 ->getLimitedValue();
3746 if (Alignment) {
3747 cast<AllocaInst>(replacement)
3748 ->setAlignment(Align(Alignment));
3749 }
3750#if LLVM_VERSION_MAJOR < 17
3751 if (call.getContext().supportsTypedPointers()) {
3752 if (anti->getType()->getPointerElementType() != elTy)
3753 replacement = bb.CreatePointerCast(
3754 replacement,
3755 getUnqual(anti->getType()->getPointerElementType()));
3756 }
3757#endif
3758 auto PT = cast<PointerType>(anti->getType());
3759 if (PT->getAddressSpace()) {
3760 replacement = bb.CreateAddrSpaceCast(replacement, PT);
3761 cast<Instruction>(replacement)
3762 ->setMetadata(
3763 "enzyme_backstack",
3764 MDNode::get(replacement->getContext(), {}));
3765 }
3766 gutils->replaceAWithB(cast<Instruction>(anti), replacement);
3767 bb.SetInsertPoint(cast<Instruction>(anti)->getNextNode());
3768 gutils->erase(cast<Instruction>(anti));
3769 return replacement;
3770 };
3771
3772 auto replacement =
3773 applyChainRule(call.getType(), bb, rule, anti);
3774 anti = replacement;
3775 }
3776 }
3777
3778 if (Mode == DerivativeMode::ReverseModeCombined ||
3779 (Mode == DerivativeMode::ReverseModePrimal && forwardsShadow) ||
3780 (Mode == DerivativeMode::ReverseModeGradient &&
3781 backwardsShadow) ||
3782 (Mode == DerivativeMode::ForwardModeSplit && backwardsShadow)) {
3783 if (!inLoop) {
3784 assert(zeroed);
3785 }
3786 }
3787 }
3788 gutils->invertedPointers.insert(
3789 std::make_pair(&call, InvertedPointerVH(gutils, anti)));
3790 }
3791 endAnti:;
3792 if (((Mode == DerivativeMode::ReverseModeCombined && shouldFree()) ||
3793 (Mode == DerivativeMode::ReverseModeGradient && shouldFree()) ||
3794 (Mode == DerivativeMode::ForwardModeSplit && shouldFree())) &&
3795 !isAlloca) {
3796 IRBuilder<> Builder2(&call);
3797 getReverseBuilder(Builder2);
3798 assert(anti);
3799 Value *tofree = lookup(anti, Builder2);
3800 assert(tofree);
3801 assert(tofree->getType());
3802 for (size_t i = 0; i < gutils->getWidth(); i++) {
3803 Value *tofree_i =
3804 gutils->getWidth() == 1
3805 ? tofree
3806 : GradientUtils::extractMeta(Builder2, tofree, i);
3807
3808 auto CI = freeKnownAllocation(Builder2, tofree_i, funcName, dbgLoc,
3809 gutils->TLI, &call, gutils);
3810 if (CI) {
3811 CI->addAttributeAtIndex(AttributeList::FirstArgIndex,
3812 Attribute::NonNull);
3813 bool combined = Mode == DerivativeMode::ReverseModeCombined;
3814 auto ident = MDNode::getDistinct(
3815 CI->getContext(),
3816 {ConstantAsMetadata::get(
3817 combined ? ConstantInt::getTrue(CI->getContext())
3818 : ConstantInt::getFalse(CI->getContext()))});
3819 Value *anti_i =
3820 gutils->getWidth() == 1
3821 ? anti
3822 : GradientUtils::extractMeta(Builder2, anti, i);
3823 cast<Instruction>(anti_i)->setMetadata(
3824 "enzyme_cache_alloc", MDNode::get(CI->getContext(), {ident}));
3825 CI->setMetadata("enzyme_cache_free",
3826 MDNode::get(CI->getContext(), {ident}));
3827 }
3828 }
3829 }
3830 } else if (Mode == DerivativeMode::ForwardMode ||
3831 Mode == DerivativeMode::ForwardModeError) {
3832 IRBuilder<> Builder2(&call);
3833 getForwardBuilder(Builder2);
3834
3835 SmallVector<Value *, 2> args;
3836 for (unsigned i = 0; i < call.arg_size(); ++i) {
3837 auto arg = call.getArgOperand(i);
3838 args.push_back(gutils->getNewFromOriginal(arg));
3839 }
3840
3841 uint64_t idx = 0;
3842 Value *prev = gutils->getNewFromOriginal(&call);
3843 auto rule = [&]() {
3844 SmallVector<ValueType, 2> BundleTypes(args.size(), ValueType::Primal);
3845
3846 auto Defs = gutils->getInvertedBundles(&call, BundleTypes, Builder2,
3847 /*lookup*/ false);
3848
3849 CallInst *CI = Builder2.CreateCall(
3850 call.getFunctionType(), call.getCalledFunction(), args, Defs);
3851 CI->setAttributes(call.getAttributes());
3852 CI->setCallingConv(call.getCallingConv());
3853 CI->setTailCallKind(call.getTailCallKind());
3854 CI->setDebugLoc(dbgLoc);
3855
3856 if (funcName == "julia.gc_alloc_obj" ||
3857 funcName == "jl_gc_alloc_typed" ||
3858 funcName == "ijl_gc_alloc_typed") {
3859 if (EnzymeShadowAllocRewrite) {
3860 bool used = unnecessaryInstructions.find(&call) ==
3861 unnecessaryInstructions.end();
3862 EnzymeShadowAllocRewrite(wrap(CI), gutils, wrap(&call), idx,
3863 wrap(prev), used);
3864 }
3865 }
3866 idx++;
3867 prev = CI;
3868 return CI;
3869 };
3870
3871 Value *CI = applyChainRule(call.getType(), Builder2, rule);
3872
3873 auto found = gutils->invertedPointers.find(&call);
3874 PHINode *placeholder = cast<PHINode>(&*found->second);
3875
3876 gutils->invertedPointers.erase(found);
3877 gutils->replaceAWithB(placeholder, CI);
3878 gutils->erase(placeholder);
3879 gutils->invertedPointers.insert(
3880 std::make_pair(&call, InvertedPointerVH(gutils, CI)));
3881 }
3882 }
3883
3884 // Cache and rematerialization irrelevant for forward mode.
3885 if (Mode == DerivativeMode::ForwardMode ||
3886 Mode == DerivativeMode::ForwardModeError) {
3887 eraseIfUnused(call);
3888 return true;
3889 }
3890
3891 std::map<UsageKey, bool> Seen = gutils->populateSeenFromKnownRecompute();
3892 bool primalNeededInReverse =
3893 Mode == DerivativeMode::ForwardMode ||
3894 Mode == DerivativeMode::ForwardModeError
3895 ? false
3896 : DifferentialUseAnalysis::is_value_needed_in_reverse<
3897 QueryType::Primal>(gutils, &call, Mode, Seen, oldUnreachable);
3898
3899 // If we explicitly decided we need this in the reverse pass, mark it as
3900 // such.
3901 {
3902 auto found = gutils->knownRecomputeHeuristic.find(&call);
3903 if (found != gutils->knownRecomputeHeuristic.end() && !found->second) {
3904 primalNeededInReverse = true;
3905 }
3906 }
3907 bool cacheWholeAllocation = gutils->needsCacheWholeAllocation(&call);
3908 if (cacheWholeAllocation) {
3909 primalNeededInReverse = true;
3910 }
3911
3912 auto restoreFromStack = [&](MDNode *MD) {
3913 IRBuilder<> B(newCall);
3914 Value *Size;
3915 if (funcName == "malloc")
3916 Size = call.getArgOperand(0);
3917 else if (funcName == "julia.gc_alloc_obj" ||
3918 funcName == "jl_gc_alloc_typed" ||
3919 funcName == "ijl_gc_alloc_typed")
3920 Size = call.getArgOperand(1);
3921 else
3922 llvm_unreachable("Unknown allocation to upgrade");
3923 Size = gutils->getNewFromOriginal(Size);
3924
3925 if (isa<ConstantInt>(Size)) {
3926 B.SetInsertPoint(gutils->inversionAllocs);
3927 }
3928 Type *elTy = Type::getInt8Ty(call.getContext());
3929 if (MD->getNumOperands() == 2) {
3930 elTy = (Type *)cast<ConstantInt>(
3931 cast<ConstantAsMetadata>(MD->getOperand(1))->getValue())
3932 ->getLimitedValue();
3933 Value *tsize = ConstantInt::get(Size->getType(),
3934 (gutils->newFunc->getParent()
3935 ->getDataLayout()
3936 .getTypeAllocSizeInBits(elTy) +
3937 7) /
3938 8);
3939 Size = B.CreateUDiv(Size, tsize, "", /*exact*/ true);
3940 }
3941 Instruction *I = nullptr;
3942#if LLVM_VERSION_MAJOR < 17
3943 if (call.getContext().supportsTypedPointers()) {
3944 for (auto U : call.users()) {
3945 if (hasMetadata(cast<Instruction>(U), "enzyme_caststack")) {
3946 if (MD->getNumOperands() == 1) {
3947 elTy = U->getType()->getPointerElementType();
3948 Value *tsize = ConstantInt::get(
3949 Size->getType(), (gutils->newFunc->getParent()
3950 ->getDataLayout()
3951 .getTypeAllocSizeInBits(elTy) +
3952 7) /
3953 8);
3954
3955 Size = B.CreateUDiv(Size, tsize, "", /*exact*/ true);
3956 }
3957 I = gutils->getNewFromOriginal(cast<Instruction>(U));
3958 break;
3959 }
3960 }
3961 }
3962#endif
3963 Value *replacement = B.CreateAlloca(elTy, Size);
3964 for (auto MD : {"enzyme_active", "enzyme_inactive", "enzyme_type",
3965 "enzymejl_allocart", "enzymejl_allocart_name",
3966 "enzymejl_gc_alloc_rt"})
3967 if (auto M = call.getMetadata(MD))
3968 cast<AllocaInst>(replacement)->setMetadata(MD, M);
3969 if (I)
3970 replacement->takeName(I);
3971 else
3972 replacement->takeName(newCall);
3973 auto Alignment =
3974 cast<ConstantInt>(
3975 cast<ConstantAsMetadata>(MD->getOperand(0))->getValue())
3976 ->getLimitedValue();
3977 // Don't set zero alignment
3978 if (Alignment) {
3979 cast<AllocaInst>(replacement)->setAlignment(Align(Alignment));
3980 }
3981#if LLVM_VERSION_MAJOR < 17
3982 if (call.getContext().supportsTypedPointers()) {
3983 if (call.getType()->getPointerElementType() != elTy)
3984 replacement = B.CreatePointerCast(
3985 replacement, getUnqual(call.getType()->getPointerElementType()));
3986 }
3987#endif
3988 auto PT = cast<PointerType>(call.getType());
3989 if (PT->getAddressSpace()) {
3990 replacement = B.CreateAddrSpaceCast(replacement, PT);
3991 cast<Instruction>(replacement)
3992 ->setMetadata("enzyme_backstack",
3993 MDNode::get(replacement->getContext(), {}));
3994 }
3995 gutils->replaceAWithB(newCall, replacement);
3996 gutils->erase(newCall);
3997 };
3998
3999 // Don't erase any allocation that is being rematerialized.
4000 {
4001 auto found = gutils->rematerializableAllocations.find(&call);
4002 if (found != gutils->rematerializableAllocations.end()) {
4003 // If rematerializing (e.g. needed in reverse, but not needing
4004 // the whole allocation):
4005 if (gutils->allocationsToBeRematerialized.count(&call)) {
4006 assert(!unnecessaryValues.count(&call));
4007 // if rematerialize, don't ever cache and downgrade to stack
4008 // allocation where possible. Note that for allocations which are
4009 // within a loop, we will create the rematerialized allocation in the
4010 // rematerialied loop. Note that what matters here is whether the
4011 // actual call itself here is inside the loop, not whether the
4012 // rematerialization is loop level. This is because one can have a
4013 // loop level cache, but a function level allocation (e.g. for stack
4014 // allocas). If we deleted it here, we would have no allocation!
4015 auto AllocationLoop = gutils->OrigLI->getLoopFor(call.getParent());
4016 // An allocation within a loop, must definitionally be a loop level
4017 // allocation (but not always the other way around.
4018 if (AllocationLoop)
4019 assert(found->second.LI);
4020 if (auto MD = hasMetadata(&call, "enzyme_fromstack")) {
4021 if (Mode == DerivativeMode::ReverseModeGradient && AllocationLoop) {
4022 gutils->rematerializedPrimalOrShadowAllocations.push_back(
4023 newCall);
4024 } else {
4025 restoreFromStack(MD);
4026 }
4027 return true;
4028 }
4029
4030 // No need to free GC.
4031 if (EnzymeJuliaAddrLoad && isa<PointerType>(call.getType()) &&
4032 cast<PointerType>(call.getType())->getAddressSpace() == 10) {
4033 if (Mode == DerivativeMode::ReverseModeGradient && AllocationLoop)
4034 gutils->rematerializedPrimalOrShadowAllocations.push_back(
4035 newCall);
4036 return true;
4037 }
4038
4039 // Otherwise if in reverse pass, free the newly created allocation.
4040 if (Mode == DerivativeMode::ReverseModeGradient ||
4041 Mode == DerivativeMode::ReverseModeCombined ||
4042 Mode == DerivativeMode::ForwardModeSplit) {
4043 IRBuilder<> Builder2(&call);
4044 getReverseBuilder(Builder2);
4045 auto dbgLoc = gutils->getNewFromOriginal(call.getDebugLoc());
4046 auto freecall = freeKnownAllocation(
4047 Builder2, lookup(newCall, Builder2), funcName, dbgLoc,
4048 gutils->TLI, &call, gutils);
4049 if (freecall) {
4050 auto ident = MDNode::getDistinct(
4051 freecall->getContext(),
4052 {ConstantAsMetadata::get(
4053 ConstantInt::getTrue(freecall->getContext()))});
4054 newCall->setMetadata(
4055 "enzyme_cache_alloc",
4056 MDNode::get(freecall->getContext(), {ident}));
4057 freecall->setMetadata(
4058 "enzyme_cache_free",
4059 MDNode::get(freecall->getContext(), {ident}));
4060 }
4061 if (Mode == DerivativeMode::ReverseModeGradient && AllocationLoop)
4062 gutils->rematerializedPrimalOrShadowAllocations.push_back(
4063 newCall);
4064 return true;
4065 }
4066 // If in primal, do nothing (keeping the original caching behavior)
4067 if (Mode == DerivativeMode::ReverseModePrimal)
4068 return true;
4069 } else if (!cacheWholeAllocation) {
4070 if (unnecessaryValues.count(&call)) {
4071 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4072 return true;
4073 }
4074 // If not caching allocation and not needed in the reverse, we can
4075 // use the original freeing behavior for the function. If in the
4076 // reverse pass we should not recreate this allocation.
4077 if (Mode == DerivativeMode::ReverseModeGradient)
4078 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4079 else if (auto MD = hasMetadata(&call, "enzyme_fromstack")) {
4080 restoreFromStack(MD);
4081 }
4082 return true;
4083 }
4084 }
4085 }
4086
4087 // If an allocation is not needed in the reverse, maintain the original
4088 // free behavior and do not rematerialize this for the reverse. However,
4089 // this is only safe to perform for allocations with a guaranteed free
4090 // as can we can only guarantee that we don't erase those frees.
4091 bool hasPDFree = gutils->allocationsWithGuaranteedFree.count(&call);
4092 if (!primalNeededInReverse && hasPDFree) {
4093 if (unnecessaryValues.count(&call)) {
4094 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4095 return true;
4096 }
4097 if (Mode == DerivativeMode::ReverseModeGradient ||
4098 Mode == DerivativeMode::ForwardModeSplit) {
4099 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4100 } else {
4101 if (auto MD = hasMetadata(&call, "enzyme_fromstack")) {
4102 restoreFromStack(MD);
4103 }
4104 }
4105 return true;
4106 }
4107
4108 // If an object is managed by the GC do not preserve it for later free,
4109 // Thus it only needs caching if there is a need for it in the reverse.
4110 if (EnzymeJuliaAddrLoad && isa<PointerType>(call.getType()) &&
4111 cast<PointerType>(call.getType())->getAddressSpace() == 10) {
4112 if (!subretused) {
4113 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4114 return true;
4115 }
4116 if (!primalNeededInReverse) {
4117 if (Mode == DerivativeMode::ReverseModeGradient ||
4118 Mode == DerivativeMode::ForwardModeSplit) {
4119 auto pn = BuilderZ.CreatePHI(call.getType(), 1,
4120 call.getName() + "_replacementJ");
4121 gutils->fictiousPHIs[pn] = &call;
4122 gutils->replaceAWithB(newCall, pn);
4123 gutils->erase(newCall);
4124 }
4125 } else if (Mode != DerivativeMode::ReverseModeCombined) {
4126 gutils->cacheForReverse(BuilderZ, newCall,
4127 getIndex(&call, CacheType::Self, BuilderZ));
4128 }
4129 return true;
4130 }
4131
4132 if (EnzymeFreeInternalAllocations)
4133 hasPDFree = true;
4134
4135 // TODO enable this if we need to free the memory
4136 // NOTE THAT TOPLEVEL IS THERE SIMPLY BECAUSE THAT WAS PREVIOUS ATTITUTE
4137 // TO FREE'ing
4138 if ((primalNeededInReverse &&
4139 !gutils->unnecessaryIntermediates.count(&call)) ||
4140 hasPDFree) {
4141 if (hasPDFree && Mode == DerivativeMode::ReverseModePrimal) {
4142 auto ident =
4143 MDNode::getDistinct(newCall->getContext(),
4144 {ConstantAsMetadata::get(ConstantInt::getFalse(
4145 newCall->getContext()))});
4146 newCall->setMetadata("enzyme_cache_alloc",
4147 MDNode::get(newCall->getContext(), {ident}));
4148 }
4149 Value *nop = gutils->cacheForReverse(
4150 BuilderZ, newCall, getIndex(&call, CacheType::Self, BuilderZ));
4151 if (hasPDFree &&
4152 ((Mode == DerivativeMode::ReverseModeGradient && shouldFree()) ||
4153 Mode == DerivativeMode::ReverseModeCombined ||
4154 (Mode == DerivativeMode::ForwardModeSplit && shouldFree()))) {
4155 IRBuilder<> Builder2(&call);
4156 getReverseBuilder(Builder2);
4157 auto dbgLoc = gutils->getNewFromOriginal(call.getDebugLoc());
4158 auto freecall =
4159 freeKnownAllocation(Builder2, lookup(nop, Builder2), funcName,
4160 dbgLoc, gutils->TLI, &call, gutils);
4161 if (freecall) {
4162 bool combined = Mode == DerivativeMode::ReverseModeCombined;
4163 auto ident = MDNode::getDistinct(
4164 freecall->getContext(),
4165 {ConstantAsMetadata::get(
4166 combined ? ConstantInt::getTrue(freecall->getContext())
4167 : ConstantInt::getFalse(freecall->getContext()))});
4168 if (combined)
4169 newCall->setMetadata("enzyme_cache_alloc",
4170 MDNode::get(freecall->getContext(), {ident}));
4171 freecall->setMetadata("enzyme_cache_free",
4172 MDNode::get(freecall->getContext(), {ident}));
4173 }
4174 }
4175 } else if (Mode == DerivativeMode::ReverseModeGradient ||
4176 Mode == DerivativeMode::ReverseModeCombined ||
4177 Mode == DerivativeMode::ForwardModeSplit) {
4178 // Note that here we cannot simply replace with null as users who
4179 // try to find the shadow pointer will use the shadow of null rather
4180 // than the true shadow of this
4181 auto pn = BuilderZ.CreatePHI(call.getType(), 1,
4182 call.getName() + "_replacementB");
4183 gutils->fictiousPHIs[pn] = &call;
4184 gutils->replaceAWithB(newCall, pn);
4185 gutils->erase(newCall);
4186 }
4187
4188 return true;
4189 }
4190
4191 if (funcName == "julia.gc_loaded") {
4192 if (gutils->isConstantValue(&call)) {
4193 eraseIfUnused(call);
4194 return true;
4195 }
4196 auto ifound = gutils->invertedPointers.find(&call);
4197 assert(ifound != gutils->invertedPointers.end());
4198
4199 if (auto placeholder = dyn_cast<PHINode>(&*ifound->second)) {
4200
4201 bool needShadow = DifferentialUseAnalysis::is_value_needed_in_reverse<
4202 QueryType::Shadow>(gutils, &call, Mode, oldUnreachable);
4203 if (!needShadow) {
4204 gutils->invertedPointers.erase(ifound);
4205 gutils->erase(placeholder);
4206 eraseIfUnused(call);
4207 return true;
4208 }
4209
4210 gutils->invertedPointers.erase(ifound);
4211 auto res = gutils->invertPointerM(&call, BuilderZ);
4212
4213 gutils->replaceAWithB(placeholder, res);
4214 gutils->erase(placeholder);
4215 }
4216 eraseIfUnused(call);
4217
4218 return true;
4219 }
4220
4221 if (funcName == "julia.pointer_from_objref") {
4222 if (gutils->isConstantValue(&call)) {
4223 eraseIfUnused(call);
4224 return true;
4225 }
4226
4227 auto ifound = gutils->invertedPointers.find(&call);
4228 assert(ifound != gutils->invertedPointers.end());
4229
4230 auto placeholder = cast<PHINode>(&*ifound->second);
4231
4232 bool needShadow =
4233 DifferentialUseAnalysis::is_value_needed_in_reverse<QueryType::Shadow>(
4234 gutils, &call, Mode, oldUnreachable);
4235 if (!needShadow) {
4236 gutils->invertedPointers.erase(ifound);
4237 gutils->erase(placeholder);
4238 eraseIfUnused(call);
4239 return true;
4240 }
4241
4242 Value *ptrshadow = gutils->invertPointerM(call.getArgOperand(0), BuilderZ);
4243
4244 Value *val = applyChainRule(
4245 call.getType(), BuilderZ,
4246 [&](Value *v) -> Value * { return BuilderZ.CreateCall(called, {v}); },
4247 ptrshadow);
4248
4249 gutils->replaceAWithB(placeholder, val);
4250 gutils->erase(placeholder);
4251 eraseIfUnused(call);
4252 return true;
4253 }
4254 if (funcName.contains("__enzyme_todense")) {
4255 if (gutils->isConstantValue(&call)) {
4256 eraseIfUnused(call);
4257 return true;
4258 }
4259
4260 auto ifound = gutils->invertedPointers.find(&call);
4261 assert(ifound != gutils->invertedPointers.end());
4262
4263 auto placeholder = cast<PHINode>(&*ifound->second);
4264
4265 bool needShadow =
4266 DifferentialUseAnalysis::is_value_needed_in_reverse<QueryType::Shadow>(
4267 gutils, &call, Mode, oldUnreachable);
4268 if (!needShadow) {
4269 gutils->invertedPointers.erase(ifound);
4270 gutils->erase(placeholder);
4271 eraseIfUnused(call);
4272 return true;
4273 }
4274
4275 // The shadow pointer is built from the shadows of the extra arguments.
4276 // An inactive argument (e.g. an index) is passed as is to every lane.
4277 SmallVector<Value *, 3> args;
4278 SmallVector<bool, 3> argIsShadow;
4279 for (size_t i = 0; i < 2; i++) {
4280 args.push_back(gutils->getNewFromOriginal(call.getArgOperand(i)));
4281 argIsShadow.push_back(false);
4282 }
4283 for (size_t i = 2; i < call.arg_size(); ++i) {
4284 auto arg = call.getArgOperand(i);
4285 if (gutils->isConstantValue(arg)) {
4286 args.push_back(gutils->getNewFromOriginal(arg));
4287 argIsShadow.push_back(false);
4288 } else {
4289 args.push_back(gutils->invertPointerM(arg, BuilderZ));
4290 argIsShadow.push_back(true);
4291 }
4292 }
4293
4294 Value *res = UndefValue::get(gutils->getShadowType(call.getType()));
4295 if (gutils->getWidth() == 1) {
4296 res = BuilderZ.CreateCall(called, args);
4297 } else {
4298 for (size_t w = 0; w < gutils->getWidth(); ++w) {
4299 SmallVector<Value *, 3> targs = {args[0], args[1]};
4300 for (size_t i = 2; i < call.arg_size(); ++i)
4301 targs.push_back(argIsShadow[i]
4302 ? GradientUtils::extractMeta(BuilderZ, args[i], w)
4303 : args[i]);
4304
4305 auto tres = BuilderZ.CreateCall(called, targs);
4306 res = BuilderZ.CreateInsertValue(res, tres, w);
4307 }
4308 }
4309
4310 gutils->replaceAWithB(placeholder, res);
4311 gutils->erase(placeholder);
4312 eraseIfUnused(call);
4313 return true;
4314 }
4315
4316 if (funcName == "memcpy" || funcName == "memmove") {
4317 auto ID = (funcName == "memcpy") ? Intrinsic::memcpy : Intrinsic::memmove;
4318 visitMemTransferCommon(ID, /*srcAlign*/ MaybeAlign(1),
4319 /*dstAlign*/ MaybeAlign(1), call,
4320 call.getArgOperand(0), call.getArgOperand(1),
4321 gutils->getNewFromOriginal(call.getArgOperand(2)),
4322 ConstantInt::getFalse(call.getContext()));
4323 return true;
4324 }
4325 if (funcName == "memset" || funcName == "memset_pattern16" ||
4326 funcName == "__memset_chk") {
4327 visitMemSetCommon(call);
4328 return true;
4329 }
4330 if (funcName == "enzyme_zerotype") {
4331 IRBuilder<> BuilderZ(&call);
4332 getForwardBuilder(BuilderZ);
4333
4334 bool backwardsShadow = false;
4335 bool forwardsShadow = true;
4336 for (auto pair : gutils->backwardsOnlyShadows) {
4337 if (pair.second.stores.count(&call)) {
4338 backwardsShadow = true;
4339 forwardsShadow = pair.second.primalInitialize;
4340 if (auto inst = dyn_cast<Instruction>(pair.first))
4341 if (!forwardsShadow && pair.second.LI &&
4342 pair.second.LI->contains(inst->getParent()))
4343 backwardsShadow = false;
4344 }
4345 }
4346
4347 bool forceErase =
4348 !((Mode == DerivativeMode::ReverseModePrimal && forwardsShadow) ||
4349 (Mode == DerivativeMode::ReverseModeCombined && forwardsShadow) ||
4350 (Mode == DerivativeMode::ReverseModeGradient && backwardsShadow) ||
4351 (Mode == DerivativeMode::ForwardModeSplit && backwardsShadow));
4352
4353 if (forceErase)
4354 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4355 else
4356 eraseIfUnused(call);
4357
4358 Value *orig_op0 = call.getArgOperand(0);
4359
4360 // If constant destination then no operation needs doing
4361 if (gutils->isConstantValue(orig_op0)) {
4362 return true;
4363 }
4364
4365 if (!forceErase) {
4366 Value *op0 = gutils->invertPointerM(orig_op0, BuilderZ);
4367 Value *op1 = gutils->getNewFromOriginal(call.getArgOperand(1));
4368 Value *op2 = gutils->getNewFromOriginal(call.getArgOperand(2));
4369 auto Defs = gutils->getInvertedBundles(
4370 &call, {ValueType::Shadow, ValueType::Primal, ValueType::Primal},
4371 BuilderZ, /*lookup*/ false);
4372
4373 applyChainRule(
4374 BuilderZ,
4375 [&](Value *op0) {
4376 SmallVector<Value *, 4> args = {op0, op1, op2};
4377 auto cal =
4378 BuilderZ.CreateCall(call.getCalledFunction(), args, Defs);
4379 llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
4380 ToCopy2.push_back(LLVMContext::MD_noalias);
4381 cal->copyMetadata(call, ToCopy2);
4382 cal->setAttributes(call.getAttributes());
4383 if (auto m = hasMetadata(&call, "enzyme_zerostack"))
4384 cal->setMetadata("enzyme_zerostack", m);
4385 cal->setCallingConv(call.getCallingConv());
4386 cal->setTailCallKind(call.getTailCallKind());
4387 cal->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
4388 },
4389 op0);
4390 }
4391 return true;
4392 }
4393 if (funcName == "cuStreamCreate") {
4394 Value *val = nullptr;
4395 llvm::Type *PT = getInt8PtrTy(call.getContext());
4396#if LLVM_VERSION_MAJOR < 17
4397 if (call.getContext().supportsTypedPointers()) {
4398 if (isa<PointerType>(call.getArgOperand(0)->getType()))
4399 PT = call.getArgOperand(0)->getType()->getPointerElementType();
4400 }
4401#endif
4402 if (Mode == DerivativeMode::ReverseModePrimal ||
4403 Mode == DerivativeMode::ReverseModeCombined) {
4404 val = gutils->getNewFromOriginal(call.getOperand(0));
4405 if (!isa<PointerType>(val->getType()))
4406 val = BuilderZ.CreateIntToPtr(val, getUnqual(PT));
4407 val = BuilderZ.CreateLoad(PT, val);
4408 val = gutils->cacheForReverse(BuilderZ, val,
4409 getIndex(&call, CacheType::Tape, BuilderZ));
4410
4411 } else if (Mode == DerivativeMode::ReverseModeGradient) {
4412 PHINode *toReplace =
4413 BuilderZ.CreatePHI(PT, 1, call.getName() + "_psxtmp");
4414 val = gutils->cacheForReverse(BuilderZ, toReplace,
4415 getIndex(&call, CacheType::Tape, BuilderZ));
4416 }
4417 if (Mode == DerivativeMode::ReverseModeGradient ||
4418 Mode == DerivativeMode::ReverseModeCombined) {
4419 IRBuilder<> Builder2(&call);
4420 getReverseBuilder(Builder2);
4421 val = gutils->lookupM(val, Builder2);
4422 auto FreeFunc = getOrInsertPerCallingConv(
4423 *gutils->newFunc->getParent(), called, "cuStreamDestroy",
4424 FunctionType::get(call.getType(), {PT}, false));
4425 Value *nargs[] = {val};
4426 Builder2.CreateCall(FreeFunc, nargs);
4427 }
4428 if (Mode == DerivativeMode::ReverseModeGradient)
4429 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4430 return true;
4431 }
4432 if (funcName == "cuStreamDestroy") {
4433 if (Mode == DerivativeMode::ReverseModeGradient ||
4434 Mode == DerivativeMode::ReverseModeCombined)
4435 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4436 return true;
4437 }
4438 if (funcName == "cuStreamSynchronize") {
4439 if (Mode == DerivativeMode::ReverseModeGradient ||
4440 Mode == DerivativeMode::ReverseModeCombined) {
4441 IRBuilder<> Builder2(&call);
4442 getReverseBuilder(Builder2);
4443 Value *nargs[] = {gutils->lookupM(
4444 gutils->getNewFromOriginal(call.getOperand(0)), Builder2)};
4445 auto callval = call.getCalledOperand();
4446 Builder2.CreateCall(call.getFunctionType(), callval, nargs);
4447 }
4448 if (Mode == DerivativeMode::ReverseModeGradient)
4449 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4450 return true;
4451 }
4452 if (funcName == "posix_memalign" || funcName == "cuMemAllocAsync" ||
4453 funcName == "cuMemAlloc" || funcName == "cuMemAlloc_v2" ||
4454 funcName == "cudaMalloc" || funcName == "cudaMallocAsync" ||
4455 funcName == "cudaMallocHost" || funcName == "cudaMallocFromPoolAsync") {
4456 bool constval = gutils->isConstantInstruction(&call);
4457
4458 Value *val;
4459 llvm::Type *PT = getInt8PtrTy(call.getContext());
4460#if LLVM_VERSION_MAJOR < 17
4461 if (call.getContext().supportsTypedPointers()) {
4462 if (isa<PointerType>(call.getArgOperand(0)->getType()))
4463 PT = call.getArgOperand(0)->getType()->getPointerElementType();
4464 }
4465#endif
4466 if (!constval) {
4467 Value *stream = nullptr;
4468 if (funcName == "cuMemAllocAsync")
4469 stream = gutils->getNewFromOriginal(call.getArgOperand(2));
4470 else if (funcName == "cudaMallocAsync")
4471 stream = gutils->getNewFromOriginal(call.getArgOperand(2));
4472 else if (funcName == "cudaMallocFromPoolAsync")
4473 stream = gutils->getNewFromOriginal(call.getArgOperand(3));
4474
4475 auto M = gutils->newFunc->getParent();
4476
4477 if (Mode == DerivativeMode::ReverseModePrimal ||
4478 Mode == DerivativeMode::ReverseModeCombined ||
4479 Mode == DerivativeMode::ForwardMode ||
4480 Mode == DerivativeMode::ForwardModeError) {
4481 Value *ptrshadow =
4482 gutils->invertPointerM(call.getArgOperand(0), BuilderZ);
4483 SmallVector<Value *, 1> args;
4484 SmallVector<ValueType, 1> valtys;
4485 args.push_back(ptrshadow);
4486 valtys.push_back(ValueType::Shadow);
4487 for (size_t i = 1; i < call.arg_size(); ++i) {
4488 args.push_back(gutils->getNewFromOriginal(call.getArgOperand(i)));
4489 valtys.push_back(ValueType::Primal);
4490 }
4491
4492 auto Defs = gutils->getInvertedBundles(&call, valtys, BuilderZ,
4493 /*lookup*/ false);
4494
4495 val = applyChainRule(
4496 PT, BuilderZ,
4497 [&](Value *ptrshadow) {
4498 args[0] = ptrshadow;
4499
4500 BuilderZ.CreateCall(called, args, Defs);
4501 if (!isa<PointerType>(ptrshadow->getType()))
4502 ptrshadow = BuilderZ.CreateIntToPtr(ptrshadow, getUnqual(PT));
4503 Value *val = BuilderZ.CreateLoad(PT, ptrshadow);
4504
4505 auto dst_arg =
4506 BuilderZ.CreateBitCast(val, getInt8PtrTy(call.getContext()));
4507
4508 auto val_arg =
4509 ConstantInt::get(Type::getInt8Ty(call.getContext()), 0);
4510 auto len_arg = gutils->getNewFromOriginal(
4511 call.getArgOperand((funcName == "posix_memalign") ? 2 : 1));
4512
4513 if (funcName == "posix_memalign" ||
4514 funcName == "cudaMallocHost") {
4515 BuilderZ.CreateMemSet(dst_arg, val_arg, len_arg, MaybeAlign());
4516 } else if (funcName == "cudaMalloc") {
4517 Type *tys[] = {PT, val_arg->getType(), len_arg->getType()};
4518 auto F = getOrInsertPerCallingConv(
4519 *M, called, "cudaMemset",
4520 FunctionType::get(call.getType(), tys, false));
4521 Value *nargs[] = {dst_arg, val_arg, len_arg};
4522 auto memset = cast<CallInst>(BuilderZ.CreateCall(F, nargs));
4523 memset->addParamAttr(0, Attribute::NonNull);
4524 } else if (funcName == "cudaMallocAsync" ||
4525 funcName == "cudaMallocFromPoolAsync") {
4526 Type *tys[] = {PT, val_arg->getType(), len_arg->getType(),
4527 stream->getType()};
4528 auto F = getOrInsertPerCallingConv(
4529 *M, called, "cudaMemsetAsync",
4530 FunctionType::get(call.getType(), tys, false));
4531 Value *nargs[] = {dst_arg, val_arg, len_arg, stream};
4532 auto memset = cast<CallInst>(BuilderZ.CreateCall(F, nargs));
4533 memset->addParamAttr(0, Attribute::NonNull);
4534 } else if (funcName == "cuMemAllocAsync") {
4535 Type *tys[] = {PT, val_arg->getType(), len_arg->getType(),
4536 stream->getType()};
4537 auto F = getOrInsertPerCallingConv(
4538 *M, called, "cuMemsetD8Async",
4539 FunctionType::get(call.getType(), tys, false));
4540 Value *nargs[] = {dst_arg, val_arg, len_arg, stream};
4541 auto memset = cast<CallInst>(BuilderZ.CreateCall(F, nargs));
4542 memset->addParamAttr(0, Attribute::NonNull);
4543 } else if (funcName == "cuMemAlloc" ||
4544 funcName == "cuMemAlloc_v2") {
4545 Type *tys[] = {PT, val_arg->getType(), len_arg->getType()};
4546 auto F = getOrInsertPerCallingConv(
4547 *M, called,
4548 funcName == "cuMemAlloc_v2" ? "cuMemsetD8_v2"
4549 : "cuMemsetD8",
4550 FunctionType::get(call.getType(), tys, false));
4551 Value *nargs[] = {dst_arg, val_arg, len_arg};
4552 auto memset = cast<CallInst>(BuilderZ.CreateCall(F, nargs));
4553 memset->addParamAttr(0, Attribute::NonNull);
4554 } else {
4555 llvm_unreachable("unhandled allocation");
4556 }
4557 return val;
4558 },
4559 ptrshadow);
4560
4561 if (Mode != DerivativeMode::ForwardMode &&
4562 Mode != DerivativeMode::ForwardModeError)
4563 val = gutils->cacheForReverse(
4564 BuilderZ, val, getIndex(&call, CacheType::Tape, BuilderZ));
4565 } else if (Mode == DerivativeMode::ReverseModeGradient) {
4566 PHINode *toReplace = BuilderZ.CreatePHI(gutils->getShadowType(PT), 1,
4567 call.getName() + "_psxtmp");
4568 val = gutils->cacheForReverse(
4569 BuilderZ, toReplace, getIndex(&call, CacheType::Tape, BuilderZ));
4570 }
4571
4572 if (Mode == DerivativeMode::ReverseModeCombined ||
4573 Mode == DerivativeMode::ReverseModeGradient) {
4574 if (shouldFree()) {
4575 IRBuilder<> Builder2(&call);
4576 getReverseBuilder(Builder2);
4577 Value *tofree = gutils->lookupM(val, Builder2, ValueToValueMapTy(),
4578 /*tryLegalRecompute*/ false);
4579
4580 Type *VoidTy = Type::getVoidTy(M->getContext());
4581 Type *IntPtrTy = getInt8PtrTy(M->getContext());
4582
4583 Value *streamL = nullptr;
4584 if (stream)
4585 streamL = gutils->lookupM(stream, Builder2);
4586
4587 applyChainRule(
4588 BuilderZ,
4589 [&](Value *tofree) {
4590 if (funcName == "posix_memalign") {
4591 auto FreeFunc =
4592 M->getOrInsertFunction("free", VoidTy, IntPtrTy);
4593 Builder2.CreateCall(FreeFunc, tofree);
4594 } else if (funcName == "cuMemAllocAsync") {
4595 auto FreeFunc = getOrInsertPerCallingConv(
4596 *M, called, "cuMemFreeAsync",
4597 FunctionType::get(VoidTy, {IntPtrTy, streamL->getType()},
4598 false));
4599 Value *nargs[] = {tofree, streamL};
4600 Builder2.CreateCall(FreeFunc, nargs);
4601 } else if (funcName == "cuMemAlloc" ||
4602 funcName == "cuMemAlloc_v2") {
4603 auto FreeFunc = getOrInsertPerCallingConv(
4604 *M, called, "cuMemFree",
4605 FunctionType::get(VoidTy, {IntPtrTy}, false));
4606 Value *nargs[] = {tofree};
4607 Builder2.CreateCall(FreeFunc, nargs);
4608 } else if (funcName == "cudaMalloc") {
4609 auto FreeFunc = getOrInsertPerCallingConv(
4610 *M, called, "cudaFree",
4611 FunctionType::get(VoidTy, {IntPtrTy}, false));
4612 Value *nargs[] = {tofree};
4613 Builder2.CreateCall(FreeFunc, nargs);
4614 } else if (funcName == "cudaMallocAsync" ||
4615 funcName == "cudaMallocFromPoolAsync") {
4616 auto FreeFunc = getOrInsertPerCallingConv(
4617 *M, called, "cudaFreeAsync",
4618 FunctionType::get(VoidTy, {IntPtrTy, streamL->getType()},
4619 false));
4620 Value *nargs[] = {tofree, streamL};
4621 Builder2.CreateCall(FreeFunc, nargs);
4622 } else if (funcName == "cudaMallocHost") {
4623 auto FreeFunc = getOrInsertPerCallingConv(
4624 *M, called, "cudaFreeHost",
4625 FunctionType::get(VoidTy, {IntPtrTy}, false));
4626 Value *nargs[] = {tofree};
4627 Builder2.CreateCall(FreeFunc, nargs);
4628 } else
4629 llvm_unreachable("unknown function to free");
4630 },
4631 tofree);
4632 }
4633 }
4634 }
4635
4636 // TODO enable this if we need to free the memory
4637 // NOTE THAT TOPLEVEL IS THERE SIMPLY BECAUSE THAT WAS PREVIOUS ATTITUTE
4638 // TO FREE'ing
4639 if (Mode == DerivativeMode::ReverseModeGradient) {
4640 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4641 } else if (Mode == DerivativeMode::ReverseModePrimal) {
4642 // if (is_value_needed_in_reverse<Primal>(
4643 // TR, gutils, orig, /*topLevel*/ Mode ==
4644 // DerivativeMode::Both))
4645 // {
4646
4647 // gutils->cacheForReverse(BuilderZ, newCall,
4648 // getIndex(orig, CacheType::Self, BuilderZ));
4649 //} else if (Mode != DerivativeMode::Forward) {
4650 // Note that here we cannot simply replace with null as users who try
4651 // to find the shadow pointer will use the shadow of null rather than
4652 // the true shadow of this
4653 //}
4654 } else if (Mode == DerivativeMode::ReverseModeCombined && shouldFree()) {
4655 IRBuilder<> Builder2(newCall->getNextNode());
4656 auto ptrv = gutils->getNewFromOriginal(call.getOperand(0));
4657 if (!isa<PointerType>(ptrv->getType()))
4658 ptrv = BuilderZ.CreateIntToPtr(ptrv, getUnqual(PT));
4659 auto load = Builder2.CreateLoad(PT, ptrv, "posix_preread");
4660 Builder2.SetInsertPoint(&call);
4661 getReverseBuilder(Builder2);
4662 auto tofree = gutils->lookupM(load, Builder2, ValueToValueMapTy(),
4663 /*tryLegal*/ false);
4664 Value *streamL = nullptr;
4665 if (funcName == "cuMemAllocAsync")
4666 streamL = gutils->getNewFromOriginal(call.getArgOperand(2));
4667 else if (funcName == "cudaMallocAsync")
4668 streamL = gutils->getNewFromOriginal(call.getArgOperand(2));
4669 else if (funcName == "cudaMallocFromPoolAsync")
4670 streamL = gutils->getNewFromOriginal(call.getArgOperand(3));
4671 if (streamL)
4672 streamL = gutils->lookupM(streamL, Builder2);
4673
4674 auto M = gutils->newFunc->getParent();
4675 Type *VoidTy = Type::getVoidTy(M->getContext());
4676 Type *IntPtrTy = getInt8PtrTy(M->getContext());
4677
4678 if (funcName == "posix_memalign") {
4679 auto FreeFunc = M->getOrInsertFunction("free", VoidTy, IntPtrTy);
4680 Builder2.CreateCall(FreeFunc, tofree);
4681 } else if (funcName == "cuMemAllocAsync") {
4682 auto FreeFunc = getOrInsertPerCallingConv(
4683 *M, called, "cuMemFreeAsync",
4684 FunctionType::get(VoidTy, {IntPtrTy, streamL->getType()}, false));
4685 Value *nargs[] = {tofree, streamL};
4686 Builder2.CreateCall(FreeFunc, nargs);
4687 } else if (funcName == "cuMemAlloc" || funcName == "cuMemAlloc_v2") {
4688 auto FreeFunc = getOrInsertPerCallingConv(
4689 *M, called, "cuMemFree",
4690 FunctionType::get(VoidTy, {IntPtrTy}, false));
4691 Value *nargs[] = {tofree};
4692 Builder2.CreateCall(FreeFunc, nargs);
4693 } else if (funcName == "cudaMalloc") {
4694 auto FreeFunc = getOrInsertPerCallingConv(
4695 *M, called, "cudaFree",
4696 FunctionType::get(VoidTy, {IntPtrTy}, false));
4697 Value *nargs[] = {tofree};
4698 Builder2.CreateCall(FreeFunc, nargs);
4699 } else if (funcName == "cudaMallocAsync" ||
4700 funcName == "cudaMallocFromPoolAsync") {
4701 auto FreeFunc = getOrInsertPerCallingConv(
4702 *M, called, "cudaFreeAsync",
4703 FunctionType::get(VoidTy, {IntPtrTy, streamL->getType()}, false));
4704 Value *nargs[] = {tofree, streamL};
4705 Builder2.CreateCall(FreeFunc, nargs);
4706 } else if (funcName == "cudaMallocHost") {
4707 auto FreeFunc = getOrInsertPerCallingConv(
4708 *M, called, "cudaFreeHost",
4709 FunctionType::get(VoidTy, {IntPtrTy}, false));
4710 Value *nargs[] = {tofree};
4711 Builder2.CreateCall(FreeFunc, nargs);
4712 } else
4713 llvm_unreachable("unknown function to free");
4714 }
4715
4716 return true;
4717 }
4718
4719 // Remove free's in forward pass so the memory can be used in the reverse
4720 // pass
4721 if (isDeallocationFunction(funcName, gutils->TLI)) {
4722 assert(gutils->invertedPointers.find(&call) ==
4723 gutils->invertedPointers.end());
4724
4725 if (Mode == DerivativeMode::ForwardMode ||
4726 Mode == DerivativeMode::ForwardModeError) {
4727 if (!gutils->isConstantValue(call.getArgOperand(0))) {
4728 IRBuilder<> Builder2(&call);
4729 getForwardBuilder(Builder2);
4730 auto origfree = call.getArgOperand(0);
4731 auto newfree = gutils->getNewFromOriginal(call.getArgOperand(0));
4732 auto tofree = gutils->invertPointerM(origfree, Builder2);
4733
4734 Function *free = getOrInsertCheckedFree(
4735 *call.getModule(), &call, newfree->getType(), gutils->getWidth());
4736
4737 bool used = true;
4738 if (auto instArg = dyn_cast<Instruction>(call.getArgOperand(0)))
4739 used = unnecessaryInstructions.find(instArg) ==
4740 unnecessaryInstructions.end();
4741
4742 SmallVector<Value *, 3> args;
4743 if (used)
4744 args.push_back(newfree);
4745 else
4746 args.push_back(
4747 Constant::getNullValue(call.getArgOperand(0)->getType()));
4748
4749 auto rule = [&args](Value *tofree) { args.push_back(tofree); };
4750 applyChainRule(Builder2, rule, tofree);
4751
4752 for (size_t i = 1; i < call.arg_size(); i++) {
4753 args.push_back(gutils->getNewFromOriginal(call.getArgOperand(i)));
4754 }
4755
4756 auto frees = Builder2.CreateCall(free->getFunctionType(), free, args);
4757 frees->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
4758
4759 eraseIfUnused(call);
4760 return true;
4761 }
4762 eraseIfUnused(call);
4763 }
4764 auto callval = call.getCalledOperand();
4765
4766 for (auto rmat : gutils->backwardsOnlyShadows) {
4767 if (gutils->allocationsToBeRematerialized.count(rmat.first) &&
4768 rmat.second.frees.count(&call)) {
4769 bool shouldFree = false;
4770 if (rmat.second.primalInitialize) {
4771 if (Mode == DerivativeMode::ReverseModePrimal)
4772 shouldFree = true;
4773 }
4774
4775 if (shouldFree) {
4776 IRBuilder<> Builder2(&call);
4777 getForwardBuilder(Builder2);
4778 auto origfree = call.getArgOperand(0);
4779 auto tofree = gutils->invertPointerM(origfree, Builder2);
4780 if (tofree != origfree) {
4781 SmallVector<Value *, 2> args = {tofree};
4782 CallInst *CI =
4783 Builder2.CreateCall(call.getFunctionType(), callval, args);
4784 CI->setAttributes(call.getAttributes());
4785 }
4786 }
4787 break;
4788 }
4789 }
4790
4791 // If a rematerializable allocation.
4792 for (auto rmat : gutils->rematerializableAllocations) {
4793 if (gutils->allocationsToBeRematerialized.count(rmat.first) &&
4794 rmat.second.frees.count(&call)) {
4795 // Leave the original free behavior since this won't be used
4796 // in the reverse pass in split mode
4797 if (Mode == DerivativeMode::ReverseModePrimal) {
4798 eraseIfUnused(call);
4799 return true;
4800 } else if (Mode == DerivativeMode::ReverseModeGradient) {
4801 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4802 return true;
4803 } else {
4804 assert(Mode == DerivativeMode::ReverseModeCombined);
4805 std::map<UsageKey, bool> Seen =
4806 gutils->populateSeenFromKnownRecompute();
4807 bool primalNeededInReverse =
4808 DifferentialUseAnalysis::is_value_needed_in_reverse<
4809 QueryType::Primal>(gutils, rmat.first, Mode, Seen,
4810 oldUnreachable);
4811 {
4812 auto found = gutils->knownRecomputeHeuristic.find(rmat.first);
4813 if (found != gutils->knownRecomputeHeuristic.end()) {
4814 if (!found->second) {
4815 primalNeededInReverse = true;
4816 }
4817 }
4818 }
4819 bool cacheWholeAllocation =
4820 gutils->needsCacheWholeAllocation(rmat.first);
4821 if (cacheWholeAllocation) {
4822 primalNeededInReverse = true;
4823 }
4824 // If in a loop context, maintain the same free behavior, unless
4825 // caching whole allocation.
4826 if (!cacheWholeAllocation) {
4827 eraseIfUnused(call);
4828 return true;
4829 }
4830 assert(!unnecessaryValues.count(rmat.first));
4831 (void)primalNeededInReverse;
4832 assert(primalNeededInReverse);
4833 }
4834 }
4835 }
4836
4837 if (gutils->forwardDeallocations.count(&call)) {
4838 if (Mode == DerivativeMode::ReverseModeGradient) {
4839 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4840 } else
4841 eraseIfUnused(call);
4842 return true;
4843 }
4844
4845 if (gutils->postDominatingFrees.count(&call)) {
4846 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4847 return true;
4848 }
4849
4850 if (call.getMetadata("enzyme_cache_free")) {
4851 bool hasGuaranteedFree = false;
4852 for (const auto &pair : gutils->allocationsWithGuaranteedFree) {
4853 if (pair.second.count(&call)) {
4854 hasGuaranteedFree = true;
4855 break;
4856 }
4857 }
4858 if (!hasGuaranteedFree) {
4859 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4860 return true;
4861 }
4862 }
4863
4864 llvm::Value *val = getBaseObject(call.getArgOperand(0));
4865 if (isa<ConstantPointerNull>(val)) {
4866 llvm::errs() << "removing free of null pointer\n";
4867 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4868 return true;
4869 }
4870
4871 // TODO HANDLE FREE
4872 llvm::errs() << "freeing without malloc " << *val << " in " << call << "\n";
4873 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4874 return true;
4875 }
4876
4877 if (call.hasFnAttr("enzyme_sample")) {
4878 if (Mode != DerivativeMode::ReverseModeCombined &&
4879 Mode != DerivativeMode::ReverseModeGradient)
4880 return true;
4881
4882 bool constval = gutils->isConstantInstruction(&call);
4883
4884 if (constval)
4885 return true;
4886
4887 IRBuilder<> Builder2(&call);
4888 getReverseBuilder(Builder2);
4889
4890 auto trace = call.getArgOperand(call.arg_size() - 1);
4891 auto address = call.getArgOperand(0);
4892
4893 auto dtrace = lookup(gutils->getNewFromOriginal(trace), Builder2);
4894 auto daddress = lookup(gutils->getNewFromOriginal(address), Builder2);
4895
4896 Value *dchoice;
4897 if (TR.query(&call)[{-1}].isPossiblePointer()) {
4898 dchoice = gutils->invertPointerM(&call, Builder2);
4899 } else {
4900 dchoice = diffe(&call, Builder2);
4901 }
4902
4903 if (call.hasMetadata("enzyme_gradient_setter")) {
4904 auto gradient_setter = cast<Function>(
4905 cast<ValueAsMetadata>(
4906 call.getMetadata("enzyme_gradient_setter")->getOperand(0).get())
4907 ->getValue());
4908
4909 TraceUtils::InsertChoiceGradient(
4910 Builder2, gradient_setter->getFunctionType(), gradient_setter,
4911 daddress, dchoice, dtrace);
4912 }
4913
4914 return true;
4915 }
4916
4917 if (call.hasFnAttr("enzyme_insert_argument")) {
4918 IRBuilder<> Builder2(&call);
4919 getReverseBuilder(Builder2);
4920
4921 auto name = call.getArgOperand(0);
4922 auto arg = call.getArgOperand(1);
4923 auto trace = call.getArgOperand(2);
4924
4925 auto gradient_setter = cast<Function>(
4926 cast<ValueAsMetadata>(
4927 call.getMetadata("enzyme_gradient_setter")->getOperand(0).get())
4928 ->getValue());
4929
4930 auto dtrace = lookup(gutils->getNewFromOriginal(trace), Builder2);
4931 auto dname = lookup(gutils->getNewFromOriginal(name), Builder2);
4932 Value *darg;
4933
4934 if (TR.query(arg)[{-1}].isPossiblePointer()) {
4935 darg = gutils->invertPointerM(arg, Builder2);
4936 } else {
4937 darg = diffe(arg, Builder2);
4938 }
4939
4940 TraceUtils::InsertArgumentGradient(Builder2,
4941 gradient_setter->getFunctionType(),
4942 gradient_setter, dname, darg, dtrace);
4943 return true;
4944 }
4945
4946 return false;
4947}
createMPIForwardCall
static void createMPIForwardCall(CallInst &call, ArrayRef< unsigned > ShadowArgs, GradientUtils *gutils, IRBuilder<> &Builder2)
Forward mode of an MPI call whose derivative is the same call with the buffers ShadowArgs replaced by...
Definition CallDerivatives.cpp:55
classifyAtomicModifyOp
static llvm::AtomicRMWInst::BinOp classifyAtomicModifyOp(llvm::Function *op, unsigned &opArgNo, llvm::Type *&FT)
Definition CallDerivatives.cpp:2303
EnzymeShadowAllocRewrite
void(* EnzymeShadowAllocRewrite)(LLVMValueRef, void *, LLVMValueRef, uint64_t, LLVMValueRef, uint8_t)
Definition CallDerivatives.cpp:32
hasSingleGCPreserveEnd
static bool hasSingleGCPreserveEnd(CallInst *Begin, GradientUtils *gutils)
Whether Begin has exactly one llvm.julia.gc_preserve_end outside the abort-only blocks the reverse pa...
Definition CallDerivatives.cpp:39
DIFFE_TYPE
DIFFE_TYPE
Potential differentiable argument classifications.
Definition EnzymeCallMarkers.h:50
isSum
CallInst * isSum(llvm::Value *v)
Definition FunctionUtils.cpp:5020
isDeallocationFunction
static bool isDeallocationFunction(const llvm::StringRef name, const llvm::TargetLibraryInfo &TLI)
Return whether a given function is a known C/C++ memory deallocation function For updating below one ...
Definition LibraryFuncs.h:139
isAllocationFunction
static bool isAllocationFunction(const llvm::StringRef name, const llvm::TargetLibraryInfo &TLI)
Return whether a given function is a known C/C++ memory allocation function For updating below one sh...
Definition LibraryFuncs.h:46
zeroKnownAllocation
static void zeroKnownAllocation(llvm::IRBuilder<> &bb, llvm::Value *toZero, llvm::ArrayRef< llvm::Value * > argValues, const llvm::StringRef funcName, const llvm::TargetLibraryInfo &TLI, llvm::CallInst *orig)
Definition LibraryFuncs.h:225
cached
bool cached
Definition ActivityAnalysis.cpp:179
MPIInactiveCommAllocators
const llvm::StringMap< size_t > MPIInactiveCommAllocators
Definition ActivityAnalysis.cpp:160
str
static std::string str(AugmentedStruct c)
Definition EnzymeLogic.h:62
looseTypeAnalysis
llvm::cl::opt< bool > looseTypeAnalysis
shadowHandlers
StringMap< std::function< Value *(IRBuilder<> &, CallInst *, ArrayRef< Value * >, GradientUtils *)> > shadowHandlers
Definition GradientUtils.cpp:71
freeKnownAllocation
llvm::CallInst * freeKnownAllocation(llvm::IRBuilder<> &builder, llvm::Value *tofree, llvm::StringRef allocationfn, const llvm::DebugLoc &debuglocation, const llvm::TargetLibraryInfo &TLI, llvm::CallInst *orig, GradientUtils *gutils)
Perform the corresponding deallocation of tofree, given it was allocated by allocationfn.
Definition GradientUtils.cpp:10102
MD_ToCopy
SmallVector< unsigned int, 9 > MD_ToCopy
Definition GradientUtils.cpp:131
EnzymeFreeInternalAllocations
llvm::cl::opt< bool > EnzymeFreeInternalAllocations
isMemFreeLibMFunction
static bool isMemFreeLibMFunction(llvm::StringRef str, llvm::Intrinsic::ID *ID=nullptr)
Definition TypeAnalysis.h:61
CreateDealloc
CallInst * CreateDealloc(llvm::IRBuilder<> &Builder, llvm::Value *ToFree)
Definition Utils.cpp:871
getOrInsertDifferentialMPI_Wait
llvm::Function * getOrInsertDifferentialMPI_Wait(llvm::Module &M, ArrayRef< llvm::Type * > T, Type *reqType, StringRef caller)
Definition Utils.cpp:2494
getOrInsertDifferentialWaitallSave
llvm::Function * getOrInsertDifferentialWaitallSave(llvm::Module &M, ArrayRef< llvm::Type * > T, PointerType *reqType)
Definition Utils.cpp:2426
getOrInsertOpFloatSum
llvm::Value * getOrInsertOpFloatSum(llvm::Module &M, llvm::Function *templateFn, llvm::Type *OpPtr, llvm::Type *OpType, ConcreteType CT, llvm::Type *intType, IRBuilder<> &B2)
Definition Utils.cpp:2601
extractBLAS
llvm::Optional< BlasInfo > extractBLAS(llvm::StringRef in)
Definition Utils.cpp:3871
CreateAllocation
Value * CreateAllocation(IRBuilder<> &Builder, llvm::Type *T, Value *Count, const Twine &Name, CallInst **caller, Instruction **ZeroMem, bool isDefault)
Definition Utils.cpp:714
getOrInsertPerCallingConv
FunctionCallee getOrInsertPerCallingConv(Module &M, Function *templateFn, StringRef callee, FunctionType *FT)
Definition Utils.cpp:2294
CustomErrorHandler
LLVMValueRef(* CustomErrorHandler)(const char *, LLVMValueRef, ErrorType, const void *, LLVMValueRef, LLVMBuilderRef)
Definition Utils.cpp:62
EmitNoDerivativeError
llvm::Value * EmitNoDerivativeError(const std::string &message, llvm::Instruction &inst, GradientUtils *gutils, llvm::IRBuilder<> &Builder2, llvm::Value *condition)
Definition Utils.cpp:4593
getFast
llvm::FastMathFlags getFast()
Get LLVM fast math flags.
Definition Utils.cpp:4026
getUndefinedValueForType
llvm::Constant * getUndefinedValueForType(llvm::Module &M, llvm::Type *T, bool forceZero)
Definition Utils.cpp:3918
getOrInsertCheckedFree
Function * getOrInsertCheckedFree(Module &M, CallInst *call, Type *Ty, unsigned width)
Definition Utils.cpp:2304
canonicalizeMPIName
static llvm::StringRef canonicalizeMPIName(llvm::StringRef Name)
Canonicalize the name of an MPI routine across calling conventions: both the C convention ("MPI_Recv"...
Definition Utils.h:2780
getRenamedPerCallingConv
static std::string getRenamedPerCallingConv(llvm::StringRef caller, llvm::StringRef callee)
Definition Utils.h:2812
startsWith
static bool startsWith(llvm::StringRef string, llvm::StringRef prefix)
Definition Utils.h:847
getUnqual
static llvm::PointerType * getUnqual(llvm::Type *T)
Definition Utils.h:1329
getMPIMemberPtr
static llvm::Value * getMPIMemberPtr(llvm::IRBuilder<> &B, llvm::Value *V, llvm::Type *T)
Definition Utils.h:1398
getInt8PtrTy
static llvm::PointerType * getInt8PtrTy(llvm::LLVMContext &Context, unsigned AddressSpace=0)
Definition Utils.h:1324
getIntrinsicDeclaration
static llvm::Function * getIntrinsicDeclaration(llvm::Module *M, llvm::Intrinsic::ID id, llvm::ArrayRef< llvm::Type * > Tys={})
Definition Utils.h:2541
getBaseObject
static llvm::Value * getBaseObject(llvm::Value *V, bool offsetAllowed=true)
Definition Utils.h:1711
isFortranMPICall
static bool isFortranMPICall(llvm::StringRef Name)
True if Name uses a Fortran ABI mangling of an MPI routine (trailing underscore(s),...
Definition Utils.h:2798
hasMetadata
static llvm::MDNode * hasMetadata(const llvm::GlobalObject *O, llvm::StringRef kind)
Check if a global has metadata.
Definition Utils.h:483
getMPIHelper
static llvm::StructType * getMPIHelper(llvm::LLVMContext &Context)
Definition Utils.h:1381
EnzymeJuliaAddrLoad
llvm::cl::opt< bool > EnzymeJuliaAddrLoad
ErrorType::MixedActivityError
@ MixedActivityError
EmitWarningAlways
void EmitWarningAlways(llvm::StringRef RemarkName, const llvm::DiagnosticLocation &Loc, const llvm::Function &F, const Args &...args)
Emit a diagnostic which is always shown, unlike EmitWarning which requires the user to opt in with -R...
Definition Utils.h:289
CreateSelect
static llvm::Value * CreateSelect(llvm::IRBuilder<> &Builder2, llvm::Value *cmp, llvm::Value *tval, llvm::Value *fval, const llvm::Twine &Name="")
Definition Utils.h:2283
getFuncNameFromCall
static llvm::StringRef getFuncNameFromCall(const llvm::CallBase *op)
Definition Utils.h:1470
MixedActivityHint
static constexpr const char * MixedActivityHint
Remedy appended to mismatched-activity warnings, which are emitted when a value Enzyme proved inactiv...
Definition Utils.h:315
ValueType
ValueType
Classification of value as an original program variable, a derivative variable, neither,...
Definition Utils.h:543
AdjointGenerator::lookup
llvm::Value * lookup(llvm::Value *val, llvm::IRBuilder<> &Builder)
Definition AdjointGenerator.h:2445
AdjointGenerator::DifferentiableMemCopyFloats
void DifferentiableMemCopyFloats(llvm::CallInst &call, llvm::Value *origArg, llvm::Value *dsto, llvm::Value *srco, llvm::Value *len_arg, llvm::IRBuilder<> &Builder2, llvm::ArrayRef< llvm::OperandBundleDef > ReverseDefs)
Definition AdjointGenerator.h:5070
AdjointGenerator::handleKnownCallDerivatives
bool handleKnownCallDerivatives(llvm::CallInst &call, llvm::Function *called, llvm::StringRef funcName, bool subsequent_calls_may_write, const std::vector< bool > &overwritten_args, llvm::CallInst *const newCall)
Definition CallDerivatives.cpp:2356
AdjointGenerator::setDiffe
void setDiffe(llvm::Value *val, llvm::Value *dif, llvm::IRBuilder<> &Builder)
Definition AdjointGenerator.h:2401
AdjointGenerator::MPI_TYPE_SIZE
llvm::Value * MPI_TYPE_SIZE(llvm::Value *DT, llvm::IRBuilder<> &B, llvm::Type *intType, llvm::Function *caller)
Definition AdjointGenerator.h:230
AdjointGenerator::applyChainRule
llvm::Value * applyChainRule(llvm::Type *diffType, llvm::IRBuilder<> &Builder, Func rule, Args... args)
Unwraps a vector derivative from its internal representation and applies a function f to each element...
Definition AdjointGenerator.h:2410
AdjointGenerator::shadowStoreSchedule
void shadowStoreSchedule(llvm::Instruction &I, llvm::Value *orig_ptr, bool &forwardsShadow, bool &backwardsShadow)
Which pass(es) emit the shadow counterpart of the store-like instruction I into the duplicated memory...
Definition AdjointGenerator.h:140
AdjointGenerator::resolveShadowPlaceholder
void resolveShadowPlaceholder(llvm::Instruction &I, llvm::Value *shadow=nullptr)
Resolve the shadow placeholder of the active instruction I in a pass which does not otherwise compute...
Definition AdjointGenerator.h:168
AdjointGenerator::eraseIfUnused
void eraseIfUnused(llvm::Instruction &I, bool erase=true, bool check=true)
Definition AdjointGenerator.h:114
AdjointGenerator::getForwardBuilder
void getForwardBuilder(llvm::IRBuilder<> &Builder2)
Definition AdjointGenerator.h:2392
AdjointGenerator::getReverseBuilder
void getReverseBuilder(llvm::IRBuilder<> &Builder2, bool original=true)
Definition AdjointGenerator.h:2388
AdjointGenerator::MPI_COMM_SIZE
llvm::Value * MPI_COMM_SIZE(llvm::Value *comm, llvm::IRBuilder<> &B, llvm::Type *rankTy, llvm::Function *caller)
Definition AdjointGenerator.h:370
AdjointGenerator::MPI_COMM_RANK
llvm::Value * MPI_COMM_RANK(llvm::Value *comm, llvm::IRBuilder<> &B, llvm::Type *rankTy, llvm::Function *caller)
Definition AdjointGenerator.h:315
AdjointGenerator::handleAdjointForIntrinsic
bool handleAdjointForIntrinsic(llvm::Intrinsic::ID ID, llvm::Instruction &I, llvm::SmallVectorImpl< llvm::Value * > &orig_ops)
Definition AdjointGenerator.h:4105
AdjointGenerator::handleMPI
void handleMPI(llvm::CallInst &call, llvm::Function *called, llvm::StringRef funcName)
Definition CallDerivatives.cpp:95
AdjointGenerator::addAtomicRMWValueAdjoint
void addAtomicRMWValueAdjoint(llvm::Instruction &I, llvm::Value *ptr, llvm::Value *val, llvm::Type *loadTy, llvm::Type *FT, bool negate, llvm::Align align, llvm::AtomicOrdering order, llvm::SyncScope::ID ssid, bool isVolatile)
Reverse-mode adjoint of the value operand val of an atomic modification new = old + val (negate == fa...
Definition AdjointGenerator.h:200
CacheUtility::newFunc
llvm::Function *const newFunc
The function whose instructions we are caching.
Definition CacheUtility.h:147
CacheUtility::TLI
llvm::TargetLibraryInfo & TLI
Various analysis results of newFunc.
Definition CacheUtility.h:150
CacheUtility::inversionAllocs
llvm::BasicBlock * inversionAllocs
Definition CacheUtility.h:161
ConcreteType
Concrete SubType of a given value.
Definition ConcreteType.h:41
GradientUtils::unnecessaryIntermediates
llvm::SmallPtrSet< llvm::Instruction *, 4 > unnecessaryIntermediates
Definition GradientUtils.h:186
GradientUtils::invertedPointers
llvm::ValueMap< const llvm::Value *, InvertedPointerVH > invertedPointers
Definition GradientUtils.h:143
GradientUtils::getNewFromOriginal
llvm::DebugLoc getNewFromOriginal(const llvm::DebugLoc L) const
GradientUtils::extractMeta
static llvm::Value * extractMeta(llvm::IRBuilder<> &Builder, llvm::Value *Agg, unsigned off, const llvm::Twine &name="")
Helper routine to extract a nested element from a struct/array. This is.
GradientUtils::TapesToPreventRecomputation
llvm::SmallPtrSet< llvm::Instruction *, 4 > TapesToPreventRecomputation
A set of tape extractions to enforce a cache of rather than attempting to recompute.
Definition GradientUtils.h:179
GradientUtils::legalRecompute
bool legalRecompute(const llvm::Value *val, const llvm::ValueToValueMapTy &available, llvm::IRBuilder<> *BuilderM, bool reverse=false, bool legalRecomputeCache=true) const
Definition GradientUtils.cpp:3973
GradientUtils::knownRecomputeHeuristic
std::map< const llvm::Value *, bool > knownRecomputeHeuristic
Definition GradientUtils.h:357
GradientUtils::addReverseBlock
llvm::BasicBlock * addReverseBlock(llvm::BasicBlock *currentBlock, llvm::Twine const &name, bool forkCache=true, bool push=true)
Definition GradientUtils.cpp:9797
GradientUtils::allocationsToBeRematerialized
llvm::SmallPtrSet< llvm::Value *, 4 > allocationsToBeRematerialized
Definition GradientUtils.h:310
GradientUtils::rematerializableAllocations
llvm::ValueMap< llvm::Value *, Rematerializer > rematerializableAllocations
Definition GradientUtils.h:309
GradientUtils::populateSeenFromKnownRecompute
std::map< UsageKey, bool > populateSeenFromKnownRecompute() const
Definition GradientUtils.h:341
GradientUtils::getWidth
unsigned getWidth()
Definition GradientUtils.h:411
GradientUtils::reverseBlocks
std::map< llvm::BasicBlock *, llvm::SmallVector< llvm::BasicBlock *, 4 > > reverseBlocks
Map of primal block to corresponding block(s) in reverse.
Definition GradientUtils.h:173
GradientUtils::oldFunc
llvm::Function * oldFunc
Definition GradientUtils.h:142
GradientUtils::getInvertedBundles
llvm::SmallVector< llvm::OperandBundleDef, 2 > getInvertedBundles(llvm::CallInst *orig, llvm::ArrayRef< ValueType > types, llvm::IRBuilder<> &Builder2, bool lookup, const llvm::ValueToValueMapTy &available=llvm::ValueToValueMapTy())
Definition GradientUtils.cpp:307
GradientUtils::replaceAWithB
void replaceAWithB(llvm::Value *A, llvm::Value *B, bool storeInCache=false) override
Replace this instruction both in LLVM modules and any local data-structures.
Definition GradientUtils.cpp:9842
GradientUtils::lookupM
llvm::Value * lookupM(llvm::Value *val, llvm::IRBuilder<> &BuilderM, const llvm::ValueToValueMapTy &incoming_availalble=llvm::ValueToValueMapTy(), bool tryLegalRecomputeCheck=true, llvm::BasicBlock *scope=nullptr) override
High-level utility to get the value an instruction at a new location specified by BuilderM.
Definition GradientUtils.cpp:6835
GradientUtils::invertPointerM
llvm::Value * invertPointerM(llvm::Value *val, llvm::IRBuilder<> &BuilderM)
GradientUtils::getReturnDiffeType
DIFFE_TYPE getReturnDiffeType(llvm::Value *orig, bool *primalReturnUsedP, bool *shadowReturnUsedP, DerivativeMode cmode) const
Definition GradientUtils.cpp:4539
GradientUtils::rematerializedPrimalOrShadowAllocations
llvm::SmallVector< llvm::Instruction *, 1 > rematerializedPrimalOrShadowAllocations
Definition GradientUtils.h:504
GradientUtils::reverseBlockToPrimal
std::map< llvm::BasicBlock *, llvm::BasicBlock * > reverseBlockToPrimal
Map of block in reverse to corresponding primal block.
Definition GradientUtils.h:175
GradientUtils::isConstantInstruction
bool isConstantInstruction(const llvm::Instruction *inst) const
Definition GradientUtils.cpp:9083
GradientUtils::getShadowType
static llvm::Type * getShadowType(llvm::Type *ty, unsigned width)
GradientUtils::cacheForReverse
llvm::Value * cacheForReverse(llvm::IRBuilder<> &BuilderQ, llvm::Value *malloc, int idx, bool replace=true)
Definition GradientUtils.cpp:2699
GradientUtils::backwardsOnlyShadows
llvm::ValueMap< llvm::Value *, ShadowRematerializer > backwardsOnlyShadows
Only loaded from and stored to (not captured), mapped to the stores (and memset).
Definition GradientUtils.h:315
GradientUtils::isConstantValue
bool isConstantValue(llvm::Value *val) const
Definition GradientUtils.cpp:9033
GradientUtils::notForAnalysis
llvm::SmallPtrSet< llvm::BasicBlock *, 4 > notForAnalysis
Definition GradientUtils.h:152
GradientUtils::erase
void erase(llvm::Instruction *I) override
Erase this instruction both from LLVM modules and any local data-structures.
Definition GradientUtils.cpp:9870
TraceUtils::InsertChoiceGradient
static llvm::CallInst * InsertChoiceGradient(llvm::IRBuilder<> &Builder, llvm::FunctionType *interface_type, llvm::Value *interface_function, llvm::Value *address, llvm::Value *choice, llvm::Value *trace)
Definition TraceUtils.cpp:312
TraceUtils::InsertArgumentGradient
static llvm::CallInst * InsertArgumentGradient(llvm::IRBuilder<> &Builder, llvm::FunctionType *interface_type, llvm::Value *interface_function, llvm::Value *name, llvm::Value *argument, llvm::Value *trace)
Definition TraceUtils.cpp:328
TypeResults::firstPointer
ConcreteType firstPointer(size_t num, llvm::Value *val, llvm::Instruction *I, GradientUtils *gutils=nullptr, llvm::IRBuilder<> *BuilderIfShouldErr=nullptr, bool pointerIntSame=false) const
Returns whether in the first num bytes there is pointer, int, float, or none If pointerIntSame is set...
Definition TypeAnalysis.cpp:6742
TypeTree
Class representing the underlying types of values as sequences of offsets to a ConcreteType.
Definition TypeTree.h:72
DifferentialUseAnalysis::is_value_needed_in_reverse
bool is_value_needed_in_reverse(const GradientUtils *gutils, const llvm::Value *inst, DerivativeMode mode, std::map< UsageKey, bool > &seen, const llvm::SmallPtrSetImpl< llvm::BasicBlock * > &oldUnreachable)
Definition DifferentialUseAnalysis.h:92
Generated on Mon Oct 12 2026 02:52:55 for Enzyme by doxygen 1.12.0