Enzyme main
Loading...
Searching...
No Matches
AdjointGenerator.h
Go to the documentation of this file.
1//===- AdjointGenerator.h - Implementation of Adjoint's of instructions --===//
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 an instruction visitor AdjointGenerator that generates
22// the corresponding augmented forward pass code, and adjoints for all
23// LLVM instructions.
24//
25//===----------------------------------------------------------------------===//
26
27#ifndef ENZYME_ADJOINT_GENERATOR_H
28#define ENZYME_ADJOINT_GENERATOR_H
29
30#include "llvm/ADT/ArrayRef.h"
31#include "llvm/ADT/STLExtras.h"
32#include "llvm/ADT/SmallSet.h"
33#include "llvm/ADT/SmallVector.h"
34#include "llvm/Analysis/ValueTracking.h"
35#include "llvm/IR/Constants.h"
36#include "llvm/IR/DerivedTypes.h"
37#include "llvm/IR/IntrinsicsX86.h"
38#include "llvm/IR/Value.h"
39#include "llvm/Transforms/Utils/BasicBlockUtils.h"
40#include "llvm/Transforms/Utils/Cloning.h"
41
42#include "DiffeGradientUtils.h"
44#include "EnzymeLogic.h"
45#include "FunctionUtils.h"
46#include "GradientUtils.h"
47#include "LibraryFuncs.h"
48#include "TraceUtils.h"
49#include "TypeAnalysis/TBAA.h"
50
51#define DEBUG_TYPE "enzyme"
52
53// Helper instruction visitor that generates adjoints
54class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
55private:
56 // Type of code being generated (forward, reverse, or both)
57 const DerivativeMode Mode;
58
59 GradientUtils *const gutils;
60 llvm::ArrayRef<DIFFE_TYPE> constant_args;
61 DIFFE_TYPE retType;
62 TypeResults &TR = gutils->TR;
63 std::function<unsigned(llvm::Instruction *, CacheType, llvm::IRBuilder<> &)>
64 getIndex;
65 const std::map<llvm::CallInst *, std::pair<bool, const std::vector<bool>>>
66 overwritten_args_map;
67 const AugmentedReturn *augmentedReturn;
68 const std::map<llvm::ReturnInst *, llvm::StoreInst *> *replacedReturns;
69
70 const llvm::SmallPtrSetImpl<const llvm::Value *> &unnecessaryValues;
71 const llvm::SmallPtrSetImpl<const llvm::Instruction *>
72 &unnecessaryInstructions;
73 const llvm::SmallPtrSetImpl<const llvm::Instruction *> &unnecessaryStores;
74 const llvm::SmallPtrSetImpl<llvm::BasicBlock *> &oldUnreachable;
75
76public:
78 DerivativeMode Mode, GradientUtils *gutils,
79 llvm::ArrayRef<DIFFE_TYPE> constant_args, DIFFE_TYPE retType,
80 std::function<unsigned(llvm::Instruction *, CacheType,
81 llvm::IRBuilder<> &)>
82 getIndex,
83 const std::map<llvm::CallInst *, std::pair<bool, const std::vector<bool>>>
84 overwritten_args_map,
85 const AugmentedReturn *augmentedReturn,
86 const std::map<llvm::ReturnInst *, llvm::StoreInst *> *replacedReturns,
87 const llvm::SmallPtrSetImpl<const llvm::Value *> &unnecessaryValues,
88 const llvm::SmallPtrSetImpl<const llvm::Instruction *>
89 &unnecessaryInstructions,
90 const llvm::SmallPtrSetImpl<const llvm::Instruction *> &unnecessaryStores,
91 const llvm::SmallPtrSetImpl<llvm::BasicBlock *> &oldUnreachable)
92 : Mode(Mode), gutils(gutils), constant_args(constant_args),
93 retType(retType), getIndex(getIndex),
94 overwritten_args_map(overwritten_args_map),
95 augmentedReturn(augmentedReturn), replacedReturns(replacedReturns),
96 unnecessaryValues(unnecessaryValues),
97 unnecessaryInstructions(unnecessaryInstructions),
98 unnecessaryStores(unnecessaryStores), oldUnreachable(oldUnreachable) {
99 using namespace llvm;
100
101 assert(TR.getFunction() == gutils->oldFunc);
102 for (auto &pair : TR.analyzer->analysis) {
103 if (auto in = dyn_cast<Instruction>(pair.first)) {
104 if (in->getParent()->getParent() != gutils->oldFunc) {
105 llvm::errs() << "inf: " << *in->getParent()->getParent() << "\n";
106 llvm::errs() << "gutils->oldFunc: " << *gutils->oldFunc << "\n";
107 llvm::errs() << "in: " << *in << "\n";
108 }
109 assert(in->getParent()->getParent() == gutils->oldFunc);
110 }
111 }
112 }
113
114 void eraseIfUnused(llvm::Instruction &I, bool erase = true,
115 bool check = true) {
116 using namespace llvm;
117
118 bool used =
119 unnecessaryInstructions.find(&I) == unnecessaryInstructions.end();
120 if (!used) {
121 // if decided to cache a value, preserve it here for later
122 // replacement in EnzymeLogic
123 auto found = gutils->knownRecomputeHeuristic.find(&I);
124 if (found != gutils->knownRecomputeHeuristic.end() && !found->second)
125 used = true;
126 }
127 auto iload = gutils->getNewFromOriginal((llvm::Value *)&I);
128 if (used && check)
129 return;
130
131 if (auto newi = dyn_cast<Instruction>(iload))
132 gutils->eraseWithPlaceholder(newi, &I, "_replacementA", erase);
133 }
134
135 /// Which pass(es) emit the shadow counterpart of the store-like
136 /// instruction I into the duplicated memory at orig_ptr: the augmented
137 /// primal (forwardsShadow) and/or the reverse pass (backwardsShadow, for
138 /// shadow allocations only materialized there), and neither of them for
139 /// shadow arguments that must not be written.
140 void shadowStoreSchedule(llvm::Instruction &I, llvm::Value *orig_ptr,
141 bool &forwardsShadow, bool &backwardsShadow) {
142 backwardsShadow = false;
143 forwardsShadow = true;
144 for (auto pair : gutils->backwardsOnlyShadows) {
145 if (pair.second.stores.count(&I)) {
146 backwardsShadow = true;
147 forwardsShadow = pair.second.primalInitialize;
148 if (auto inst = llvm::dyn_cast<llvm::Instruction>(pair.first))
149 if (!forwardsShadow && pair.second.LI &&
150 pair.second.LI->contains(inst->getParent()))
151 backwardsShadow = false;
152 }
153 }
154
155 if (auto arg = llvm::dyn_cast<llvm::Argument>(getBaseObject(orig_ptr))) {
156 unsigned argNo = arg->getArgNo();
157 if (argNo < gutils->nowrite_shadows.size() &&
158 gutils->nowrite_shadows[argNo]) {
159 forwardsShadow = false;
160 backwardsShadow = false;
161 }
162 }
163 }
164
165 /// Resolve the shadow placeholder of the active instruction I in a pass
166 /// which does not otherwise compute its shadow, replacing it by the given
167 /// value (a zero shadow if null).
168 void resolveShadowPlaceholder(llvm::Instruction &I,
169 llvm::Value *shadow = nullptr) {
170 using namespace llvm;
171 auto ifound = gutils->invertedPointers.find(&I);
172 if (ifound == gutils->invertedPointers.end())
173 return;
174 if (!shadow)
175 shadow = Constant::getNullValue(gutils->getShadowType(I.getType()));
176 auto placeholder = cast<PHINode>(&*ifound->second);
177 gutils->invertedPointers.erase(ifound);
178 gutils->replaceAWithB(placeholder, shadow);
179 gutils->erase(placeholder);
180 gutils->invertedPointers.insert(
181 std::make_pair((const Value *)&I, InvertedPointerVH(gutils, shadow)));
182 }
183
184 /// The ordering with which the reverse pass may read a location that the
185 /// primal modified atomically with the given ordering.
186 static llvm::AtomicOrdering reverseLoadOrdering(llvm::AtomicOrdering order) {
187 if (order == llvm::AtomicOrdering::Release)
188 return llvm::AtomicOrdering::Monotonic;
189 if (order == llvm::AtomicOrdering::AcquireRelease)
190 return llvm::AtomicOrdering::Acquire;
191 return order;
192 }
193
194 /// Reverse-mode adjoint of the value operand val of an atomic modification
195 /// new = old + val (negate == false) or new = old - val (negate == true) of
196 /// the active location ptr by the instruction I, whose result is inactive:
197 /// the adjoint of val is (the negation of) the contents of the shadow
198 /// location, read as loadTy (holding the bits of the floating point type
199 /// FT) with the ordering downgraded for the reverse pass.
200 void addAtomicRMWValueAdjoint(llvm::Instruction &I, llvm::Value *ptr,
201 llvm::Value *val, llvm::Type *loadTy,
202 llvm::Type *FT, bool negate, llvm::Align align,
203 llvm::AtomicOrdering order,
204 llvm::SyncScope::ID ssid, bool isVolatile) {
205 using namespace llvm;
206 assert(!gutils->isConstantValue(ptr));
207 IRBuilder<> Builder2(&I);
208 getReverseBuilder(Builder2);
209 Value *ip = gutils->invertPointerM(ptr, Builder2);
210 ip = lookup(ip, Builder2);
211 order = reverseLoadOrdering(order);
212
213 auto rule = [&](Value *ip) -> Value * {
214 LoadInst *dif1 = Builder2.CreateLoad(loadTy, ip, isVolatile);
215 dif1->setAlignment(align);
216 dif1->setOrdering(order);
217 dif1->setSyncScopeID(ssid);
218 Value *res = dif1;
219 if (negate) {
220 res = Builder2.CreateBitCast(res, FT);
221 res = Builder2.CreateFNeg(res);
222 }
223 return Builder2.CreateBitCast(res, val->getType());
224 };
225 Value *diff = applyChainRule(val->getType(), Builder2, rule, ip);
226
227 addToDiffe(val, diff, Builder2, FT->getScalarType());
228 }
229
230 llvm::Value *MPI_TYPE_SIZE(llvm::Value *DT, llvm::IRBuilder<> &B,
231 llvm::Type *intType, llvm::Function *caller) {
232 using namespace llvm;
233
234 // Fortran ABI ("mpi_type_size_"): all arguments are passed by reference
235 // and the call takes an extra trailing `ierr` argument.
236 if (isFortranMPICall(caller->getName())) {
237 Type *i32 = Type::getInt32Ty(DT->getContext());
238 Type *pargs[] = {getInt8PtrTy(DT->getContext()), getUnqual(i32),
239 getUnqual(i32)};
240 auto FT =
241 FunctionType::get(Type::getVoidTy(DT->getContext()), pargs, false);
242 IRBuilder<> AllocaBuilder(gutils->inversionAllocs);
243 auto alloc = AllocaBuilder.CreateAlloca(i32);
244 auto ierr = AllocaBuilder.CreateAlloca(i32);
245 llvm::Value *args[] = {DT, alloc, ierr};
246 if (DT->getType() != pargs[0])
247 args[0] = B.CreateBitCast(args[0], pargs[0]);
248 B.CreateCall(
249 B.GetInsertBlock()->getParent()->getParent()->getOrInsertFunction(
250 getRenamedPerCallingConv(caller->getName(), "MPI_Type_size"), FT),
251 args);
252 return B.CreateLoad(i32, alloc);
253 }
254
255 if (DT->getType()->isIntegerTy())
256 DT = B.CreateIntToPtr(DT, getInt8PtrTy(DT->getContext()));
257
258 if (Constant *C = dyn_cast<Constant>(DT)) {
259 while (ConstantExpr *CE = dyn_cast<ConstantExpr>(C)) {
260 C = CE->getOperand(0);
261 }
262 if (auto GV = dyn_cast<GlobalVariable>(C)) {
263 if (GV->getName() == "ompi_mpi_double") {
264 return ConstantInt::get(intType, 8, false);
265 } else if (GV->getName() == "ompi_mpi_float") {
266 return ConstantInt::get(intType, 4, false);
267 }
268 }
269 }
270 Type *pargs[] = {getInt8PtrTy(DT->getContext()), getUnqual(intType)};
271 auto FT = FunctionType::get(intType, pargs, false);
272 auto alloc = IRBuilder<>(gutils->inversionAllocs).CreateAlloca(intType);
273 llvm::Value *args[] = {DT, alloc};
274 if (DT->getType() != pargs[0])
275 args[0] = B.CreateBitCast(args[0], pargs[0]);
276 AttributeList AL;
277 AL = AL.addParamAttribute(DT->getContext(), 0,
278 Attribute::AttrKind::ReadOnly);
279 AL = addFunctionNoCapture(DT->getContext(), AL, 0);
280 AL =
281 AL.addParamAttribute(DT->getContext(), 0, Attribute::AttrKind::NoAlias);
282 AL =
283 AL.addParamAttribute(DT->getContext(), 0, Attribute::AttrKind::NonNull);
284 AL = AL.addParamAttribute(DT->getContext(), 1,
285 Attribute::AttrKind::WriteOnly);
286 AL = addFunctionNoCapture(DT->getContext(), AL, 1);
287 AL =
288 AL.addParamAttribute(DT->getContext(), 1, Attribute::AttrKind::NoAlias);
289 AL =
290 AL.addParamAttribute(DT->getContext(), 1, Attribute::AttrKind::NonNull);
291 AL = AL.addAttributeAtIndex(DT->getContext(), AttributeList::FunctionIndex,
292 Attribute::AttrKind::NoUnwind);
293 AL = AL.addAttributeAtIndex(DT->getContext(), AttributeList::FunctionIndex,
294 Attribute::AttrKind::NoFree);
295 AL = AL.addAttributeAtIndex(DT->getContext(), AttributeList::FunctionIndex,
296 Attribute::AttrKind::NoSync);
297 AL = AL.addAttributeAtIndex(DT->getContext(), AttributeList::FunctionIndex,
298 Attribute::AttrKind::WillReturn);
299 auto CI = B.CreateCall(
300 B.GetInsertBlock()->getParent()->getParent()->getOrInsertFunction(
301 getRenamedPerCallingConv(caller->getName(), "MPI_Type_size"), FT,
302 AL),
303 args);
304#if LLVM_VERSION_MAJOR >= 16
305 CI->setOnlyAccessesArgMemory();
306#else
307 CI->addAttributeAtIndex(AttributeList::FunctionIndex,
308 Attribute::ArgMemOnly);
309#endif
310 return B.CreateLoad(intType, alloc);
311 }
312
313 // To be double-checked against the functionality needed and the respective
314 // implementation in Adjoint-MPI
315 llvm::Value *MPI_COMM_RANK(llvm::Value *comm, llvm::IRBuilder<> &B,
316 llvm::Type *rankTy, llvm::Function *caller) {
317 using namespace llvm;
318
319 // Fortran ABI ("mpi_comm_rank_"): all arguments are passed by reference
320 // and the call takes an extra trailing `ierr` argument.
321 if (isFortranMPICall(caller->getName())) {
322 Type *i32 = Type::getInt32Ty(comm->getContext());
323 Type *pargs[] = {getInt8PtrTy(comm->getContext()), getUnqual(i32),
324 getUnqual(i32)};
325 auto FT =
326 FunctionType::get(Type::getVoidTy(comm->getContext()), pargs, false);
327 IRBuilder<> AllocaBuilder(gutils->inversionAllocs);
328 auto alloc = AllocaBuilder.CreateAlloca(i32);
329 auto ierr = AllocaBuilder.CreateAlloca(i32);
330 llvm::Value *args[] = {comm, alloc, ierr};
331 if (comm->getType() != pargs[0])
332 args[0] = B.CreateBitCast(args[0], pargs[0]);
333 B.CreateCall(
334 B.GetInsertBlock()->getParent()->getParent()->getOrInsertFunction(
335 getRenamedPerCallingConv(caller->getName(), "MPI_Comm_rank"), FT),
336 args);
337 return B.CreateLoad(i32, alloc);
338 }
339
340 Type *pargs[] = {comm->getType(), getUnqual(rankTy)};
341 auto FT = FunctionType::get(rankTy, pargs, false);
342 auto &context = comm->getContext();
343 auto alloc = IRBuilder<>(gutils->inversionAllocs).CreateAlloca(rankTy);
344 AttributeList AL;
345 AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::ReadOnly);
346 AL = addFunctionNoCapture(context, AL, 0);
347 AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NoAlias);
348 AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NonNull);
349 AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::WriteOnly);
350 AL = addFunctionNoCapture(context, AL, 1);
351 AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NoAlias);
352 AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NonNull);
353 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
354 Attribute::AttrKind::NoUnwind);
355 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
356 Attribute::AttrKind::NoFree);
357 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
358 Attribute::AttrKind::NoSync);
359 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
360 Attribute::AttrKind::WillReturn);
361 llvm::Value *args[] = {comm, alloc};
362 B.CreateCall(
363 B.GetInsertBlock()->getParent()->getParent()->getOrInsertFunction(
364 getRenamedPerCallingConv(caller->getName(), "MPI_Comm_rank"), FT,
365 AL),
366 args);
367 return B.CreateLoad(rankTy, alloc);
368 }
369
370 llvm::Value *MPI_COMM_SIZE(llvm::Value *comm, llvm::IRBuilder<> &B,
371 llvm::Type *rankTy, llvm::Function *caller) {
372 using namespace llvm;
373
374 // Fortran ABI ("mpi_comm_size_"): all arguments are passed by reference
375 // and the call takes an extra trailing `ierr` argument.
376 if (isFortranMPICall(caller->getName())) {
377 Type *i32 = Type::getInt32Ty(comm->getContext());
378 Type *pargs[] = {getInt8PtrTy(comm->getContext()), getUnqual(i32),
379 getUnqual(i32)};
380 auto FT =
381 FunctionType::get(Type::getVoidTy(comm->getContext()), pargs, false);
382 IRBuilder<> AllocaBuilder(gutils->inversionAllocs);
383 auto alloc = AllocaBuilder.CreateAlloca(i32);
384 auto ierr = AllocaBuilder.CreateAlloca(i32);
385 llvm::Value *args[] = {comm, alloc, ierr};
386 if (comm->getType() != pargs[0])
387 args[0] = B.CreateBitCast(args[0], pargs[0]);
388 B.CreateCall(
389 B.GetInsertBlock()->getParent()->getParent()->getOrInsertFunction(
390 getRenamedPerCallingConv(caller->getName(), "MPI_Comm_size"), FT),
391 args);
392 return B.CreateLoad(i32, alloc);
393 }
394
395 Type *pargs[] = {comm->getType(), getUnqual(rankTy)};
396 auto FT = FunctionType::get(rankTy, pargs, false);
397 auto &context = comm->getContext();
398 auto alloc = IRBuilder<>(gutils->inversionAllocs).CreateAlloca(rankTy);
399 AttributeList AL;
400 AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::ReadOnly);
401 AL = addFunctionNoCapture(context, AL, 0);
402 AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NoAlias);
403 AL = AL.addParamAttribute(context, 0, Attribute::AttrKind::NonNull);
404 AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::WriteOnly);
405 AL = addFunctionNoCapture(context, AL, 1);
406 AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NoAlias);
407 AL = AL.addParamAttribute(context, 1, Attribute::AttrKind::NonNull);
408 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
409 Attribute::AttrKind::NoUnwind);
410 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
411 Attribute::AttrKind::NoFree);
412 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
413 Attribute::AttrKind::NoSync);
414 AL = AL.addAttributeAtIndex(context, AttributeList::FunctionIndex,
415 Attribute::AttrKind::WillReturn);
416 llvm::Value *args[] = {comm, alloc};
417 B.CreateCall(
418 B.GetInsertBlock()->getParent()->getParent()->getOrInsertFunction(
419 getRenamedPerCallingConv(caller->getName(), "MPI_Comm_size"), FT,
420 AL),
421 args);
422 return B.CreateLoad(rankTy, alloc);
423 }
424
425 void visitInstruction(llvm::Instruction &inst) {
426 using namespace llvm;
427
428 // TODO explicitly handle all instructions rather than using the catch all
429 // below
430
431 switch (inst.getOpcode()) {
432#include "InstructionDerivatives.inc"
433 default:
434 break;
435 }
436
437 if (gutils->isConstantInstruction(&inst)) {
439 return;
440 eraseIfUnused(inst);
441 return;
442 }
443
444 std::string s;
445 llvm::raw_string_ostream ss(s);
446 ss << "in Mode: " << to_string(Mode) << "\n";
447 ss << "cannot handle unknown instruction\n" << inst;
448 IRBuilder<> Builder2(&inst);
449 getForwardBuilder(Builder2);
450 EmitNoDerivativeError(ss.str(), inst, gutils, Builder2);
451 if (!gutils->isConstantValue(&inst)) {
452 if (Mode == DerivativeMode::ForwardMode ||
455 setDiffe(&inst,
456 Constant::getNullValue(gutils->getShadowType(inst.getType())),
457 Builder2);
458 }
459 if (!inst.getType()->isVoidTy()) {
460 for (auto &U :
461 make_early_inc_range(gutils->getNewFromOriginal(&inst)->uses())) {
462 U.set(UndefValue::get(inst.getType()));
463 }
464 }
465 eraseIfUnused(inst, /*erase*/ true, /*check*/ false);
466 return;
467 }
468
469 // Common function for falling back to the implementation
470 // of dual propagation, as available in invertPointerM.
471 void forwardModeInvertedPointerFallback(llvm::Instruction &I) {
472 using namespace llvm;
473
474 if (gutils->isConstantValue(&I)) {
475 // A constant value may have a shadow cached by invertPointerM. It is a
476 // real shadow, not a placeholder, so leave it for reuse.
477 return;
478 }
479
480 auto found = gutils->invertedPointers.find(&I);
481 assert(found != gutils->invertedPointers.end());
482 auto placeholder = cast<PHINode>(&*found->second);
483 gutils->invertedPointers.erase(found);
484
486 gutils, &I, Mode, oldUnreachable)) {
487 gutils->erase(placeholder);
488 return;
489 }
490
491 IRBuilder<> Builder2(&I);
492 getForwardBuilder(Builder2);
493
494 auto toset = gutils->invertPointerM(&I, Builder2);
495
496 assert(toset != placeholder);
497
498 gutils->replaceAWithB(placeholder, toset);
499 placeholder->replaceAllUsesWith(toset);
500 gutils->erase(placeholder);
501 gutils->invertedPointers.insert(
502 std::make_pair((const Value *)&I, InvertedPointerVH(gutils, toset)));
503 return;
504 }
505
506 void visitAllocaInst(llvm::AllocaInst &I) {
507 eraseIfUnused(I);
508 switch (Mode) {
513 return;
514 }
515 default:
516 return;
517 }
518 }
519
520 void visitICmpInst(llvm::ICmpInst &I) { eraseIfUnused(I); }
521
522 void visitFCmpInst(llvm::FCmpInst &I) { eraseIfUnused(I); }
523
524 void visitLoadLike(llvm::Instruction &I, llvm::MaybeAlign alignment,
525 bool constantval, llvm::Value *mask = nullptr,
526 llvm::Value *orig_maskInit = nullptr) {
527 using namespace llvm;
528
529 auto &DL = gutils->newFunc->getParent()->getDataLayout();
530 auto LoadSize = (DL.getTypeSizeInBits(I.getType()) + 1) / 8;
531
532 assert(Mode == DerivativeMode::ForwardMode ||
534 assert(Mode == DerivativeMode::ForwardMode ||
536 gutils->can_modref_map->find(&I) != gutils->can_modref_map->end());
537 bool can_modref = (Mode == DerivativeMode::ForwardMode ||
539 ? false
540 : gutils->can_modref_map->find(&I)->second;
541
542 constantval |= gutils->isConstantValue(&I);
543
544 Type *type = gutils->getShadowType(I.getType());
545 (void)type;
546
547 auto *newi = dyn_cast<Instruction>(gutils->getNewFromOriginal(&I));
548
549 SmallVector<Metadata *, 1> scopeMD = {
550 gutils->getDerivativeAliasScope(I.getOperand(0), -1)};
551 if (auto prev = I.getMetadata(LLVMContext::MD_alias_scope)) {
552 for (auto &M : cast<MDNode>(prev)->operands()) {
553 scopeMD.push_back(M);
554 }
555 }
556 auto scope = MDNode::get(I.getContext(), scopeMD);
557 newi->setMetadata(LLVMContext::MD_alias_scope, scope);
558
559 SmallVector<Metadata *, 1> MDs;
560 for (size_t j = 0; j < gutils->getWidth(); j++) {
561 MDs.push_back(gutils->getDerivativeAliasScope(I.getOperand(0), j));
562 }
563 if (auto prev = I.getMetadata(LLVMContext::MD_noalias)) {
564 for (auto &M : cast<MDNode>(prev)->operands()) {
565 MDs.push_back(M);
566 }
567 }
568 auto noscope = MDNode::get(I.getContext(), MDs);
569 newi->setMetadata(LLVMContext::MD_noalias, noscope);
570
571 auto vd = TR.query(&I);
572
573 IRBuilder<> BuilderZ(newi);
574 if (!vd.isKnown()) {
575 auto ET = I.getType();
576 if (looseTypeAnalysis || true) {
577 vd = defaultTypeTreeForLLVM(ET, &I);
578 if (EmitWarningEnabled(I.getContext())) {
579 std::string str;
580 raw_string_ostream ss(str);
581 ss << "Cannot deduce type of load " << I << ", assumed " << vd.str()
582 << "\n";
583 EmitWarning("CannotDeduceType", I, ss.str());
584 }
585 goto known;
586 }
587 {
588 std::string str;
589 raw_string_ostream ss(str);
590 ss << "Cannot deduce type of load " << I;
591 EmitNoTypeError(str, I, gutils, BuilderZ);
592 }
593 known:;
594 }
595
596 if (Mode == DerivativeMode::ForwardMode ||
599 if (!constantval) {
600 auto found = gutils->invertedPointers.find(&I);
601 assert(found != gutils->invertedPointers.end());
602 Instruction *placeholder = cast<Instruction>(&*found->second);
603 assert(placeholder->getType() == type);
604 gutils->invertedPointers.erase(found);
605
606 // only make shadow where caching needed
608 QueryType::Shadow>(gutils, &I, Mode, oldUnreachable)) {
609 gutils->erase(placeholder);
610 return;
611 }
612
613 if (can_modref) {
614 if (vd[{-1}].isPossiblePointer()) {
615 Value *newip = gutils->cacheForReverse(
616 BuilderZ, placeholder,
617 getIndex(&I, CacheType::Shadow, BuilderZ));
618 assert(newip->getType() == type);
619 gutils->invertedPointers.insert(std::make_pair(
620 (const Value *)&I, InvertedPointerVH(gutils, newip)));
621 } else {
622 gutils->erase(placeholder);
623 }
624 } else {
625 Value *newip = gutils->invertPointerM(&I, BuilderZ);
626 if (gutils->runtimeActivity && vd[{-1}].isFloat()) {
627 // TODO handle mask
628 assert(!mask);
629
630 auto rule = [&](Value *inop, Value *newip) -> Value * {
631 Value *shadow = BuilderZ.CreateICmpNE(
632 gutils->getNewFromOriginal(I.getOperand(0)), inop);
633 newip = CreateSelect(BuilderZ, shadow, newip,
634 Constant::getNullValue(newip->getType()));
635 return newip;
636 };
637 newip = applyChainRule(
638 I.getType(), BuilderZ, rule,
639 gutils->invertPointerM(I.getOperand(0), BuilderZ), newip);
640 }
641 assert(newip->getType() == type);
642 placeholder->replaceAllUsesWith(newip);
643 gutils->erase(placeholder);
644 gutils->invertedPointers.erase(&I);
645 gutils->invertedPointers.insert(std::make_pair(
646 (const Value *)&I, InvertedPointerVH(gutils, newip)));
647 }
648 }
649 return;
650 }
651
652 //! Store inverted pointer loads that need to be cached for use in reverse
653 //! pass
654 if (vd[{-1}].isPossiblePointer()) {
655 auto found = gutils->invertedPointers.find(&I);
656 if (found != gutils->invertedPointers.end()) {
657 Instruction *placeholder = cast<Instruction>(&*found->second);
658 assert(placeholder->getType() == type);
659 gutils->invertedPointers.erase(found);
660
661 if (!constantval) {
662 Value *newip = nullptr;
663
664 // TODO: In the case of fwd mode this should be true if the loaded
665 // value itself is used as a pointer.
667 QueryType::Shadow>(gutils, &I, Mode, oldUnreachable);
668
669 switch (Mode) {
670
673 if (!needShadow) {
674 gutils->erase(placeholder);
675 } else {
676 newip = gutils->invertPointerM(&I, BuilderZ);
677 assert(newip->getType() == type);
678 if (Mode == DerivativeMode::ReverseModePrimal && can_modref &&
680 QueryType::Shadow>(gutils, &I,
682 oldUnreachable)) {
683 gutils->cacheForReverse(
684 BuilderZ, newip, getIndex(&I, CacheType::Shadow, BuilderZ));
685 }
686 placeholder->replaceAllUsesWith(newip);
687 gutils->erase(placeholder);
688 gutils->invertedPointers.insert(std::make_pair(
689 (const Value *)&I, InvertedPointerVH(gutils, newip)));
690 }
691 break;
692 }
696 assert(0 && "impossible branch");
697 return;
698 }
700 if (!needShadow) {
701 gutils->erase(placeholder);
702 } else {
703 // only make shadow where caching needed
704 if (can_modref) {
705 newip = gutils->cacheForReverse(
706 BuilderZ, placeholder,
707 getIndex(&I, CacheType::Shadow, BuilderZ));
708 assert(newip->getType() == type);
709 gutils->invertedPointers.insert(std::make_pair(
710 (const Value *)&I, InvertedPointerVH(gutils, newip)));
711 } else {
712 newip = gutils->invertPointerM(&I, BuilderZ);
713 assert(newip->getType() == type);
714 placeholder->replaceAllUsesWith(newip);
715 gutils->erase(placeholder);
716 gutils->invertedPointers.insert(std::make_pair(
717 (const Value *)&I, InvertedPointerVH(gutils, newip)));
718 }
719 }
720 break;
721 }
722 }
723
724 } else {
725 gutils->erase(placeholder);
726 }
727 }
728 }
729
730 Value *inst = newi;
731
732 //! Store loads that need to be cached for use in reverse pass
733
734 // Only cache value here if caching decision isn't precomputed.
735 // Otherwise caching will be done inside EnzymeLogic.cpp at
736 // the end of the function jointly.
737 if (Mode != DerivativeMode::ForwardMode &&
739 !gutils->knownRecomputeHeuristic.count(&I) && can_modref &&
740 !gutils->unnecessaryIntermediates.count(&I)) {
741 // we can pre initialize all the knownRecomputeHeuristic values to false
742 // (not needing) as we may assume that minCutCache already preserves
743 // everything it requires.
744 std::map<UsageKey, bool> Seen = gutils->populateSeenFromKnownRecompute();
745 bool primalNeededInReverse = false;
746 {
747 auto found = gutils->knownRecomputeHeuristic.find(&I);
748 if (found != gutils->knownRecomputeHeuristic.end()) {
749 if (!found->second) {
750 primalNeededInReverse = true;
751 }
752 }
753 }
754 auto cacheMode = (Mode == DerivativeMode::ReverseModePrimal)
756 : Mode;
757 primalNeededInReverse |=
759 QueryType::Primal>(gutils, &I, cacheMode, Seen, oldUnreachable);
760 if (primalNeededInReverse) {
761 inst = gutils->cacheForReverse(BuilderZ, newi,
762 getIndex(&I, CacheType::Self, BuilderZ));
763 (void)inst;
764 assert(inst->getType() == I.getType());
765
769 assert(inst != newi);
770 } else {
771 assert(inst == newi);
772 }
773 }
774 }
775
777 return;
778
779 if (constantval)
780 return;
781
783 // Assume that non enzyme_shadow globals are inactive
784 // If we ever store to a global variable, we will error if it doesn't
785 // have a shadow This allows functions who only read global memory to
786 // have their derivative computed Note that this is too aggressive for
787 // general programs as if the global aliases with an argument something
788 // that is written to, then we will have a logical error
789 if (auto arg = dyn_cast<GlobalVariable>(I.getOperand(0))) {
790 if (!hasMetadata(arg, "enzyme_shadow")) {
791 return;
792 }
793 }
794 }
795
796 // Only propagate if instruction is active. The value can be active and not
797 // the instruction if the value is a potential pointer. This may not be
798 // caught by type analysis is the result does not have a known type.
799 if (!gutils->isConstantInstruction(&I)) {
800 switch (Mode) {
804 assert(0 && "impossible branch");
805 return;
806 }
809
810 IRBuilder<> Builder2(&I);
811 getReverseBuilder(Builder2);
812
813 Value *prediff = nullptr;
814
815 for (ssize_t i = -1; i < (ssize_t)LoadSize; ++i) {
816 if (vd[{(int)i}].isFloat()) {
817 prediff = diffe(&I, Builder2);
818 break;
819 }
820 }
821
822 Value *premask = nullptr;
823
824 if (prediff && mask) {
825 premask = lookup(mask, Builder2);
826 }
827
828 if (prediff)
829 ((DiffeGradientUtils *)gutils)
830 ->addToInvertedPtrDiffe(&I, &I, vd, LoadSize, I.getOperand(0),
831 prediff, Builder2, alignment, premask);
832
833 unsigned start = 0;
834 unsigned size = LoadSize;
835
836 while (1) {
837 unsigned nextStart = size;
838
839 auto dt = vd[{-1}];
840 for (size_t i = start; i < size; ++i) {
841 bool Legal = true;
842 dt.checkedOrIn(vd[{(int)i}], /*PointerIntSame*/ true, Legal);
843 if (!Legal) {
844 nextStart = i;
845 break;
846 }
847 }
848 if (!dt.isKnown()) {
849 std::string str;
850 raw_string_ostream ss(str);
851 ss << "Cannot deduce type of load " << I;
852 ss << " vd:" << vd.str() << " start:" << start << " size: " << size
853 << " dt:" << dt.str() << "\n";
854 EmitNoTypeError(str, I, gutils, BuilderZ);
855 continue;
856 }
857 assert(dt.isKnown());
858
859 if (Type *isfloat = dt.isFloat()) {
860 if (premask && !gutils->isConstantValue(orig_maskInit)) {
861 // Masked partial type is unhanled.
862 if (premask)
863 assert(start == 0 && nextStart == LoadSize);
864 addToDiffe(orig_maskInit, prediff, Builder2, isfloat,
865 Builder2.CreateNot(premask));
866 }
867 }
868
869 if (nextStart == size)
870 break;
871 start = nextStart;
872 }
873 break;
874 }
876 break;
877 }
878 }
879 }
880
881 void visitLoadInst(llvm::LoadInst &LI) {
882 using namespace llvm;
883
884 // If a load of an omp init argument, don't cache for reverse
885 // and don't do any adjoint propagation (assumed integral)
886 for (auto U : LI.getPointerOperand()->users()) {
887 if (auto CI = dyn_cast<CallInst>(U)) {
888 if (auto F = CI->getCalledFunction()) {
889 if (F->getName() == "__kmpc_for_static_init_4" ||
890 F->getName() == "__kmpc_for_static_init_4u" ||
891 F->getName() == "__kmpc_for_static_init_8" ||
892 F->getName() == "__kmpc_for_static_init_8u") {
893 eraseIfUnused(LI);
894 return;
895 }
896 }
897 }
898 }
899
900 auto alignment = LI.getAlign();
901 auto &DL = gutils->newFunc->getParent()->getDataLayout();
902
903 bool constantval = parseTBAA(LI, DL, nullptr)[{-1}].isIntegral();
904 visitLoadLike(LI, alignment, constantval);
905 eraseIfUnused(LI);
906 }
907
908 void visitAtomicRMWInst(llvm::AtomicRMWInst &I) {
909 using namespace llvm;
910
911 if (gutils->isConstantInstruction(&I) && gutils->isConstantValue(&I)) {
914 eraseIfUnused(I, /*erase*/ true, /*check*/ false);
915 } else {
916 eraseIfUnused(I);
917 }
918 return;
919 }
920
921 IRBuilder<> BuilderZ(&I);
922 getForwardBuilder(BuilderZ);
923
924 switch (I.getOperation()) {
925 case AtomicRMWInst::FAdd:
926 case AtomicRMWInst::FSub: {
927
928 if (Mode == DerivativeMode::ForwardMode ||
931 auto rule = [&](Value *ptr, Value *dif) -> Value * {
932 if (dif == nullptr)
933 dif = Constant::getNullValue(I.getType());
934 if (!gutils->isConstantInstruction(&I)) {
935 assert(ptr);
936 AtomicRMWInst *rmw = nullptr;
937 rmw = BuilderZ.CreateAtomicRMW(I.getOperation(), ptr, dif,
938 I.getAlign(), I.getOrdering(),
939 I.getSyncScopeID());
940 rmw->setVolatile(I.isVolatile());
941 if (gutils->isConstantValue(&I))
942 return Constant::getNullValue(dif->getType());
943 else
944 return rmw;
945 } else {
946 assert(gutils->isConstantValue(&I));
947 return Constant::getNullValue(dif->getType());
948 }
949 };
950
951 Value *diff = applyChainRule(
952 I.getType(), BuilderZ, rule,
953 gutils->isConstantValue(I.getPointerOperand())
954 ? nullptr
955 : gutils->invertPointerM(I.getPointerOperand(), BuilderZ),
956 gutils->isConstantValue(I.getValOperand())
957 ? nullptr
958 : gutils->invertPointerM(I.getValOperand(), BuilderZ,
959 TR.query(&I)));
960 if (!gutils->isConstantValue(&I))
961 setDiffe(&I, diff, BuilderZ);
962 return;
963 }
965 eraseIfUnused(I);
966 return;
967 }
970 gutils->isConstantValue(&I)) {
971 if (!gutils->isConstantValue(I.getValOperand())) {
973 I, I.getPointerOperand(), I.getValOperand(), I.getType(),
974 I.getType(), /*negate*/ I.getOperation() == AtomicRMWInst::FSub,
975 I.getAlign(), I.getOrdering(), I.getSyncScopeID(),
976 I.isVolatile());
977 }
979 eraseIfUnused(I, /*erase*/ true, /*check*/ false);
980 } else
981 eraseIfUnused(I);
982 return;
983 }
984 break;
985 }
986 default:
987 break;
988 }
989
990 if (looseTypeAnalysis) {
991 auto &DL = gutils->newFunc->getParent()->getDataLayout();
992 auto valType = I.getValOperand()->getType();
993 auto storeSize = DL.getTypeSizeInBits(valType) / 8;
994 auto fp = TR.firstPointer(storeSize, I.getPointerOperand(), &I, gutils,
995 /*errifnotfound*/ nullptr,
996 /*pointerIntSame*/ true);
997 if (!fp.isKnown() && valType->isIntOrIntVectorTy()) {
1000 eraseIfUnused(I, /*erase*/ true, /*check*/ false);
1001 } else
1002 eraseIfUnused(I);
1003 return;
1004 }
1005 }
1006 std::string s;
1007 llvm::raw_string_ostream ss(s);
1008 ss << *I.getParent()->getParent() << "\n" << I << "\n";
1009 ss << " Active atomic inst not yet handled";
1010 EmitNoDerivativeError(ss.str(), I, gutils, BuilderZ);
1011 if (!gutils->isConstantValue(&I)) {
1012 if (Mode == DerivativeMode::ForwardMode ||
1015 setDiffe(&I, Constant::getNullValue(gutils->getShadowType(I.getType())),
1016 BuilderZ);
1017 }
1018 if (!I.getType()->isVoidTy()) {
1019 for (auto &U :
1020 make_early_inc_range(gutils->getNewFromOriginal(&I)->uses())) {
1021 U.set(UndefValue::get(I.getType()));
1022 }
1023 }
1024 eraseIfUnused(I, /*erase*/ true, /*check*/ false);
1025 return;
1026 }
1027
1028 void visitStoreInst(llvm::StoreInst &SI) {
1029 using namespace llvm;
1030
1031 // If a store of an omp init argument, don't delete in reverse
1032 // and don't do any adjoint propagation (assumed integral)
1033 for (auto U : SI.getPointerOperand()->users()) {
1034 if (auto CI = dyn_cast<CallInst>(U)) {
1035 if (auto F = CI->getCalledFunction()) {
1036 if (F->getName() == "__kmpc_for_static_init_4" ||
1037 F->getName() == "__kmpc_for_static_init_4u" ||
1038 F->getName() == "__kmpc_for_static_init_8" ||
1039 F->getName() == "__kmpc_for_static_init_8u") {
1040 return;
1041 }
1042 }
1043 }
1044 }
1045 auto align = SI.getAlign();
1046
1047 visitCommonStore(SI, SI.getPointerOperand(), SI.getValueOperand(), align,
1048 SI.isVolatile(), SI.getOrdering(), SI.getSyncScopeID(),
1049 /*mask=*/nullptr);
1050
1051 bool forceErase = false;
1053 // Since we won't redo the store in the reverse pass, do not
1054 // force the write barrier.
1055 forceErase = true;
1056 for (const auto &pair : gutils->rematerializableAllocations) {
1057 // However, if we are rematerailizing the allocationa and not
1058 // inside the loop level rematerialization, we do still need the
1059 // reverse passes ``fake primal'' store and therefore write barrier
1060 if (gutils->allocationsToBeRematerialized.count(pair.first) &&
1061 pair.second.stores.count(&SI) &&
1062 (!pair.second.LI || !pair.second.LI->contains(&SI))) {
1063 forceErase = false;
1064 }
1065 }
1066 }
1067 if (forceErase)
1068 eraseIfUnused(SI, /*erase*/ true, /*check*/ false);
1069 else
1070 eraseIfUnused(SI);
1071 }
1072
1073 void visitCommonStore(llvm::Instruction &I, llvm::Value *orig_ptr,
1074 llvm::Value *orig_val, llvm::MaybeAlign prevalign,
1075 bool isVolatile, llvm::AtomicOrdering ordering,
1076 llvm::SyncScope::ID syncScope, llvm::Value *mask) {
1077 using namespace llvm;
1078
1079 Value *val = gutils->getNewFromOriginal(orig_val);
1080 Type *valType = orig_val->getType();
1081
1082 auto &DL = gutils->newFunc->getParent()->getDataLayout();
1083
1084 if (unnecessaryStores.count(&I)) {
1085 return;
1086 }
1087
1088 if (gutils->isConstantValue(orig_ptr)) {
1089 return;
1090 }
1091
1092 SmallVector<Metadata *, 1> scopeMD = {
1093 gutils->getDerivativeAliasScope(orig_ptr, -1)};
1094 SmallVector<Metadata *, 1> prevScopes;
1095 if (auto prev = I.getMetadata(LLVMContext::MD_alias_scope)) {
1096 for (auto &M : cast<MDNode>(prev)->operands()) {
1097 scopeMD.push_back(M);
1098 prevScopes.push_back(M);
1099 }
1100 }
1101 auto scope = MDNode::get(I.getContext(), scopeMD);
1102 auto NewI = gutils->getNewFromOriginal(&I);
1103 NewI->setMetadata(LLVMContext::MD_alias_scope, scope);
1104
1105 SmallVector<Metadata *, 1> MDs;
1106 SmallVector<Metadata *, 1> prevNoAlias;
1107 for (size_t j = 0; j < gutils->getWidth(); j++) {
1108 MDs.push_back(gutils->getDerivativeAliasScope(orig_ptr, j));
1109 }
1110 if (auto prev = I.getMetadata(LLVMContext::MD_noalias)) {
1111 for (auto &M : cast<MDNode>(prev)->operands()) {
1112 MDs.push_back(M);
1113 prevNoAlias.push_back(M);
1114 }
1115 }
1116 auto noscope = MDNode::get(I.getContext(), MDs);
1117 NewI->setMetadata(LLVMContext::MD_noalias, noscope);
1118
1119 bool constantval = gutils->isConstantValue(orig_val) ||
1120 parseTBAA(I, DL, nullptr)[{-1}].isIntegral();
1121
1122 IRBuilder<> BuilderZ(NewI);
1123 BuilderZ.setFastMathFlags(getFast());
1124
1125 // TODO allow recognition of other types that could contain pointers [e.g.
1126 // {void*, void*} or <2 x i64> ]
1127 auto storeSize = (DL.getTypeSizeInBits(valType) + 7) / 8;
1128
1129 auto vd = TR.query(orig_ptr).Lookup(storeSize, DL);
1130
1131 if (!vd.isKnown()) {
1132 if (looseTypeAnalysis || true) {
1133 vd = defaultTypeTreeForLLVM(valType, &I);
1134 if (EmitWarningEnabled(I.getContext())) {
1135 std::string str;
1136 raw_string_ostream ss(str);
1137 ss << "Cannot deduce type of store " << I << ", assumed " << vd.str()
1138 << "\n";
1139 EmitWarning("CannotDeduceType", I, ss.str());
1140 }
1141 goto known;
1142 }
1143 {
1144 std::string str;
1145 raw_string_ostream ss(str);
1146 ss << "Cannot deduce type of store " << I;
1147 EmitNoTypeError(str, I, gutils, BuilderZ);
1148 }
1149 return;
1150 known:;
1151 }
1152
1153 if (Mode == DerivativeMode::ForwardMode ||
1155
1156 bool anyPointer = false;
1157 // Only need the full type in forward mode, if storing a constant
1158 // and therefore may need to zero some floats.
1159 if (constantval) {
1160 // Bytes of the destination that nothing else types (e.g. the zero
1161 // lanes of a constant vector, which type analysis leaves as Anything
1162 // and does not propagate to memory) take the stored value's own type.
1163 TypeTree valTT = TR.query(orig_val);
1164 for (size_t i = 0; i < storeSize;) {
1165 ConcreteType ct = vd[{(int)i}];
1166 if (!ct.isKnown())
1167 ct = valTT[{(int)i}];
1168 if (auto flt = ct.isFloat()) {
1169 i += DL.getTypeSizeInBits(flt) / 8;
1170 continue;
1171 }
1172 if (ct == BaseType::Pointer) {
1173 anyPointer = true;
1174 i += DL.getPointerSizeInBits() / 8;
1175 continue;
1176 }
1177 if (ct == BaseType::Integer || ct == BaseType::Anything) {
1178 i++;
1179 continue;
1180 }
1181 std::string str;
1182 raw_string_ostream ss(str);
1183 ss << "Cannot deduce type of store " << I << vd.str()
1184 << " size: " << storeSize << " at " << i;
1185 EmitNoTypeError(str, I, gutils, BuilderZ);
1186 return;
1187 }
1188 }
1189
1190 Value *diff = nullptr;
1191 bool needs_writebarrier = false;
1192 if (!gutils->runtimeActivity && constantval) {
1193 if (anyPointer && vd[{-1, -1}] != BaseType::Integer) {
1194 if (!isa<UndefValue>(orig_val) &&
1195 !isa<ConstantPointerNull>(orig_val)) {
1196 std::string str;
1197 raw_string_ostream ss(str);
1198 ss << "Mismatched activity for: " << I
1199 << " const val: " << *orig_val;
1200 if (CustomErrorHandler) {
1201 diff = unwrap(CustomErrorHandler(
1202 str.c_str(), wrap(&I), ErrorType::MixedActivityError, gutils,
1203 wrap(orig_val), wrap(&BuilderZ)));
1204 if (diff)
1205 needs_writebarrier = true;
1206 } else
1207 EmitWarningAlways("MixedActivityError", I, ss.str(),
1209 }
1210 }
1211 }
1212
1213 // The shadow of integer-only memory mirrors the primal.
1214 if (!diff && vd.isKnown() && !vd.anyPointer(orig_val, DL) &&
1215 !vd.anyFloat(orig_val, DL)) {
1216 auto rule = [&val]() { return val; };
1217 diff = applyChainRule(valType, BuilderZ, rule);
1218 }
1219
1220 // TODO type analyze
1221 if (!diff) {
1222 diff = gutils->invertPointerM(orig_val, BuilderZ, vd);
1223 }
1224
1225 gutils->setPtrDiffe(&I, orig_ptr, diff, BuilderZ, prevalign, 0, storeSize,
1226 isVolatile, ordering, syncScope, mask, prevNoAlias,
1227 prevScopes, needs_writebarrier);
1228
1229 return;
1230 }
1231
1232 unsigned start = 0;
1233
1234 IRBuilder<> Builder2(&I);
1235 BasicBlock *merge = nullptr;
1238 getReverseBuilder(Builder2);
1239
1240 while (1) {
1241 unsigned nextStart = storeSize;
1242
1243 auto dt = vd[{-1}];
1244 for (size_t i = start; i < storeSize; ++i) {
1245 auto nex = vd[{(int)i}];
1246 if ((nex == BaseType::Anything && dt.isFloat()) ||
1247 (dt == BaseType::Anything && nex.isFloat())) {
1248 nextStart = i;
1249 break;
1250 }
1251 bool Legal = true;
1252 dt.checkedOrIn(nex, /*PointerIntSame*/ true, Legal);
1253 if (!Legal) {
1254 nextStart = i;
1255 break;
1256 }
1257 }
1258 unsigned size = nextStart - start;
1259 if (!dt.isKnown()) {
1260
1261 std::string str;
1262 raw_string_ostream ss(str);
1263 ss << "Cannot deduce type of store " << I << vd.str()
1264 << " start: " << start << " size: " << size
1265 << " storeSize: " << storeSize;
1266 EmitNoTypeError(str, I, gutils, BuilderZ);
1267 break;
1268 }
1269
1270 MaybeAlign align;
1271 if (prevalign) {
1272 if (start % prevalign->value() == 0)
1273 align = prevalign;
1274 else
1275 align = Align(1);
1276 }
1277 //! Storing a floating point value
1278 if (Type *FT = dt.isFloat()) {
1279 //! Only need to update the reverse function
1280 switch (Mode) {
1282 break;
1285
1286 if (!merge && gutils->runtimeActivity) {
1287 auto basePtr = getBaseObject(orig_ptr);
1288
1289 // If runtime activity, first see if we can prove that the
1290 // shadow/primal are distinct statically as they are
1291 // allocas/mallocs, if not compare the pointers and conditionally
1292 // execute.
1293 if (!isa<AllocaInst>(basePtr) &&
1294 !isAllocationCall(basePtr, gutils->TLI)) {
1295 auto shadow_ptr =
1296 lookup(gutils->invertPointerM(orig_ptr, Builder2), Builder2);
1297 auto primal_ptr =
1298 lookup(gutils->getNewFromOriginal(orig_ptr), Builder2);
1299 if (gutils->getWidth() != 1) {
1300 shadow_ptr = gutils->extractMeta(Builder2, shadow_ptr, 0);
1301 }
1302 Value *shadow = Builder2.CreateICmpNE(primal_ptr, shadow_ptr);
1303
1304 BasicBlock *current = Builder2.GetInsertBlock();
1305 BasicBlock *conditional = gutils->addReverseBlock(
1306 current, current->getName() + "_active");
1307 merge = gutils->addReverseBlock(conditional,
1308 current->getName() + "_amerge");
1309 Builder2.CreateCondBr(shadow, conditional, merge);
1310 Builder2.SetInsertPoint(conditional);
1311 }
1312 }
1313
1314 if (constantval) {
1315 gutils->setPtrDiffe(
1316 &I, orig_ptr,
1317 Constant::getNullValue(gutils->getShadowType(valType)),
1318 Builder2, align, start, size, isVolatile, ordering, syncScope,
1319 mask, prevNoAlias, prevScopes);
1320 } else {
1321 Value *diff;
1322 Value *maskL = mask;
1323 if (!mask) {
1324 Value *dif1Ptr =
1325 lookup(gutils->invertPointerM(orig_ptr, Builder2), Builder2);
1326
1327 size_t idx = 0;
1328 auto rule = [&](Value *dif1Ptr) {
1329 LoadInst *dif1 =
1330 Builder2.CreateLoad(valType, dif1Ptr, isVolatile);
1331 if (align)
1332 dif1->setAlignment(*align);
1333 dif1->setOrdering(ordering);
1334 dif1->setSyncScopeID(syncScope);
1335
1336 SmallVector<Metadata *, 1> scopeMD = {
1337 gutils->getDerivativeAliasScope(orig_ptr, idx)};
1338 for (auto M : prevScopes)
1339 scopeMD.push_back(M);
1340
1341 SmallVector<Metadata *, 1> MDs;
1342 for (ssize_t j = -1; j < gutils->getWidth(); j++) {
1343 if (j != (ssize_t)idx)
1344 MDs.push_back(gutils->getDerivativeAliasScope(orig_ptr, j));
1345 }
1346 for (auto M : prevNoAlias)
1347 MDs.push_back(M);
1348
1349 dif1->setMetadata(LLVMContext::MD_alias_scope,
1350 MDNode::get(I.getContext(), scopeMD));
1351 dif1->setMetadata(LLVMContext::MD_noalias,
1352 MDNode::get(I.getContext(), MDs));
1353 dif1->setMetadata(LLVMContext::MD_tbaa,
1354 I.getMetadata(LLVMContext::MD_tbaa));
1355 dif1->setMetadata(LLVMContext::MD_tbaa_struct,
1356 I.getMetadata(LLVMContext::MD_tbaa_struct));
1357 idx++;
1358 return dif1;
1359 };
1360
1361 diff = applyChainRule(valType, Builder2, rule, dif1Ptr);
1362 } else {
1363 maskL = lookup(mask, Builder2);
1364 Type *tys[] = {valType, orig_ptr->getType()};
1365 auto F = getIntrinsicDeclaration(gutils->oldFunc->getParent(),
1366 Intrinsic::masked_load, tys);
1367 Value *alignv =
1368 ConstantInt::get(Type::getInt32Ty(mask->getContext()),
1369 align ? align->value() : 0);
1370 Value *ip =
1371 lookup(gutils->invertPointerM(orig_ptr, Builder2), Builder2);
1372
1373 auto rule = [&](Value *ip) {
1374 Value *args[] = {ip, alignv, maskL,
1375 Constant::getNullValue(valType)};
1376 diff = Builder2.CreateCall(F, args);
1377 return diff;
1378 };
1379
1380 diff = applyChainRule(valType, Builder2, rule, ip);
1381 }
1382
1383 gutils->setPtrDiffe(
1384 &I, orig_ptr,
1385 Constant::getNullValue(gutils->getShadowType(valType)),
1386 Builder2, align, start, size, isVolatile, ordering, syncScope,
1387 mask, prevNoAlias, prevScopes);
1388 ((DiffeGradientUtils *)gutils)
1389 ->addToDiffe(orig_val, diff, Builder2, FT, start, size, {},
1390 maskL);
1391 }
1392 break;
1393 }
1397
1398 Type *diffeTy = gutils->getShadowType(valType);
1399
1400 Value *diff = constantval
1401 ? Constant::getNullValue(diffeTy)
1402 : gutils->invertPointerM(orig_val, BuilderZ, vd);
1403 gutils->setPtrDiffe(&I, orig_ptr, diff, BuilderZ, align, start, size,
1404 isVolatile, ordering, syncScope, mask,
1405 prevNoAlias, prevScopes);
1406
1407 break;
1408 }
1409 }
1410
1411 //! Storing an integer or pointer
1412 } else {
1413 //! Only need to update the forward function
1414
1415 // Don't reproduce mpi null requests
1416 if (constantval)
1417 if (Constant *C = dyn_cast<Constant>(orig_val)) {
1418 while (ConstantExpr *CE = dyn_cast<ConstantExpr>(C)) {
1419 C = CE->getOperand(0);
1420 }
1421 if (auto GV = dyn_cast<GlobalVariable>(C)) {
1422 if (GV->getName() == "ompi_request_null") {
1423 continue;
1424 }
1425 }
1426 }
1427
1428 bool forwardsShadow, backwardsShadow;
1429 shadowStoreSchedule(I, orig_ptr, forwardsShadow, backwardsShadow);
1430
1431 if ((Mode == DerivativeMode::ReverseModePrimal && forwardsShadow) ||
1432 (Mode == DerivativeMode::ReverseModeGradient && backwardsShadow) ||
1433 (Mode == DerivativeMode::ForwardModeSplit && backwardsShadow) ||
1435 (forwardsShadow || backwardsShadow)) ||
1438
1439 Value *valueop = nullptr;
1440
1441 bool needs_writebarrier = false;
1442 if (constantval) {
1443 if (!gutils->runtimeActivity) {
1444 if (dt.isPossiblePointer() && vd[{-1, -1}] != BaseType::Integer) {
1445 if (!isa<UndefValue>(orig_val) &&
1446 !isa<ConstantPointerNull>(orig_val)) {
1447 std::string str;
1448 raw_string_ostream ss(str);
1449 ss << "Mismatched activity for: " << I
1450 << " const val: " << *orig_val;
1451 if (CustomErrorHandler) {
1452 valueop = unwrap(CustomErrorHandler(
1453 str.c_str(), wrap(&I), ErrorType::MixedActivityError,
1454 gutils, wrap(orig_val), wrap(&BuilderZ)));
1455 if (valueop)
1456 needs_writebarrier = true;
1457 } else
1458 EmitWarningAlways("MixedActivityError", I, ss.str(),
1460 }
1461 }
1462 }
1463 if (!valueop) {
1464 valueop = val;
1465 if (gutils->getWidth() > 1) {
1466 Value *array =
1467 UndefValue::get(gutils->getShadowType(val->getType()));
1468 for (unsigned i = 0; i < gutils->getWidth(); ++i) {
1469 array = BuilderZ.CreateInsertValue(array, val, {i});
1470 }
1471 valueop = array;
1472 }
1473 }
1474 } else {
1475 valueop = gutils->invertPointerM(orig_val, BuilderZ, vd);
1476 }
1477 gutils->setPtrDiffe(&I, orig_ptr, valueop, BuilderZ, align, start,
1478 size, isVolatile, ordering, syncScope, mask,
1479 prevNoAlias, prevScopes, needs_writebarrier);
1480 }
1481 }
1482
1483 if (nextStart == storeSize)
1484 break;
1485 start = nextStart;
1486 }
1487
1488 if (merge) {
1489 Builder2.CreateBr(merge);
1490 Builder2.SetInsertPoint(merge);
1491 }
1492 }
1493
1494 void visitGetElementPtrInst(llvm::GetElementPtrInst &gep) {
1495 eraseIfUnused(gep);
1496 switch (Mode) {
1501 return;
1502 }
1503 default:
1504 return;
1505 }
1506 }
1507
1508 void visitPHINode(llvm::PHINode &phi) {
1509 eraseIfUnused(phi);
1510
1511 switch (Mode) {
1515 return;
1516 }
1521 return;
1522 }
1523 }
1524 }
1525
1526 void visitCastInst(llvm::CastInst &I) {
1527 using namespace llvm;
1528
1529 eraseIfUnused(I);
1530
1531 switch (Mode) {
1533 return;
1534 }
1537 if (gutils->isConstantInstruction(&I))
1538 return;
1539
1540 if (I.getType()->isPointerTy() ||
1541 I.getOpcode() == CastInst::CastOps::PtrToInt)
1542 return;
1543
1544 Value *orig_op0 = I.getOperand(0);
1545 Value *op0 = gutils->getNewFromOriginal(orig_op0);
1546
1547 IRBuilder<> Builder2(&I);
1548 getReverseBuilder(Builder2);
1549
1550 if (!gutils->isConstantValue(orig_op0)) {
1551 size_t size = 1;
1552 if (orig_op0->getType()->isSized())
1553 size =
1554 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
1555 orig_op0->getType()) +
1556 7) /
1557 8;
1558 Type *FT = TR.addingType(size, orig_op0);
1559 if (!FT && looseTypeAnalysis) {
1560 if (auto ET = I.getSrcTy()->getScalarType())
1561 if (ET->isFPOrFPVectorTy()) {
1562 FT = ET;
1563 EmitWarning("CannotDeduceType", I,
1564 "failed to deduce adding type of cast ", I,
1565 " assumed ", FT, " from src");
1566 }
1567 }
1568 if (!FT && looseTypeAnalysis) {
1569 if (auto ET = I.getDestTy()->getScalarType())
1570 if (ET->isFPOrFPVectorTy()) {
1571 FT = ET;
1572 EmitWarning("CannotDeduceType", I,
1573 "failed to deduce adding type of cast ", I,
1574 " assumed ", FT, " from dst");
1575 }
1576 }
1577 if (!FT) {
1578 if (TR.query(orig_op0)[{-1}] == BaseType::Integer &&
1579 TR.query(&I)[{-1}] == BaseType::Integer)
1580 return;
1581 if (looseTypeAnalysis) {
1582 if (auto ET = I.getSrcTy()->getScalarType())
1583 if (ET->isIntOrIntVectorTy()) {
1584 EmitWarning("CannotDeduceType", I,
1585 "failed to deduce adding type of cast ", I,
1586 " assumed integral from src");
1587 return;
1588 }
1589 }
1590 std::string str;
1591 raw_string_ostream ss(str);
1592 ss << "Cannot deduce adding type (cast) of " << I;
1593 EmitNoTypeError(str, I, gutils, Builder2);
1594 }
1595
1596 if (FT) {
1597
1598 auto rule = [&](Value *dif) {
1599 if (I.getOpcode() == CastInst::CastOps::FPTrunc ||
1600 I.getOpcode() == CastInst::CastOps::FPExt) {
1601 return Builder2.CreateFPCast(dif, op0->getType());
1602 } else if (I.getOpcode() == CastInst::CastOps::BitCast) {
1603 return Builder2.CreateBitCast(dif, op0->getType());
1604 } else if (I.getOpcode() == CastInst::CastOps::Trunc) {
1605 // TODO CHECK THIS
1606 return Builder2.CreateZExt(dif, op0->getType());
1607 } else {
1608 std::string s;
1609 llvm::raw_string_ostream ss(s);
1610 ss << *I.getParent()->getParent() << "\n";
1611 ss << "cannot handle above cast " << I << "\n";
1612 EmitNoDerivativeError(ss.str(), I, gutils, Builder2);
1613 return (llvm::Value *)UndefValue::get(op0->getType());
1614 }
1615 };
1616
1617 Value *dif = diffe(&I, Builder2);
1618 Value *diff = applyChainRule(op0->getType(), Builder2, rule, dif);
1619
1620 addToDiffe(orig_op0, diff, Builder2, FT);
1621 }
1622 }
1623
1624 Type *diffTy = gutils->getShadowType(I.getType());
1625 setDiffe(&I, Constant::getNullValue(diffTy), Builder2);
1626
1627 break;
1628 }
1633 return;
1634 }
1635 }
1636 }
1637
1638 void visitSelectInst(llvm::SelectInst &SI) {
1639 eraseIfUnused(SI);
1640
1641 switch (Mode) {
1643 return;
1646 if (gutils->isConstantInstruction(&SI))
1647 return;
1648 if (SI.getType()->isPointerTy())
1649 return;
1651 return;
1652 }
1657 return;
1658 }
1659 }
1660 }
1661
1662 void createSelectInstAdjoint(llvm::SelectInst &SI) {
1663 using namespace llvm;
1664
1665 Value *op0 = gutils->getNewFromOriginal(SI.getOperand(0));
1666 Value *orig_op1 = SI.getOperand(1);
1667 Value *op1 = gutils->getNewFromOriginal(orig_op1);
1668 Value *orig_op2 = SI.getOperand(2);
1669 Value *op2 = gutils->getNewFromOriginal(orig_op2);
1670
1671 // TODO fix all the reverse builders
1672 IRBuilder<> Builder2(&SI);
1673 getReverseBuilder(Builder2);
1674
1675 Value *dif1 = nullptr;
1676 Value *dif2 = nullptr;
1677
1678 size_t size = 1;
1679 if (orig_op1->getType()->isSized())
1680 size = (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
1681 orig_op1->getType()) +
1682 7) /
1683 8;
1684 // Required loopy phi = [in, BO, BO, ..., BO]
1685 // 1) phi is only used in this B0
1686 // 2) BO dominates all latches
1687 // 3) phi == B0 whenever not coming from preheader [implies 2]
1688 // 4) [optional but done for ease] one exit to make it easier to
1689 // calculation the product at that point
1690 for (int i = 0; i < 2; i++)
1691 if (auto P0 = dyn_cast<PHINode>(SI.getOperand(i + 1))) {
1692 LoopContext lc;
1693 SmallVector<Instruction *, 4> activeUses;
1694 for (auto u : P0->users()) {
1695 if (!gutils->isConstantInstruction(cast<Instruction>(u))) {
1696 activeUses.push_back(cast<Instruction>(u));
1697 } else if (retType == DIFFE_TYPE::OUT_DIFF && isa<ReturnInst>(u))
1698 activeUses.push_back(cast<Instruction>(u));
1699 }
1700 if (activeUses.size() == 1 && activeUses[0] == &SI &&
1701 gutils->getContext(gutils->getNewFromOriginal(P0->getParent()),
1702 lc) &&
1703 gutils->getNewFromOriginal(P0->getParent()) == lc.header) {
1704 SmallVector<BasicBlock *, 1> Latches;
1705 gutils->OrigLI->getLoopFor(P0->getParent())->getLoopLatches(Latches);
1706 bool allIncoming = true;
1707 for (auto Latch : Latches) {
1708 if (&SI != P0->getIncomingValueForBlock(Latch)) {
1709 allIncoming = false;
1710 break;
1711 }
1712 }
1713 if (allIncoming && lc.exitBlocks.size() == 1) {
1714 if (!gutils->isConstantValue(SI.getOperand(2 - i))) {
1715 auto addingType = TR.addingType(size, SI.getOperand(2 - i));
1716 if (addingType || !looseTypeAnalysis) {
1717 auto index = gutils->getOrInsertConditionalIndex(
1718 gutils->getNewFromOriginal(SI.getOperand(0)), lc, i == 1);
1719 IRBuilder<> EB(*lc.exitBlocks.begin());
1720 getReverseBuilder(EB, /*original=*/false);
1721 Value *inc = lookup(lc.incvar, Builder2);
1722 if (VectorType *VTy =
1723 dyn_cast<VectorType>(SI.getOperand(0)->getType())) {
1724 inc = Builder2.CreateVectorSplat(VTy->getElementCount(), inc);
1725 }
1726 Value *difSI = diffe(&SI, Builder2);
1727 Value *isSelected =
1728 Builder2.CreateICmpEQ(gutils->lookupM(index, EB), inc);
1729 Value *dif =
1730 CreateSelect(Builder2, isSelected, difSI,
1731 Constant::getNullValue(
1732 gutils->getShadowType(op1->getType())));
1733 addToDiffe(SI.getOperand(2 - i), dif, Builder2, addingType);
1734 }
1735 }
1736 return;
1737 }
1738 }
1739 }
1740
1741 if (!gutils->isConstantValue(orig_op1))
1742 dif1 = CreateSelect(
1743 Builder2, lookup(op0, Builder2), diffe(&SI, Builder2),
1744 Constant::getNullValue(gutils->getShadowType(op1->getType())),
1745 "diffe" + op1->getName());
1746 if (!gutils->isConstantValue(orig_op2))
1747 dif2 = CreateSelect(
1748 Builder2, lookup(op0, Builder2),
1749 Constant::getNullValue(gutils->getShadowType(op2->getType())),
1750 diffe(&SI, Builder2), "diffe" + op2->getName());
1751
1752 setDiffe(&SI, Constant::getNullValue(gutils->getShadowType(SI.getType())),
1753 Builder2);
1754 if (dif1) {
1755 Type *addingType = TR.addingType(size, orig_op1);
1756 if (addingType || !looseTypeAnalysis)
1757 addToDiffe(orig_op1, dif1, Builder2, addingType);
1758 else
1759 llvm::errs() << " warning: assuming integral for " << SI << "\n";
1760 }
1761 if (dif2) {
1762 Type *addingType = TR.addingType(size, orig_op2);
1763 if (addingType || !looseTypeAnalysis)
1764 addToDiffe(orig_op2, dif2, Builder2, addingType);
1765 else
1766 llvm::errs() << " warning: assuming integral for " << SI << "\n";
1767 }
1768 }
1769
1770 void visitExtractElementInst(llvm::ExtractElementInst &EEI) {
1771 using namespace llvm;
1772
1773 eraseIfUnused(EEI);
1774 switch (Mode) {
1779 return;
1780 }
1783 if (gutils->isConstantInstruction(&EEI))
1784 return;
1785 IRBuilder<> Builder2(&EEI);
1786 getReverseBuilder(Builder2);
1787
1788 Value *orig_vec = EEI.getVectorOperand();
1789
1790 if (!gutils->isConstantValue(orig_vec)) {
1791
1792 size_t size = 1;
1793 if (EEI.getType()->isSized())
1794 size =
1795 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
1796 EEI.getType()) +
1797 7) /
1798 8;
1799 auto diff = diffe(&EEI, Builder2);
1800 if (gutils->getWidth() == 1) {
1801 Value *sv[] = {gutils->getNewFromOriginal(EEI.getIndexOperand())};
1802 ((DiffeGradientUtils *)gutils)
1803 ->addToDiffe(orig_vec, diff, Builder2, TR.addingType(size, &EEI),
1804 sv);
1805 } else {
1806 for (size_t i = 0; i < gutils->getWidth(); i++) {
1807 Value *sv[] = {nullptr,
1808 gutils->getNewFromOriginal(EEI.getIndexOperand())};
1809 sv[0] = ConstantInt::get(sv[1]->getType(), i);
1810 ((DiffeGradientUtils *)gutils)
1811 ->addToDiffe(orig_vec, gutils->extractMeta(Builder2, diff, i),
1812 Builder2, TR.addingType(size, &EEI), sv);
1813 }
1814 }
1815 }
1816 setDiffe(&EEI,
1817 Constant::getNullValue(gutils->getShadowType(EEI.getType())),
1818 Builder2);
1819 return;
1820 }
1822 return;
1823 }
1824 }
1825 }
1826
1827 void visitInsertElementInst(llvm::InsertElementInst &IEI) {
1828 using namespace llvm;
1829
1830 eraseIfUnused(IEI);
1831
1832 switch (Mode) {
1837 return;
1838 }
1841 if (gutils->isConstantInstruction(&IEI))
1842 return;
1843 IRBuilder<> Builder2(&IEI);
1844 getReverseBuilder(Builder2);
1845
1846 Value *dif1 = diffe(&IEI, Builder2);
1847
1848 Value *orig_op0 = IEI.getOperand(0);
1849 Value *orig_op1 = IEI.getOperand(1);
1850 Value *op1 = gutils->getNewFromOriginal(orig_op1);
1851 Value *op2 = gutils->getNewFromOriginal(IEI.getOperand(2));
1852
1853 size_t size0 = 1;
1854 if (orig_op0->getType()->isSized())
1855 size0 =
1856 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
1857 orig_op0->getType()) +
1858 7) /
1859 8;
1860 size_t size1 = 1;
1861 if (orig_op1->getType()->isSized())
1862 size1 =
1863 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
1864 orig_op1->getType()) +
1865 7) /
1866 8;
1867
1868 if (!gutils->isConstantValue(orig_op0)) {
1869 if (gutils->getWidth() == 1) {
1870 addToDiffe(
1871 orig_op0,
1872 Builder2.CreateInsertElement(
1873 dif1,
1874 Constant::getNullValue(gutils->getShadowType(op1->getType())),
1875 lookup(op2, Builder2)),
1876 Builder2, TR.addingType(size0, orig_op0));
1877 } else {
1878 for (size_t i = 0; i < gutils->getWidth(); i++) {
1879 Value *sv[] = {ConstantInt::get(op2->getType(), i)};
1880 ((DiffeGradientUtils *)gutils)
1881 ->addToDiffe(orig_op0,
1882 Builder2.CreateInsertElement(
1883 gutils->extractMeta(Builder2, dif1, i),
1884 Constant::getNullValue(op1->getType()),
1885 lookup(op2, Builder2)),
1886 Builder2, TR.addingType(size0, orig_op0), sv);
1887 }
1888 }
1889 }
1890
1891 if (!gutils->isConstantValue(orig_op1)) {
1892 if (gutils->getWidth() == 1) {
1893 addToDiffe(orig_op1,
1894 Builder2.CreateExtractElement(dif1, lookup(op2, Builder2)),
1895 Builder2, TR.addingType(size1, orig_op1));
1896 } else {
1897 for (size_t i = 0; i < gutils->getWidth(); i++) {
1898 Value *sv[] = {ConstantInt::get(op2->getType(), i)};
1899 ((DiffeGradientUtils *)gutils)
1900 ->addToDiffe(orig_op1,
1901 Builder2.CreateExtractElement(
1902 gutils->extractMeta(Builder2, dif1, i),
1903 lookup(op2, Builder2)),
1904 Builder2, TR.addingType(size1, orig_op1), sv);
1905 }
1906 }
1907 }
1908
1909 setDiffe(&IEI,
1910 Constant::getNullValue(gutils->getShadowType(IEI.getType())),
1911 Builder2);
1912 return;
1913 }
1915 return;
1916 }
1917 }
1918 }
1919
1920 void visitShuffleVectorInst(llvm::ShuffleVectorInst &SVI) {
1921 using namespace llvm;
1922
1923 eraseIfUnused(SVI);
1924
1925 switch (Mode) {
1930 return;
1931 }
1934 if (gutils->isConstantInstruction(&SVI))
1935 return;
1936 IRBuilder<> Builder2(&SVI);
1937 getReverseBuilder(Builder2);
1938
1939 auto loaded = diffe(&SVI, Builder2);
1940 auto count =
1941 cast<VectorType>(SVI.getOperand(0)->getType())->getElementCount();
1942 assert(!count.isScalable());
1943 size_t l1 = count.getKnownMinValue();
1944 uint64_t instidx = 0;
1945
1946 for (size_t idx : SVI.getShuffleMask()) {
1947 auto opnum = (idx < l1) ? 0 : 1;
1948 auto opidx = (idx < l1) ? idx : (idx - l1);
1949
1950 if (!gutils->isConstantValue(SVI.getOperand(opnum))) {
1951 size_t size = 1;
1952 if (SVI.getOperand(opnum)->getType()->isSized())
1953 size = (gutils->newFunc->getParent()
1954 ->getDataLayout()
1955 .getTypeSizeInBits(SVI.getOperand(opnum)->getType()) +
1956 7) /
1957 8;
1958 if (gutils->getWidth() == 1) {
1959 Value *sv[] = {
1960 ConstantInt::get(Type::getInt32Ty(SVI.getContext()), opidx)};
1961 Value *toadd = Builder2.CreateExtractElement(loaded, instidx);
1962 ((DiffeGradientUtils *)gutils)
1963 ->addToDiffe(SVI.getOperand(opnum), toadd, Builder2,
1964 TR.addingType(size, SVI.getOperand(opnum)), sv);
1965 } else {
1966 for (size_t i = 0; i < gutils->getWidth(); i++) {
1967 Value *sv[] = {
1968 ConstantInt::get(Type::getInt32Ty(SVI.getContext()), i),
1969 ConstantInt::get(Type::getInt32Ty(SVI.getContext()), opidx)};
1970 Value *toadd = Builder2.CreateExtractElement(
1971 GradientUtils::extractMeta(Builder2, loaded, i), instidx);
1972 ((DiffeGradientUtils *)gutils)
1973 ->addToDiffe(SVI.getOperand(opnum), toadd, Builder2,
1974 TR.addingType(size, SVI.getOperand(opnum)), sv);
1975 }
1976 }
1977 }
1978 ++instidx;
1979 }
1980 setDiffe(&SVI,
1981 Constant::getNullValue(gutils->getShadowType(SVI.getType())),
1982 Builder2);
1983 return;
1984 }
1986 return;
1987 }
1988 }
1989 }
1990
1991 void visitExtractValueInst(llvm::ExtractValueInst &EVI) {
1992 using namespace llvm;
1993
1994 eraseIfUnused(EVI);
1995
1996 if (!gutils->isConstantValue(&EVI) && gutils->isConstantValue(&EVI)) {
1997 llvm::errs() << *gutils->oldFunc->getParent() << "\n";
1998 llvm::errs() << EVI << "\n";
1999 llvm_unreachable("Illegal activity for extractvalue");
2000 }
2001
2002 switch (Mode) {
2007 return;
2008 }
2011 if (gutils->isConstantInstruction(&EVI))
2012 return;
2013 if (EVI.getType()->isPointerTy())
2014 return;
2015 IRBuilder<> Builder2(&EVI);
2016 getReverseBuilder(Builder2);
2017
2018 Value *orig_op0 = EVI.getOperand(0);
2019
2020 auto prediff = diffe(&EVI, Builder2);
2021
2022 // todo const
2023 if (!gutils->isConstantValue(orig_op0)) {
2024 SmallVector<Value *, 4> sv;
2025 for (auto i : EVI.getIndices())
2026 sv.push_back(ConstantInt::get(Type::getInt32Ty(EVI.getContext()), i));
2027 size_t storeSize = 1;
2028 if (EVI.getType()->isSized())
2029 storeSize =
2030 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
2031 EVI.getType()) +
2032 7) /
2033 8;
2034
2035 unsigned start = 0;
2036 const auto &vd = TR.query(&EVI);
2037
2038 while (1) {
2039 unsigned nextStart = storeSize;
2040
2041 auto dt = vd[{-1}];
2042 for (size_t i = start; i < storeSize; ++i) {
2043 auto nex = vd[{(int)i}];
2044 if ((nex == BaseType::Anything && dt.isFloat()) ||
2045 (dt == BaseType::Anything && nex.isFloat())) {
2046 nextStart = i;
2047 break;
2048 }
2049 bool Legal = true;
2050 dt.checkedOrIn(nex, /*PointerIntSame*/ true, Legal);
2051 if (!Legal) {
2052 nextStart = i;
2053 break;
2054 }
2055 }
2056 unsigned size = nextStart - start;
2057 if (!dt.isKnown()) {
2058 bool found = false;
2059 if (looseTypeAnalysis) {
2060 if (EVI.getType()->isFPOrFPVectorTy()) {
2061 dt = ConcreteType(EVI.getType()->getScalarType());
2062 found = true;
2063 } else if (EVI.getType()->isIntOrIntVectorTy() ||
2064 EVI.getType()->isPointerTy()) {
2065 dt = BaseType::Integer;
2066 found = true;
2067 }
2068 }
2069 if (!found) {
2070 std::string str;
2071 raw_string_ostream ss(str);
2072 ss << "Cannot deduce type of extract " << EVI << vd.str()
2073 << " start: " << start << " size: " << size
2074 << " extractSize: " << storeSize;
2075 EmitNoTypeError(str, EVI, gutils, Builder2);
2076 }
2077 }
2078 if (auto FT = dt.isFloat())
2079 ((DiffeGradientUtils *)gutils)
2080 ->addToDiffe(orig_op0, prediff, Builder2, FT, start, size, sv,
2081 nullptr, /*ignoreFirstSlicesToDiff*/ sv.size());
2082
2083 if (nextStart == storeSize)
2084 break;
2085 start = nextStart;
2086 }
2087 }
2088
2089 setDiffe(&EVI,
2090 Constant::getNullValue(gutils->getShadowType(EVI.getType())),
2091 Builder2);
2092 return;
2093 }
2095 return;
2096 }
2097 }
2098 }
2099
2100 void visitInsertValueInst(llvm::InsertValueInst &IVI) {
2101 using namespace llvm;
2102
2103 eraseIfUnused(IVI);
2104 if (gutils->isConstantValue(&IVI))
2105 return;
2106
2108 return;
2109
2110 if (Mode == DerivativeMode::ForwardMode ||
2114 return;
2115 }
2116
2117 bool hasNonPointer = false;
2118 if (auto st = dyn_cast<StructType>(IVI.getType())) {
2119 for (unsigned i = 0; i < st->getNumElements(); ++i) {
2120 if (!st->getElementType(i)->isPointerTy()) {
2121 hasNonPointer = true;
2122 }
2123 }
2124 } else if (auto at = dyn_cast<ArrayType>(IVI.getType())) {
2125 if (!at->getElementType()->isPointerTy()) {
2126 hasNonPointer = true;
2127 }
2128 }
2129 if (!hasNonPointer)
2130 return;
2131
2132 bool floatingInsertion = false;
2133 if (auto MD = hasMetadata(&IVI, "enzyme_truetype")) {
2134 auto toIterate = parseTrueType(MD, Mode, false);
2135 for (auto &tuple : toIterate) {
2136 Type *ty = std::get<0>(tuple);
2137 if (ty && ty->isFloatingPointTy()) {
2138 floatingInsertion = true;
2139 break;
2140 }
2141 }
2142 } else {
2143 for (InsertValueInst *iv = &IVI;;) {
2144 size_t size0 = 1;
2145 if (iv->getInsertedValueOperand()->getType()->isSized() &&
2146 (iv->getInsertedValueOperand()->getType()->isIntOrIntVectorTy() ||
2147 iv->getInsertedValueOperand()->getType()->isFPOrFPVectorTy()))
2148 size0 =
2149 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
2150 iv->getInsertedValueOperand()->getType()) +
2151 7) /
2152 8;
2153 auto it = TR.intType(size0, iv->getInsertedValueOperand(), iv, gutils,
2154 /*err*/ nullptr);
2155 if (it.isFloat() || !it.isKnown()) {
2156 floatingInsertion = true;
2157 break;
2158 }
2159 Value *val = iv->getAggregateOperand();
2160 if (gutils->isConstantValue(val))
2161 break;
2162 if (auto dc = dyn_cast<InsertValueInst>(val)) {
2163 iv = dc;
2164 } else {
2165 // unsure where this came from, conservatively assume contains float
2166 floatingInsertion = true;
2167 break;
2168 }
2169 }
2170 }
2171
2172 if (!floatingInsertion)
2173 return;
2174
2175 // TODO handle pointers
2176 // TODO type analysis handle structs
2177
2178 switch (Mode) {
2182 assert(0 && "should be handled above");
2183 return;
2186 IRBuilder<> Builder2(&IVI);
2187 getReverseBuilder(Builder2);
2188
2189 Value *orig_inserted = IVI.getInsertedValueOperand();
2190 Value *orig_agg = IVI.getAggregateOperand();
2191
2192 size_t size0 = 1;
2193 if (orig_inserted->getType()->isSized())
2194 size0 =
2195 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
2196 orig_inserted->getType()) +
2197 7) /
2198 8;
2199
2200 if (!gutils->isConstantValue(orig_inserted)) {
2201 const auto &TT = TR.query(orig_inserted);
2202
2203 unsigned start = 0;
2204 Value *dindex = nullptr;
2205
2206 auto &dl = gutils->newFunc->getParent()->getDataLayout();
2207 while (1) {
2208 unsigned nextStart = size0;
2209
2210 auto dt = TT[{-1}];
2211 for (size_t i = start; i < size0;) {
2212 auto nex = TT[{(int)i}];
2213 if ((nex == BaseType::Anything && dt.isFloat()) ||
2214 (dt == BaseType::Anything && nex.isFloat())) {
2215 nextStart = i;
2216 break;
2217 }
2218 bool Legal = true;
2219 dt.checkedOrIn(nex, /*PointerIntSame*/ true, Legal);
2220 if (!Legal) {
2221 nextStart = i;
2222 break;
2223 }
2224 if (auto fltType = dt.isFloat()) {
2225 i += dl.getTypeSizeInBits(fltType) / 8;
2226 } else if (dt == BaseType::Pointer) {
2227 i += dl.getPointerSizeInBits() / 8;
2228 } else {
2229 i++;
2230 }
2231 }
2232 Type *flt = dt.isFloat();
2233 if (!dt.isKnown()) {
2234 bool found = false;
2235 if (looseTypeAnalysis) {
2236 if (orig_inserted->getType()->isFPOrFPVectorTy()) {
2237 flt = orig_inserted->getType()->getScalarType();
2238 found = true;
2239 } else if (orig_inserted->getType()->isIntOrIntVectorTy() ||
2240 orig_inserted->getType()->isPointerTy()) {
2241 flt = nullptr;
2242 found = true;
2243 }
2244 }
2245 if (!found) {
2246 std::string str;
2247 raw_string_ostream ss(str);
2248 ss << "Cannot deduce type of insertvalue ins " << IVI
2249 << " size: " << size0 << " TT: " << TT.str();
2250 EmitNoTypeError(str, IVI, gutils, Builder2);
2251 }
2252 }
2253
2254 if (flt) {
2255 if (!dindex) {
2256 auto rule = [&](Value *prediff) {
2257 return Builder2.CreateExtractValue(prediff, IVI.getIndices());
2258 };
2259 auto prediff = diffe(&IVI, Builder2);
2260 dindex = applyChainRule(orig_inserted->getType(), Builder2, rule,
2261 prediff);
2262 }
2263
2264 ((DiffeGradientUtils *)gutils)
2265 ->addToDiffe(orig_inserted, dindex, Builder2, flt, start,
2266 nextStart - start);
2267 }
2268 if (nextStart == size0)
2269 break;
2270 start = nextStart;
2271 }
2272 }
2273
2274 size_t size1 = 1;
2275 if (orig_agg->getType()->isSized())
2276 size1 =
2277 (gutils->newFunc->getParent()->getDataLayout().getTypeSizeInBits(
2278 orig_agg->getType()) +
2279 7) /
2280 8;
2281
2282 if (!gutils->isConstantValue(orig_agg)) {
2283
2284 const auto &TT = TR.query(orig_agg);
2285 const MDNode *MD = hasMetadata(&IVI, "enzyme_truetype");
2286
2287 unsigned start = 0;
2288
2289 Value *dindex = nullptr;
2290
2291 auto &dl = gutils->newFunc->getParent()->getDataLayout();
2292 while (1) {
2293 unsigned nextStart = size1;
2294
2295 auto dt = TT[{-1}];
2296 for (size_t i = start; i < size1;) {
2297 auto nex = TT[{(int)i}];
2298 if (MD) {
2299 for (size_t j = 0; j < MD->getNumOperands(); j += 2) {
2300 ConcreteType base(
2301 llvm::cast<llvm::MDString>(MD->getOperand(j))->getString(),
2302 MD->getContext());
2303 auto offset = llvm::cast<llvm::ConstantInt>(
2304 llvm::cast<llvm::ConstantAsMetadata>(
2305 MD->getOperand(j + 1))
2306 ->getValue())
2307 ->getSExtValue();
2308 if (offset == (int64_t)i) {
2309 nex = base;
2310 break;
2311 }
2312 }
2313 }
2314 if ((nex == BaseType::Anything && dt.isFloat()) ||
2315 (dt == BaseType::Anything && nex.isFloat())) {
2316 nextStart = i;
2317 break;
2318 }
2319 bool Legal = true;
2320 dt.checkedOrIn(nex, /*PointerIntSame*/ true, Legal);
2321 if (!Legal) {
2322 nextStart = i;
2323 break;
2324 }
2325 if (auto fltType = dt.isFloat()) {
2326 i += dl.getTypeSizeInBits(fltType) / 8;
2327 } else if (dt == BaseType::Pointer) {
2328 i += dl.getPointerSizeInBits() / 8;
2329 } else {
2330 i++;
2331 }
2332 }
2333 Type *flt = dt.isFloat();
2334 if (!dt.isKnown()) {
2335 bool found = false;
2336 if (looseTypeAnalysis) {
2337 if (orig_agg->getType()->isFPOrFPVectorTy()) {
2338 flt = orig_agg->getType()->getScalarType();
2339 found = true;
2340 } else if (orig_agg->getType()->isIntOrIntVectorTy() ||
2341 orig_agg->getType()->isPointerTy()) {
2342 flt = nullptr;
2343 found = true;
2344 }
2345 }
2346 if (!found) {
2347 std::string str;
2348 raw_string_ostream ss(str);
2349 ss << "Cannot deduce type of insertvalue agg " << IVI
2350 << " start: " << start << " size: " << size1
2351 << " TT: " << TT.str();
2352 EmitNoTypeError(str, IVI, gutils, Builder2);
2353 }
2354 }
2355
2356 if (flt) {
2357 if (!dindex) {
2358 auto rule = [&](Value *prediff) {
2359 return Builder2.CreateInsertValue(
2360 prediff, Constant::getNullValue(orig_inserted->getType()),
2361 IVI.getIndices());
2362 };
2363 auto prediff = diffe(&IVI, Builder2);
2364 dindex =
2365 applyChainRule(orig_agg->getType(), Builder2, rule, prediff);
2366 }
2367 ((DiffeGradientUtils *)gutils)
2368 ->addToDiffe(orig_agg, dindex, Builder2, flt, start,
2369 nextStart - start);
2370 }
2371 if (nextStart == size1)
2372 break;
2373 start = nextStart;
2374 }
2375 }
2376
2377 setDiffe(&IVI,
2378 Constant::getNullValue(gutils->getShadowType(IVI.getType())),
2379 Builder2);
2380 return;
2381 }
2383 return;
2384 }
2385 }
2386 }
2387
2388 void getReverseBuilder(llvm::IRBuilder<> &Builder2, bool original = true) {
2389 ((GradientUtils *)gutils)->getReverseBuilder(Builder2, original);
2390 }
2391
2392 void getForwardBuilder(llvm::IRBuilder<> &Builder2) {
2393 ((GradientUtils *)gutils)->getForwardBuilder(Builder2);
2394 }
2395
2396 llvm::Value *diffe(llvm::Value *val, llvm::IRBuilder<> &Builder) {
2397 assert(Mode != DerivativeMode::ReverseModePrimal);
2398 return ((DiffeGradientUtils *)gutils)->diffe(val, Builder);
2399 }
2400
2401 void setDiffe(llvm::Value *val, llvm::Value *dif,
2402 llvm::IRBuilder<> &Builder) {
2403 assert(Mode != DerivativeMode::ReverseModePrimal);
2404 ((DiffeGradientUtils *)gutils)->setDiffe(val, dif, Builder);
2405 }
2406
2407 /// Unwraps a vector derivative from its internal representation and applies a
2408 /// function f to each element. Return values of f are collected and wrapped.
2409 template <typename Func, typename... Args>
2410 llvm::Value *applyChainRule(llvm::Type *diffType, llvm::IRBuilder<> &Builder,
2411 Func rule, Args... args) {
2412 return ((GradientUtils *)gutils)
2413 ->applyChainRule(diffType, Builder, rule, args...);
2414 }
2415
2416 /// Unwraps a vector derivative from its internal representation and applies a
2417 /// function f to each element.
2418 template <typename Func, typename... Args>
2419 void applyChainRule(llvm::IRBuilder<> &Builder, Func rule, Args... args) {
2420 ((GradientUtils *)gutils)->applyChainRule(Builder, rule, args...);
2421 }
2422
2423 /// Unwraps an collection of constant vector derivatives from their internal
2424 /// representations and applies a function f to each element.
2425 template <typename Func>
2426 void applyChainRule(llvm::ArrayRef<llvm::Value *> diffs,
2427 llvm::IRBuilder<> &Builder, Func rule) {
2428 ((GradientUtils *)gutils)->applyChainRule(diffs, Builder, rule);
2429 }
2430
2431 bool shouldFree() {
2432 assert(Mode == DerivativeMode::ReverseModeCombined ||
2435 return ((DiffeGradientUtils *)gutils)->FreeMemory;
2436 }
2437
2438 llvm::SmallVector<llvm::SelectInst *, 4>
2439 addToDiffe(llvm::Value *val, llvm::Value *dif, llvm::IRBuilder<> &Builder,
2440 llvm::Type *T, llvm::Value *mask = nullptr) {
2441 return ((DiffeGradientUtils *)gutils)
2442 ->addToDiffe(val, dif, Builder, T, /*idxs*/ {}, mask);
2443 }
2444
2445 llvm::Value *lookup(llvm::Value *val, llvm::IRBuilder<> &Builder) {
2446 return gutils->lookupM(val, Builder);
2447 }
2448
2449 void visitBinaryOperator(llvm::BinaryOperator &BO) {
2450 eraseIfUnused(BO);
2451
2452 if (BO.getOpcode() == llvm::Instruction::FDiv &&
2455 !gutils->isConstantValue(&BO)) {
2456 using namespace llvm;
2457 // Required loopy phi = [in, BO, BO, ..., BO]
2458 // 1) phi is only used in this B0
2459 // 2) BO dominates all latches
2460 // 3) phi == B0 whenever not coming from preheader [implies 2]
2461 // 4) [optional but done for ease] one exit to make it easier to
2462 // calculation the product at that point
2463 Value *orig_op0 = BO.getOperand(0);
2464 if (auto P0 = dyn_cast<PHINode>(orig_op0)) {
2465 LoopContext lc;
2466 SmallVector<Instruction *, 4> activeUses;
2467 for (auto u : P0->users()) {
2468 if (!gutils->isConstantInstruction(cast<Instruction>(u))) {
2469 activeUses.push_back(cast<Instruction>(u));
2470 } else if (retType == DIFFE_TYPE::OUT_DIFF && isa<ReturnInst>(u)) {
2471 activeUses.push_back(cast<Instruction>(u));
2472 }
2473 }
2474 if (activeUses.size() == 1 && activeUses[0] == &BO &&
2475 gutils->getContext(gutils->getNewFromOriginal(P0->getParent()),
2476 lc) &&
2477 gutils->getNewFromOriginal(P0->getParent()) == lc.header) {
2478 SmallVector<BasicBlock *, 1> Latches;
2479 gutils->OrigLI->getLoopFor(P0->getParent())->getLoopLatches(Latches);
2480 bool allIncoming = true;
2481 for (auto Latch : Latches) {
2482 if (&BO != P0->getIncomingValueForBlock(Latch)) {
2483 allIncoming = false;
2484 break;
2485 }
2486 }
2487 if (allIncoming && lc.exitBlocks.size() == 1) {
2488
2489 IRBuilder<> Builder2(&BO);
2490 getReverseBuilder(Builder2);
2491
2492 Value *orig_op1 = BO.getOperand(1);
2493
2494 Value *dif1 = nullptr;
2495 Value *idiff = diffe(&BO, Builder2);
2496
2497 Type *addingType = BO.getType();
2498
2499 if (!gutils->isConstantValue(orig_op1)) {
2500 IRBuilder<> EB(*lc.exitBlocks.begin());
2501 getReverseBuilder(EB, /*original=*/false);
2502 Value *Pstart = P0->getIncomingValueForBlock(
2503 gutils->getOriginalFromNew(lc.preheader));
2504 if (gutils->isConstantValue(Pstart)) {
2505 Value *lop0 = lookup(gutils->getNewFromOriginal(&BO), EB);
2506 Value *lop1 =
2507 lookup(gutils->getNewFromOriginal(orig_op1), Builder2);
2508 auto rule = [&](Value *idiff) {
2509 auto res = Builder2.CreateFDiv(
2510 Builder2.CreateFNeg(Builder2.CreateFMul(idiff, lop0)),
2511 lop1);
2512 if (gutils->strongZero) {
2513 res = CreateSelect(
2514 Builder2,
2515 Builder2.CreateFCmpOEQ(
2516 idiff, Constant::getNullValue(idiff->getType())),
2517 idiff, res);
2518 }
2519 return res;
2520 };
2521 dif1 =
2522 applyChainRule(orig_op1->getType(), Builder2, rule, idiff);
2523 } else {
2524 auto product = gutils->getOrInsertTotalMultiplicativeProduct(
2525 gutils->getNewFromOriginal(orig_op1), lc);
2526 IRBuilder<> EB(*lc.exitBlocks.begin());
2527 getReverseBuilder(EB, /*original=*/false);
2528 Value *s = lookup(gutils->getNewFromOriginal(Pstart), Builder2);
2529 Value *lop0 = lookup(product, EB);
2530 Value *lop1 =
2531 lookup(gutils->getNewFromOriginal(orig_op1), Builder2);
2532 auto rule = [&](Value *idiff) {
2533 auto res = Builder2.CreateFDiv(
2534 Builder2.CreateFNeg(Builder2.CreateFMul(
2535 s, Builder2.CreateFDiv(idiff, lop0))),
2536 lop1);
2537 if (gutils->strongZero) {
2538 res = CreateSelect(
2539 Builder2,
2540 Builder2.CreateFCmpOEQ(
2541 idiff, Constant::getNullValue(idiff->getType())),
2542 idiff, res);
2543 }
2544 return res;
2545 };
2546 dif1 =
2547 applyChainRule(orig_op1->getType(), Builder2, rule, idiff);
2548 }
2549 addToDiffe(orig_op1, dif1, Builder2, addingType);
2550 }
2551 return;
2552 }
2553 }
2554 }
2555 }
2556
2557 {
2558 using namespace llvm;
2559 switch (BO.getOpcode()) {
2560#include "BinopDerivatives.inc"
2561 default:
2562 break;
2563 }
2564 }
2565
2566 switch (Mode) {
2569 if (gutils->isConstantInstruction(&BO))
2570 return;
2572 break;
2577 break;
2579 return;
2580 }
2581 }
2582
2583 void createBinaryOperatorAdjoint(llvm::BinaryOperator &BO) {
2584 if (gutils->isConstantInstruction(&BO)) {
2585 return;
2586 }
2587 using namespace llvm;
2588
2589 IRBuilder<> Builder2(&BO);
2590 getReverseBuilder(Builder2);
2591
2592 Value *orig_op0 = BO.getOperand(0);
2593 Value *orig_op1 = BO.getOperand(1);
2594
2595 Value *dif0 = nullptr;
2596 Value *dif1 = nullptr;
2597 Value *idiff = diffe(&BO, Builder2);
2598
2599 Type *addingType = BO.getType();
2600
2601 switch (BO.getOpcode()) {
2602 case Instruction::LShr: {
2603 if (!gutils->isConstantValue(orig_op0)) {
2604 if (auto ci = dyn_cast<ConstantInt>(orig_op1)) {
2605 size_t size = 1;
2606 if (orig_op0->getType()->isSized())
2607 size = (gutils->newFunc->getParent()
2608 ->getDataLayout()
2609 .getTypeSizeInBits(orig_op0->getType()) +
2610 7) /
2611 8;
2612
2613 if (Type *flt = TR.addingType(size, orig_op0)) {
2614 auto bits = gutils->newFunc->getParent()
2615 ->getDataLayout()
2616 .getTypeAllocSizeInBits(flt);
2617 if (ci->getSExtValue() >= (int64_t)bits &&
2618 ci->getSExtValue() % bits == 0) {
2619 auto rule = [&](Value *idiff) {
2620 return Builder2.CreateShl(idiff, ci);
2621 };
2622 dif0 = applyChainRule(orig_op0->getType(), Builder2, rule, idiff);
2623 addingType = flt;
2624 goto done;
2625 }
2626 }
2627 }
2628 }
2629 if (looseTypeAnalysis) {
2630 llvm::errs() << "warning: binary operator is integer and constant: "
2631 << BO << "\n";
2632 // if loose type analysis, assume this integer and is constant
2633 return;
2634 }
2635 goto def;
2636 }
2637 case Instruction::AShr: {
2638 if (looseTypeAnalysis) {
2639 llvm::errs() << "warning: binary operator is integer and constant: "
2640 << BO << "\n";
2641 // if loose type analysis, assume this integer and is constant
2642 return;
2643 }
2644 goto def;
2645 }
2646 case Instruction::And: {
2647 // If & against 0b10000000000 and a float the result is 0
2648 auto &dl = gutils->oldFunc->getParent()->getDataLayout();
2649
2650 auto FT = TR.query(&BO).allFloat(&BO, dl);
2651 auto eFT = FT;
2652 if (FT)
2653 for (int i = 0; i < 2; ++i) {
2654 auto CI = dyn_cast<ConstantInt>(BO.getOperand(i));
2655 if (CI && dl.getTypeSizeInBits(eFT) ==
2656 dl.getTypeSizeInBits(CI->getType())) {
2657 if (eFT->isDoubleTy() && CI->getValue() == -134217728) {
2658 setDiffe(
2659 &BO,
2660 Constant::getNullValue(gutils->getShadowType(BO.getType())),
2661 Builder2);
2662 // Derivative is zero (equivalent to rounding as just chopping off
2663 // bits of mantissa), no update
2664 return;
2665 }
2666 }
2667 }
2668 if (looseTypeAnalysis) {
2669 llvm::errs() << "warning: binary operator is integer and constant: "
2670 << BO << "\n";
2671 // if loose type analysis, assume this integer and is constant
2672 return;
2673 }
2674 goto def;
2675 }
2676 case Instruction::Xor: {
2677 auto &dl = gutils->oldFunc->getParent()->getDataLayout();
2678
2679 auto FT = TR.query(&BO).allFloat(&BO, dl);
2680 auto eFT = FT;
2681 // If ^ against 0b10000000000 and a float the result is a float
2682 if (FT)
2683 for (int i = 0; i < 2; ++i) {
2684 if (containsOnlyAtMostTopBit(BO.getOperand(i), eFT, dl, &FT)) {
2685 setDiffe(
2686 &BO,
2687 Constant::getNullValue(gutils->getShadowType(BO.getType())),
2688 Builder2);
2689 auto isZero = Builder2.CreateICmpEQ(
2690 lookup(gutils->getNewFromOriginal(BO.getOperand(i)), Builder2),
2691 Constant::getNullValue(BO.getType()));
2692 auto rule = [&](Value *idiff) {
2693 auto ext = Builder2.CreateBitCast(idiff, FT);
2694 auto neg = Builder2.CreateFNeg(ext);
2695 neg = CreateSelect(Builder2, isZero, ext, neg);
2696 neg = Builder2.CreateBitCast(neg, BO.getType());
2697 return neg;
2698 };
2699 auto bc = applyChainRule(BO.getOperand(1 - i)->getType(), Builder2,
2700 rule, idiff);
2701 addToDiffe(BO.getOperand(1 - i), bc, Builder2, FT);
2702 return;
2703 }
2704 }
2705 if (looseTypeAnalysis) {
2706 llvm::errs() << "warning: binary operator is integer and constant: "
2707 << BO << "\n";
2708 // if loose type analysis, assume this integer and is constant
2709 return;
2710 }
2711 goto def;
2712 }
2713 case Instruction::Or: {
2714 auto &dl = gutils->oldFunc->getParent()->getDataLayout();
2715
2716 auto FT = TR.query(&BO).allFloat(&BO, dl);
2717 auto eFT = FT;
2718 // If & against 0b10000000000 and a float the result is a float
2719 if (FT)
2720 for (int i = 0; i < 2; ++i) {
2721 auto CI = dyn_cast<ConstantInt>(BO.getOperand(i));
2722 if (auto CV = dyn_cast<ConstantVector>(BO.getOperand(i))) {
2723 CI = dyn_cast_or_null<ConstantInt>(CV->getSplatValue());
2724 FT = VectorType::get(FT, CV->getType()->getElementCount());
2725 }
2726 if (auto CV = dyn_cast<ConstantDataVector>(BO.getOperand(i))) {
2727 CI = dyn_cast_or_null<ConstantInt>(CV->getSplatValue());
2728 FT = VectorType::get(FT, CV->getType()->getElementCount());
2729 }
2730 if (CI && dl.getTypeSizeInBits(eFT) ==
2731 dl.getTypeSizeInBits(CI->getType())) {
2732 auto AP = CI->getValue();
2733 bool validXor = false;
2734#if LLVM_VERSION_MAJOR > 16
2735 if (AP.isZero())
2736#else
2737 if (AP.isNullValue())
2738#endif
2739 {
2740 validXor = true;
2741 } else if (
2742 !AP.isNegative() &&
2743 ((eFT->isFloatTy()
2744#if LLVM_VERSION_MAJOR > 16
2745 && (AP & ~0b01111111100000000000000000000000ULL).isZero()
2746#else
2747 && (AP & ~0b01111111100000000000000000000000ULL).isNullValue()
2748#endif
2749 ) ||
2750 (eFT->isDoubleTy()
2751#if LLVM_VERSION_MAJOR > 16
2752 &&
2753 (AP &
2754 ~0b0111111111110000000000000000000000000000000000000000000000000000ULL)
2755 .isZero()
2756#else
2757 &&
2758 (AP &
2759 ~0b0111111111110000000000000000000000000000000000000000000000000000ULL)
2760 .isNullValue()
2761#endif
2762 ))) {
2763 validXor = true;
2764 }
2765 if (validXor) {
2766 setDiffe(
2767 &BO,
2768 Constant::getNullValue(gutils->getShadowType(BO.getType())),
2769 Builder2);
2770
2771 auto arg = lookup(
2772 gutils->getNewFromOriginal(BO.getOperand(1 - i)), Builder2);
2773
2774 auto rule = [&](Value *idiff) {
2775 auto prev = Builder2.CreateOr(arg, BO.getOperand(i));
2776 prev = Builder2.CreateSub(prev, arg, "", /*NUW*/ true,
2777 /*NSW*/ false);
2778 uint64_t num = 0;
2779 if (eFT->isFloatTy()) {
2780 num = 127ULL << 23;
2781 } else {
2782 assert(eFT->isDoubleTy());
2783 num = 1023ULL << 52;
2784 }
2785 prev = Builder2.CreateAdd(
2786 prev, ConstantInt::get(prev->getType(), num, false), "",
2787 /*NUW*/ true, /*NSW*/ true);
2788 auto prevFP = Builder2.CreateBitCast(prev, FT);
2789 auto idiffFP = Builder2.CreateBitCast(idiff, FT);
2790 prev = Builder2.CreateBitCast(
2791 checkedMul(gutils->strongZero, Builder2, idiffFP, prevFP),
2792 prev->getType());
2793 return prev;
2794 };
2795
2796 Value *prev = applyChainRule(BO.getOperand(1 - i)->getType(),
2797 Builder2, rule, idiff);
2798 addToDiffe(BO.getOperand(1 - i), prev, Builder2, FT);
2799 return;
2800 }
2801 }
2802 }
2803 if (looseTypeAnalysis) {
2804 llvm::errs() << "warning: binary operator is integer and constant: "
2805 << BO << "\n";
2806 // if loose type analysis, assume this integer or is constant
2807 return;
2808 }
2809 goto def;
2810 }
2811 case Instruction::UDiv:
2812 case Instruction::URem:
2813 case Instruction::SRem:
2814 case Instruction::SDiv:
2815 case Instruction::Shl:
2816 case Instruction::Mul:
2817 case Instruction::Sub:
2818 case Instruction::Add: {
2819 if (looseTypeAnalysis) {
2820 llvm::errs()
2821 << "warning: binary operator is integer and assumed constant: "
2822 << BO << "\n";
2823 // if loose type analysis, assume this integer add is constant
2824 return;
2825 }
2826 goto def;
2827 }
2828 default:
2829 def:;
2830 std::string s;
2831 llvm::raw_string_ostream ss(s);
2832 ss << *gutils->oldFunc << "\n";
2833 for (auto &arg : gutils->oldFunc->args()) {
2834 ss << " constantarg[" << arg << "] = " << gutils->isConstantValue(&arg)
2835 << " type: " << TR.query(&arg).str() << " - vals: {";
2836 for (auto v : TR.knownIntegralValues(&arg))
2837 ss << v << ",";
2838 ss << "}\n";
2839 }
2840 for (auto &BB : *gutils->oldFunc)
2841 for (auto &I : BB) {
2842 ss << " constantinst[" << I
2843 << "] = " << gutils->isConstantInstruction(&I)
2844 << " val:" << gutils->isConstantValue(&I)
2845 << " type: " << TR.query(&I).str() << "\n";
2846 }
2847 ss << "cannot handle unknown binary operator: " << BO << "\n";
2848 EmitNoDerivativeError(ss.str(), BO, gutils, Builder2);
2849 }
2850
2851 done:;
2852 if (dif0 || dif1)
2853 setDiffe(&BO, Constant::getNullValue(gutils->getShadowType(BO.getType())),
2854 Builder2);
2855 if (dif0)
2856 addToDiffe(orig_op0, dif0, Builder2, addingType);
2857 if (dif1)
2858 addToDiffe(orig_op1, dif1, Builder2, addingType);
2859 }
2860
2861 void createBinaryOperatorDual(llvm::BinaryOperator &BO) {
2862 using namespace llvm;
2863
2864 if (gutils->isConstantInstruction(&BO)) {
2866 return;
2867 }
2868
2869 IRBuilder<> Builder2(&BO);
2870 getForwardBuilder(Builder2);
2871
2872 Value *orig_op0 = BO.getOperand(0);
2873 Value *orig_op1 = BO.getOperand(1);
2874
2875 bool constantval0 = gutils->isConstantValue(orig_op0);
2876 bool constantval1 = gutils->isConstantValue(orig_op1);
2877
2878 switch (BO.getOpcode()) {
2879 case Instruction::And: {
2880 // If & against 0b10000000000 and a float the result is 0
2881 auto &dl = gutils->oldFunc->getParent()->getDataLayout();
2882 Type *diffTy = gutils->getShadowType(BO.getType());
2883
2884 auto FT = TR.query(&BO).allFloat(&BO, dl);
2885 auto eFT = FT;
2886 if (FT)
2887 for (int i = 0; i < 2; ++i) {
2888 auto CI = dyn_cast<ConstantInt>(BO.getOperand(i));
2889 if (CI && dl.getTypeSizeInBits(eFT) ==
2890 dl.getTypeSizeInBits(CI->getType())) {
2891 if (eFT->isDoubleTy() && CI->getValue() == -134217728) {
2892 setDiffe(&BO, Constant::getNullValue(diffTy), Builder2);
2893 // Derivative is zero (equivalent to rounding as just chopping off
2894 // bits of mantissa), no update
2895 return;
2896 }
2897 }
2898 }
2899 if (looseTypeAnalysis) {
2901 llvm::errs() << "warning: binary operator is integer and constant: "
2902 << BO << "\n";
2903 // if loose type analysis, assume this integer and is constant
2904 return;
2905 }
2906 goto def;
2907 }
2908 case Instruction::Xor: {
2909 auto &dl = gutils->oldFunc->getParent()->getDataLayout();
2910
2911 auto FT = TR.query(&BO).allFloat(&BO, dl);
2912 auto eFT = FT;
2913
2914 Value *dif[2] = {constantval0 ? nullptr : diffe(orig_op0, Builder2),
2915 constantval1 ? nullptr : diffe(orig_op1, Builder2)};
2916
2917 for (int i = 0; i < 2; ++i) {
2918 if (containsOnlyAtMostTopBit(BO.getOperand(i), eFT, dl, &FT) &&
2919 dif[1 - i] && !dif[i]) {
2920 auto isZero = Builder2.CreateICmpEQ(
2921 gutils->getNewFromOriginal(BO.getOperand(i)),
2922 Constant::getNullValue(BO.getType()));
2923 auto rule = [&](Value *idiff) {
2924 auto ext = Builder2.CreateBitCast(idiff, FT);
2925 auto neg = Builder2.CreateFNeg(ext);
2926 neg = CreateSelect(Builder2, isZero, ext, neg);
2927 neg = Builder2.CreateBitCast(neg, BO.getType());
2928 return neg;
2929 };
2930 auto bc = applyChainRule(BO.getOperand(1 - i)->getType(), Builder2,
2931 rule, dif[1 - i]);
2932 setDiffe(&BO, bc, Builder2);
2933 return;
2934 }
2935 }
2936 if (looseTypeAnalysis) {
2938 llvm::errs() << "warning: binary operator is integer and constant: "
2939 << BO << "\n";
2940 // if loose type analysis, assume this integer and is constant
2941 return;
2942 }
2943 goto def;
2944 }
2945 case Instruction::Or: {
2946 auto &dl = gutils->oldFunc->getParent()->getDataLayout();
2947
2948 Value *dif[2] = {constantval0 ? nullptr : diffe(orig_op0, Builder2),
2949 constantval1 ? nullptr : diffe(orig_op1, Builder2)};
2950
2951 auto FT = TR.query(&BO).allFloat(&BO, dl);
2952 auto eFT = FT;
2953 // If & against 0b10000000000 and a float the result is a float
2954 if (FT)
2955 for (int i = 0; i < 2; ++i) {
2956 auto CI = dyn_cast<ConstantInt>(BO.getOperand(i));
2957 if (auto CV = dyn_cast<ConstantVector>(BO.getOperand(i))) {
2958 CI = dyn_cast_or_null<ConstantInt>(CV->getSplatValue());
2959 FT = VectorType::get(FT, CV->getType()->getElementCount());
2960 }
2961 if (auto CV = dyn_cast<ConstantDataVector>(BO.getOperand(i))) {
2962 CI = dyn_cast_or_null<ConstantInt>(CV->getSplatValue());
2963 FT = VectorType::get(FT, CV->getType()->getElementCount());
2964 }
2965 if (CI && dl.getTypeSizeInBits(eFT) ==
2966 dl.getTypeSizeInBits(CI->getType())) {
2967 auto AP = CI->getValue();
2968 bool validXor = false;
2969#if LLVM_VERSION_MAJOR > 16
2970 if (AP.isZero())
2971#else
2972 if (AP.isNullValue())
2973#endif
2974 {
2975 validXor = true;
2976 } else if (
2977 !AP.isNegative() &&
2978 ((eFT->isFloatTy()
2979#if LLVM_VERSION_MAJOR > 16
2980 && (AP & ~0b01111111100000000000000000000000ULL).isZero()
2981#else
2982 && (AP & ~0b01111111100000000000000000000000ULL).isNullValue()
2983#endif
2984 ) ||
2985 (eFT->isDoubleTy()
2986#if LLVM_VERSION_MAJOR > 16
2987 &&
2988 (AP &
2989 ~0b0111111111110000000000000000000000000000000000000000000000000000ULL)
2990 .isZero()
2991#else
2992 &&
2993 (AP &
2994 ~0b0111111111110000000000000000000000000000000000000000000000000000ULL)
2995 .isNullValue()
2996#endif
2997 ))) {
2998 validXor = true;
2999 }
3000 if (validXor) {
3001 auto rule = [&](Value *difi) {
3002 auto arg = gutils->getNewFromOriginal(BO.getOperand(1 - i));
3003 auto prev = Builder2.CreateOr(arg, BO.getOperand(i));
3004 prev = Builder2.CreateSub(prev, arg, "", /*NUW*/ true,
3005 /*NSW*/ false);
3006 uint64_t num = 0;
3007 if (eFT->isFloatTy()) {
3008 num = 127ULL << 23;
3009 } else {
3010 assert(eFT->isDoubleTy());
3011 num = 1023ULL << 52;
3012 }
3013 prev = Builder2.CreateAdd(
3014 prev, ConstantInt::get(prev->getType(), num, false), "",
3015 /*NUW*/ true, /*NSW*/ true);
3016 auto prevFP = Builder2.CreateBitCast(prev, FT);
3017 auto difiFP = Builder2.CreateBitCast(difi, FT);
3018 prev = Builder2.CreateBitCast(
3019 checkedMul(gutils->strongZero, Builder2, difiFP, prevFP),
3020 prev->getType());
3021
3022 return prev;
3023 };
3024
3025 auto diffe =
3026 applyChainRule(BO.getType(), Builder2, rule, dif[1 - i]);
3027 setDiffe(&BO, diffe, Builder2);
3028 return;
3029 }
3030 }
3031 }
3032 if (looseTypeAnalysis) {
3034 llvm::errs() << "warning: binary operator is integer and constant: "
3035 << BO << "\n";
3036 // if loose type analysis, assume this integer or is constant
3037 return;
3038 }
3039 goto def;
3040 }
3041 case Instruction::LShr: {
3042 if (!gutils->isConstantValue(orig_op0)) {
3043 if (auto ci = dyn_cast<ConstantInt>(orig_op1)) {
3044 size_t size = 1;
3045 if (orig_op0->getType()->isSized())
3046 size = (gutils->newFunc->getParent()
3047 ->getDataLayout()
3048 .getTypeSizeInBits(orig_op0->getType()) +
3049 7) /
3050 8;
3051
3052 if (Type *flt = TR.addingType(size, orig_op0)) {
3053 auto bits = gutils->newFunc->getParent()
3054 ->getDataLayout()
3055 .getTypeAllocSizeInBits(flt);
3056 if (ci->getSExtValue() >= (int64_t)bits &&
3057 ci->getSExtValue() % bits == 0) {
3058 auto rule = [&](Value *idiff) {
3059 return Builder2.CreateLShr(idiff, ci);
3060 };
3061 auto dif = applyChainRule(orig_op0->getType(), Builder2, rule,
3062 diffe(orig_op0, Builder2));
3063 setDiffe(&BO, dif, Builder2);
3064 return;
3065 }
3066 }
3067 }
3068 }
3069 if (looseTypeAnalysis) {
3071 llvm::errs() << "warning: binary operator is integer and constant: "
3072 << BO << "\n";
3073 // if loose type analysis, assume this integer or is constant
3074 return;
3075 }
3076 goto def;
3077 }
3078 case Instruction::AShr:
3079 case Instruction::SDiv:
3080 case Instruction::UDiv:
3081 case Instruction::SRem:
3082 case Instruction::URem:
3083 case Instruction::Shl:
3084 case Instruction::Mul:
3085 case Instruction::Sub:
3086 case Instruction::Add: {
3087 if (looseTypeAnalysis) {
3089 llvm::errs() << "warning: binary operator is integer and constant: "
3090 << BO << "\n";
3091 // if loose type analysis, assume this integer add is constant
3092 return;
3093 }
3094 goto def;
3095 }
3096 default:
3097 def:;
3098 std::string s;
3099 llvm::raw_string_ostream ss(s);
3100 ss << *gutils->oldFunc << "\n";
3101 for (auto &arg : gutils->oldFunc->args()) {
3102 ss << " constantarg[" << arg << "] = " << gutils->isConstantValue(&arg)
3103 << " type: " << TR.query(&arg).str() << " - vals: {";
3104 for (auto v : TR.knownIntegralValues(&arg))
3105 ss << v << ",";
3106 ss << "}\n";
3107 }
3108 for (auto &BB : *gutils->oldFunc)
3109 for (auto &I : BB) {
3110 ss << " constantinst[" << I
3111 << "] = " << gutils->isConstantInstruction(&I)
3112 << " val:" << gutils->isConstantValue(&I)
3113 << " type: " << TR.query(&I).str() << "\n";
3114 }
3115 ss << "cannot handle unknown binary operator: " << BO << "\n";
3116 auto rval = EmitNoDerivativeError(ss.str(), BO, gutils, Builder2);
3117 if (!rval)
3118 rval = Constant::getNullValue(gutils->getShadowType(BO.getType()));
3119 auto ifound = gutils->invertedPointers.find(&BO);
3120 if (!gutils->isConstantValue(&BO)) {
3121 if (ifound != gutils->invertedPointers.end()) {
3122 auto placeholder = cast<PHINode>(&*ifound->second);
3123 gutils->invertedPointers.erase(ifound);
3124 gutils->replaceAWithB(placeholder, rval);
3125 gutils->erase(placeholder);
3126 gutils->invertedPointers.insert(std::make_pair(
3127 (const Value *)&BO, InvertedPointerVH(gutils, rval)));
3128 }
3129 } else {
3130 assert(ifound == gutils->invertedPointers.end());
3131 }
3132 break;
3133 }
3134 }
3135
3136 void visitMemSetInst(llvm::MemSetInst &MS) { visitMemSetCommon(MS); }
3137
3138 void visitMemSetCommon(llvm::CallInst &MS) {
3139 using namespace llvm;
3140
3141 IRBuilder<> BuilderZ(&MS);
3142 getForwardBuilder(BuilderZ);
3143
3144 IRBuilder<> Builder2(&MS);
3147 getReverseBuilder(Builder2);
3148
3149 bool forceErase = false;
3151 for (const auto &pair : gutils->rematerializableAllocations) {
3152 if (gutils->allocationsToBeRematerialized.count(pair.first) &&
3153 pair.second.stores.count(&MS) && pair.second.LI) {
3154 forceErase = true;
3155 }
3156 }
3157 }
3158 if (forceErase)
3159 eraseIfUnused(MS, /*erase*/ true, /*check*/ false);
3160 else
3161 eraseIfUnused(MS);
3162
3163 Value *orig_op0 = MS.getArgOperand(0);
3164 Value *orig_op1 = MS.getArgOperand(1);
3165
3166 // If constant destination then no operation needs doing
3167 if (gutils->isConstantValue(orig_op0)) {
3168 return;
3169 }
3170
3171 bool activeValToSet = !gutils->isConstantValue(orig_op1);
3172 if (activeValToSet)
3173 if (auto CI = dyn_cast<ConstantInt>(orig_op1))
3174 if (CI->isZero())
3175 activeValToSet = false;
3176 if (activeValToSet) {
3177 std::string s;
3178 llvm::raw_string_ostream ss(s);
3179 ss << "couldn't handle non constant inst in memset to "
3180 "propagate differential to\n"
3181 << MS;
3182 EmitNoDerivativeError(ss.str(), MS, gutils, BuilderZ);
3183 }
3184
3185 if (Mode == DerivativeMode::ForwardMode ||
3187 Value *op0 = gutils->invertPointerM(orig_op0, BuilderZ);
3188 Value *op1 = gutils->getNewFromOriginal(MS.getArgOperand(1));
3189 Value *op2 = gutils->getNewFromOriginal(MS.getArgOperand(2));
3190 Value *op3 = nullptr;
3191 if (3 < MS.arg_size()) {
3192 op3 = gutils->getNewFromOriginal(MS.getOperand(3));
3193 }
3194
3195 auto Defs =
3196 gutils->getInvertedBundles(&MS,
3199 BuilderZ, /*lookup*/ false);
3200
3201 auto funcName = getFuncNameFromCall(&MS);
3203 BuilderZ,
3204 [&](Value *op0) {
3205 SmallVector<Value *, 4> args = {op0, op1, op2};
3206 if (op3)
3207 args.push_back(op3);
3208
3209 CallInst *cal;
3210 if (startsWith(funcName, "memset_pattern") ||
3211 startsWith(funcName, "llvm.experimental.memset"))
3212 cal = Builder2.CreateMemSet(
3213 op0, ConstantInt::get(Builder2.getInt8Ty(), 0), op2, {});
3214 else
3215 cal = BuilderZ.CreateCall(MS.getCalledFunction(), args, Defs);
3216
3217 llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
3218 ToCopy2.push_back(LLVMContext::MD_noalias);
3219 cal->copyMetadata(MS, ToCopy2);
3220 if (auto m = hasMetadata(&MS, "enzyme_zerostack"))
3221 cal->setMetadata("enzyme_zerostack", m);
3222 if (auto m = hasMetadata(&MS, "enzyme_truetype"))
3223 cal->setMetadata("enzyme_truetype", m);
3224
3225 if (startsWith(funcName, "memset_pattern") ||
3226 startsWith(funcName, "llvm.experimental.memset")) {
3227 AttributeList NewAttrs;
3228 for (auto idx :
3229 {AttributeList::ReturnIndex, AttributeList::FunctionIndex,
3230 AttributeList::FirstArgIndex})
3231 for (auto attr : MS.getAttributes().getAttributes(idx))
3232 NewAttrs =
3233 NewAttrs.addAttributeAtIndex(MS.getContext(), idx, attr);
3234 cal->setAttributes(NewAttrs);
3235 } else
3236 cal->setAttributes(MS.getAttributes());
3237 cal->setCallingConv(MS.getCallingConv());
3238 cal->setTailCallKind(MS.getTailCallKind());
3239 cal->setDebugLoc(gutils->getNewFromOriginal(MS.getDebugLoc()));
3240 },
3241 op0);
3242 return;
3243 }
3244
3245 bool forwardsShadow, backwardsShadow;
3246 shadowStoreSchedule(MS, MS.getOperand(0), forwardsShadow, backwardsShadow);
3247
3248 size_t size = 1;
3249 if (auto ci = dyn_cast<ConstantInt>(MS.getOperand(2))) {
3250 size = ci->getLimitedValue();
3251 }
3252
3253 // TODO note that we only handle memset of ONE type (aka memset of {int,
3254 // double} not allowed)
3255
3256 if (size == 0) {
3257 llvm::errs() << MS << "\n";
3258 }
3259 assert(size != 0);
3260
3261 // Offsets of the form Optional<floating type>, segment start, segment size
3262 std::vector<std::tuple<Type *, size_t, size_t>> toIterate;
3263
3264 // Special handling mechanism to bypass TA limitations by supporting
3265 // arbitrary sized types.
3266 if (auto MD = hasMetadata(&MS, "enzyme_truetype")) {
3267 toIterate = parseTrueType(MD, Mode, false);
3268 } else {
3269 auto &DL = gutils->newFunc->getParent()->getDataLayout();
3270 auto vd = TR.query(MS.getOperand(0)).Data0().ShiftIndices(DL, 0, size, 0);
3271
3272 if (!vd.isKnownPastPointer()) {
3273 // If unknown type results, and zeroing known undef allocation, consider
3274 // integers
3275 if (auto CI = dyn_cast<ConstantInt>(MS.getOperand(1)))
3276 if (CI->isZero()) {
3277 auto root = getBaseObject(MS.getOperand(0));
3278 bool writtenTo = false;
3279 bool undefMemory =
3280 isa<AllocaInst>(root) || isAllocationCall(root, gutils->TLI);
3281 if (auto arg = dyn_cast<Argument>(root))
3282 if (arg->hasStructRetAttr())
3283 undefMemory = true;
3284 if (undefMemory) {
3285 Instruction *cur = MS.getPrevNode();
3286 while (cur) {
3287 if (cur == root)
3288 break;
3289 if (auto MCI = dyn_cast<ConstantInt>(MS.getOperand(2))) {
3290 if (auto II = dyn_cast<IntrinsicInst>(cur)) {
3291 if (II->getCalledFunction()->getName() ==
3292 "llvm.enzyme.lifetime_start") {
3293 if (getBaseObject(II->getOperand(1)) == root) {
3294 if (auto CI2 =
3295 dyn_cast<ConstantInt>(II->getOperand(0))) {
3296 if (MCI->getValue().ule(CI2->getValue()))
3297 break;
3298 }
3299 }
3300 cur = cur->getPrevNode();
3301 continue;
3302 }
3303 // If the start of the lifetime for more memory than being
3304 // memset, its valid.
3305 if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
3306 if (getBaseObject(II->getOperand(1)) == root) {
3307 if (auto CI2 =
3308 dyn_cast<ConstantInt>(II->getOperand(0))) {
3309 if (MCI->getValue().ule(CI2->getValue()))
3310 break;
3311 }
3312 }
3313 cur = cur->getPrevNode();
3314 continue;
3315 }
3316 }
3317 }
3318 if (cur->mayWriteToMemory()) {
3319 writtenTo = true;
3320 break;
3321 }
3322 cur = cur->getPrevNode();
3323 // We've hit the start of the block, assume written to by a
3324 // previous block.
3325 if (cur == nullptr) {
3326 writtenTo = true;
3327 }
3328 }
3329
3330 if (!writtenTo) {
3332 vd.insert({-1}, BaseType::Integer);
3333 }
3334 }
3335 }
3336 }
3337
3338 if (!vd.isKnownPastPointer()) {
3339 // If unknown type results, consider the intersection of all incoming.
3340 if (isa<PHINode>(MS.getOperand(0)) ||
3341 isa<SelectInst>(MS.getOperand(0))) {
3342 SmallVector<Value *, 2> todo = {MS.getOperand(0)};
3343 bool set = false;
3344 SmallSet<Value *, 2> seen;
3345 TypeTree vd2;
3346 while (todo.size()) {
3347 Value *cur = todo.back();
3348 todo.pop_back();
3349 if (seen.count(cur))
3350 continue;
3351 seen.insert(cur);
3352 if (auto PN = dyn_cast<PHINode>(cur)) {
3353 for (size_t i = 0, end = PN->getNumIncomingValues(); i < end;
3354 i++) {
3355 todo.push_back(PN->getIncomingValue(i));
3356 }
3357 continue;
3358 }
3359 if (auto S = dyn_cast<SelectInst>(cur)) {
3360 todo.push_back(S->getTrueValue());
3361 todo.push_back(S->getFalseValue());
3362 continue;
3363 }
3364 if (auto CE = dyn_cast<ConstantExpr>(cur)) {
3365 if (CE->isCast()) {
3366 todo.push_back(CE->getOperand(0));
3367 continue;
3368 }
3369 }
3370 if (auto CI = dyn_cast<CastInst>(cur)) {
3371 todo.push_back(CI->getOperand(0));
3372 continue;
3373 }
3374 if (isa<ConstantPointerNull>(cur))
3375 continue;
3376 if (auto CI = dyn_cast<ConstantInt>(cur))
3377 if (CI->isZero())
3378 continue;
3379 auto curTT = TR.query(cur).Data0().ShiftIndices(DL, 0, size, 0);
3380 if (!set)
3381 vd2 = curTT;
3382 else
3383 vd2 &= curTT;
3384 set = true;
3385 }
3386 vd = vd2;
3387 }
3388 }
3389 if (!vd.isKnownPastPointer()) {
3390 if (looseTypeAnalysis) {
3391#if LLVM_VERSION_MAJOR < 17
3392 if (auto CI = dyn_cast<CastInst>(MS.getOperand(0))) {
3393 if (auto PT = dyn_cast<PointerType>(CI->getSrcTy())) {
3394 auto ET = PT->getPointerElementType();
3395 while (1) {
3396 if (auto ST = dyn_cast<StructType>(ET)) {
3397 if (ST->getNumElements()) {
3398 ET = ST->getElementType(0);
3399 continue;
3400 }
3401 }
3402 if (auto AT = dyn_cast<ArrayType>(ET)) {
3403 ET = AT->getElementType();
3404 continue;
3405 }
3406 break;
3407 }
3408 if (ET->isFPOrFPVectorTy()) {
3409 vd = TypeTree(ConcreteType(ET->getScalarType())).Only(0, &MS);
3410 goto known;
3411 }
3412 if (ET->isPointerTy()) {
3413 vd = TypeTree(BaseType::Pointer).Only(0, &MS);
3414 goto known;
3415 }
3416 if (ET->isIntOrIntVectorTy()) {
3417 vd = TypeTree(BaseType::Integer).Only(0, &MS);
3418 goto known;
3419 }
3420 }
3421 }
3422#endif
3423 if (auto gep = dyn_cast<GetElementPtrInst>(MS.getOperand(0))) {
3424 if (auto AT = dyn_cast<ArrayType>(gep->getSourceElementType())) {
3425 if (AT->getElementType()->isIntegerTy()) {
3426 vd = TypeTree(BaseType::Integer).Only(0, &MS);
3427 goto known;
3428 }
3429 }
3430 }
3431 EmitWarning("CannotDeduceType", MS,
3432 "failed to deduce type of memset ", MS);
3433 vd = TypeTree(BaseType::Pointer).Only(0, &MS);
3434 goto known;
3435 }
3436 std::string str;
3437 raw_string_ostream ss(str);
3438 ss << "Cannot deduce type of memset " << MS;
3439 EmitNoTypeError(str, MS, gutils, BuilderZ);
3440 return;
3441 }
3442 known:;
3443 {
3444 unsigned start = 0;
3445 while (1) {
3446 unsigned nextStart = size;
3447
3448 auto dt = vd[{-1}];
3449 for (size_t i = start; i < size; ++i) {
3450 bool Legal = true;
3451 dt.checkedOrIn(vd[{(int)i}], /*PointerIntSame*/ true, Legal);
3452 if (!Legal) {
3453 nextStart = i;
3454 break;
3455 }
3456 }
3457 if (!dt.isKnown()) {
3458 TR.dump();
3459 llvm::errs() << " vd:" << vd.str() << " start:" << start
3460 << " size: " << size << " dt:" << dt.str() << "\n";
3461 }
3462 assert(dt.isKnown());
3463 toIterate.emplace_back(dt.isFloat(), start, nextStart - start);
3464
3465 if (nextStart == size)
3466 break;
3467 start = nextStart;
3468 }
3469 }
3470 }
3471
3472#if 0
3473 unsigned dstalign = dstAlign.valueOrOne().value();
3474 unsigned srcalign = srcAlign.valueOrOne().value();
3475#endif
3476
3477 Value *op1 = gutils->getNewFromOriginal(MS.getArgOperand(1));
3478 Value *new_size = gutils->getNewFromOriginal(MS.getArgOperand(2));
3479 Value *op3 = nullptr;
3480 if (3 < MS.arg_size()) {
3481 op3 = gutils->getNewFromOriginal(MS.getOperand(3));
3482 }
3483
3484 for (auto &&[secretty_ref, seg_start_ref, seg_size_ref] : toIterate) {
3485 auto secretty = secretty_ref;
3486 auto seg_start = seg_start_ref;
3487 auto seg_size = seg_size_ref;
3488
3489 Value *length = new_size;
3490 if (seg_start != std::get<1>(toIterate.back())) {
3491 length = ConstantInt::get(new_size->getType(), seg_start + seg_size);
3492 }
3493 if (seg_start != 0)
3494 length = BuilderZ.CreateSub(
3495 length, ConstantInt::get(new_size->getType(), seg_start));
3496
3497#if 0
3498 unsigned subdstalign = dstalign;
3499 // todo make better alignment calculation
3500 if (dstalign != 0) {
3501 if (start % dstalign != 0) {
3502 dstalign = 1;
3503 }
3504 }
3505 unsigned subsrcalign = srcalign;
3506 // todo make better alignment calculation
3507 if (srcalign != 0) {
3508 if (start % srcalign != 0) {
3509 srcalign = 1;
3510 }
3511 }
3512#endif
3513
3514 Value *shadow_dst = gutils->invertPointerM(MS.getOperand(0), BuilderZ);
3515
3516 // TODO ponder forward split mode
3517 if (!secretty &&
3518 ((Mode == DerivativeMode::ReverseModePrimal && forwardsShadow) ||
3519 (Mode == DerivativeMode::ReverseModeCombined && forwardsShadow) ||
3520 (Mode == DerivativeMode::ReverseModeGradient && backwardsShadow) ||
3521 (Mode == DerivativeMode::ForwardModeSplit && backwardsShadow))) {
3522 auto Defs =
3523 gutils->getInvertedBundles(&MS,
3526 BuilderZ, /*lookup*/ false);
3527 auto rule = [&](Value *op0) {
3528 if (seg_start != 0) {
3529 Value *idxs[] = {ConstantInt::get(
3530 Type::getInt32Ty(op0->getContext()), seg_start)};
3531 op0 = BuilderZ.CreateInBoundsGEP(Type::getInt8Ty(op0->getContext()),
3532 op0, idxs);
3533 }
3534 SmallVector<Value *, 4> args = {op0, op1, length};
3535 if (op3)
3536 args.push_back(op3);
3537 auto cal = BuilderZ.CreateCall(MS.getCalledFunction(), args, Defs);
3538 llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
3539 ToCopy2.push_back(LLVMContext::MD_noalias);
3540 if (auto m = hasMetadata(&MS, "enzyme_zerostack"))
3541 cal->setMetadata("enzyme_zerostack", m);
3542 if (auto m = hasMetadata(&MS, "enzyme_truetype"))
3543 if (auto sliced = sliceTrueType(m, seg_start, seg_size))
3544 cal->setMetadata("enzyme_truetype", sliced);
3545 cal->copyMetadata(MS, ToCopy2);
3546 cal->setAttributes(MS.getAttributes());
3547 cal->setCallingConv(MS.getCallingConv());
3548 cal->setTailCallKind(MS.getTailCallKind());
3549 cal->setDebugLoc(gutils->getNewFromOriginal(MS.getDebugLoc()));
3550 };
3551
3552 applyChainRule(BuilderZ, rule, shadow_dst);
3553 }
3554 if (secretty && (Mode == DerivativeMode::ReverseModeGradient ||
3556
3557 auto Defs =
3558 gutils->getInvertedBundles(&MS,
3561 BuilderZ, /*lookup*/ true);
3562 Value *op1l = gutils->lookupM(op1, Builder2);
3563 Value *op3l = op3;
3564 if (op3l)
3565 op3l = gutils->lookupM(op3l, BuilderZ);
3566 length = gutils->lookupM(length, Builder2);
3567 auto rule = [&](Value *op0) {
3568 if (seg_start != 0) {
3569 Value *idxs[] = {ConstantInt::get(
3570 Type::getInt32Ty(op0->getContext()), seg_start)};
3571 op0 = Builder2.CreateInBoundsGEP(Type::getInt8Ty(op0->getContext()),
3572 op0, idxs);
3573 }
3574 SmallVector<Value *, 4> args = {op0, op1l, length};
3575 if (op3l)
3576 args.push_back(op3l);
3577 CallInst *cal;
3578 auto funcName = getFuncNameFromCall(&MS);
3579 if (startsWith(funcName, "memset_pattern") ||
3580 startsWith(funcName, "llvm.experimental.memset"))
3581 cal = Builder2.CreateMemSet(
3582 op0, ConstantInt::get(Builder2.getInt8Ty(), 0), length, {});
3583 else
3584 cal = Builder2.CreateCall(MS.getCalledFunction(), args, Defs);
3585 llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
3586 ToCopy2.push_back(LLVMContext::MD_noalias);
3587 cal->copyMetadata(MS, ToCopy2);
3588 if (auto m = hasMetadata(&MS, "enzyme_zerostack"))
3589 cal->setMetadata("enzyme_zerostack", m);
3590 if (auto m = hasMetadata(&MS, "enzyme_truetype"))
3591 if (auto sliced = sliceTrueType(m, seg_start, seg_size))
3592 cal->setMetadata("enzyme_truetype", sliced);
3593
3594 if (startsWith(funcName, "memset_pattern") ||
3595 startsWith(funcName, "llvm.experimental.memset")) {
3596 AttributeList NewAttrs;
3597 for (auto idx :
3598 {AttributeList::ReturnIndex, AttributeList::FunctionIndex,
3599 AttributeList::FirstArgIndex})
3600 for (auto attr : MS.getAttributes().getAttributes(idx))
3601 NewAttrs =
3602 NewAttrs.addAttributeAtIndex(MS.getContext(), idx, attr);
3603 cal->setAttributes(NewAttrs);
3604 } else
3605 cal->setAttributes(MS.getAttributes());
3606 cal->setCallingConv(MS.getCallingConv());
3607 cal->setDebugLoc(gutils->getNewFromOriginal(MS.getDebugLoc()));
3608 };
3609
3610 applyChainRule(Builder2, rule, gutils->lookupM(shadow_dst, Builder2));
3611 }
3612 }
3613 }
3614
3615 void visitMemTransferInst(llvm::MemTransferInst &MTI) {
3616 using namespace llvm;
3617 Value *isVolatile = gutils->getNewFromOriginal(MTI.getOperand(3));
3618 auto srcAlign = MTI.getSourceAlign();
3619 auto dstAlign = MTI.getDestAlign();
3620 visitMemTransferCommon(MTI.getIntrinsicID(), srcAlign, dstAlign, MTI,
3621 MTI.getOperand(0), MTI.getOperand(1),
3622 gutils->getNewFromOriginal(MTI.getOperand(2)),
3623 isVolatile);
3624 }
3625
3626 void visitMemTransferCommon(llvm::Intrinsic::ID ID, llvm::MaybeAlign srcAlign,
3627 llvm::MaybeAlign dstAlign, llvm::CallInst &MTI,
3628 llvm::Value *orig_dst, llvm::Value *orig_src,
3629 llvm::Value *new_size, llvm::Value *isVolatile) {
3630 using namespace llvm;
3631
3632 if (gutils->isConstantValue(MTI.getOperand(0))) {
3633 eraseIfUnused(MTI);
3634 return;
3635 }
3636
3637 if (unnecessaryStores.count(&MTI)) {
3638 eraseIfUnused(MTI);
3639 return;
3640 }
3641
3642 // memcpy of size 1 cannot move differentiable data [single byte copy]
3643 if (auto ci = dyn_cast<ConstantInt>(new_size)) {
3644 if (ci->getValue() == 1) {
3645 eraseIfUnused(MTI);
3646 return;
3647 }
3648 }
3649
3650 // copying into nullptr is invalid (not sure why it exists here), but we
3651 // shouldn't do it in reverse pass or shadow
3652 if (isa<ConstantPointerNull>(orig_dst) ||
3653 TR.query(orig_dst)[{-1}] == BaseType::Anything) {
3654 eraseIfUnused(MTI);
3655 return;
3656 }
3657
3658 size_t size = 1;
3659 if (auto ci = dyn_cast<ConstantInt>(new_size)) {
3660 size = ci->getLimitedValue();
3661 }
3662
3663 // TODO note that we only handle memcpy/etc of ONE type (aka memcpy of {int,
3664 // double} not allowed)
3665 if (size == 0) {
3666 eraseIfUnused(MTI);
3667 return;
3668 }
3669
3670 if ((Mode == DerivativeMode::ForwardMode ||
3672 gutils->isConstantValue(orig_dst)) {
3673 eraseIfUnused(MTI);
3674 return;
3675 }
3676
3677 // Offsets of the form Optional<floating type>, segment start, segment size
3678 std::vector<std::tuple<Type *, size_t, size_t>> toIterate;
3679 IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&MTI));
3680
3681 // Special handling mechanism to bypass TA limitations by supporting
3682 // arbitrary sized types.
3683 if (auto MD = hasMetadata(&MTI, "enzyme_truetype")) {
3684 toIterate = parseTrueType(MD, Mode,
3685 !gutils->isConstantValue(orig_src) &&
3686 !gutils->runtimeActivity);
3687 } else {
3688 auto &DL = gutils->newFunc->getParent()->getDataLayout();
3689 auto vd = TR.query(orig_dst).Data0().ShiftIndices(DL, 0, size, 0);
3690 vd |= TR.query(orig_src).Data0().PurgeAnything().ShiftIndices(DL, 0, size,
3691 0);
3692 for (size_t i = 0; i < MTI.getNumOperands(); i++)
3693 if (MTI.getOperand(i) == orig_dst)
3694 if (MTI.getAttributes().hasParamAttr(i, "enzyme_type")) {
3695 auto attr = MTI.getAttributes().getParamAttr(i, "enzyme_type");
3696 auto TT =
3697 TypeTree::parse(attr.getValueAsString(), MTI.getContext());
3698 vd |= TT.Data0().ShiftIndices(DL, 0, size, 0);
3699 break;
3700 }
3701
3702 bool errorIfNoType = true;
3703 if ((Mode == DerivativeMode::ForwardMode ||
3705 (!gutils->isConstantValue(orig_src) && !gutils->runtimeActivity)) {
3706 errorIfNoType = false;
3707 }
3708
3709 if (!vd.isKnownPastPointer()) {
3710 if (looseTypeAnalysis) {
3711 for (auto val : {orig_dst, orig_src}) {
3712#if LLVM_VERSION_MAJOR < 17
3713 if (auto CI = dyn_cast<CastInst>(val)) {
3714 if (auto PT = dyn_cast<PointerType>(CI->getSrcTy())) {
3715 auto ET = PT->getPointerElementType();
3716 while (1) {
3717 if (auto ST = dyn_cast<StructType>(ET)) {
3718 if (ST->getNumElements()) {
3719 ET = ST->getElementType(0);
3720 continue;
3721 }
3722 }
3723 if (auto AT = dyn_cast<ArrayType>(ET)) {
3724 ET = AT->getElementType();
3725 continue;
3726 }
3727 break;
3728 }
3729 if (ET->isFPOrFPVectorTy()) {
3730 vd =
3731 TypeTree(ConcreteType(ET->getScalarType())).Only(0, &MTI);
3732 goto known;
3733 }
3734 if (ET->isPointerTy()) {
3735 vd = TypeTree(BaseType::Pointer).Only(0, &MTI);
3736 goto known;
3737 }
3738 if (ET->isIntOrIntVectorTy()) {
3739 vd = TypeTree(BaseType::Integer).Only(0, &MTI);
3740 goto known;
3741 }
3742 }
3743 }
3744#endif
3745 if (auto gep = dyn_cast<GetElementPtrInst>(val)) {
3746 if (auto AT = dyn_cast<ArrayType>(gep->getSourceElementType())) {
3747 if (AT->getElementType()->isIntegerTy()) {
3748 vd = TypeTree(BaseType::Integer).Only(0, &MTI);
3749 goto known;
3750 }
3751 }
3752 }
3753 }
3754 // If the type is known, but outside of the known range
3755 // (but the memcpy size is a variable), attempt to use
3756 // the first type out of range as the memcpy type.
3757 if (size == 1 && !isa<ConstantInt>(new_size)) {
3758 for (auto ptr : {orig_dst, orig_src}) {
3759 vd = TR.query(ptr).Data0().ShiftIndices(DL, 0, -1, 0);
3760 if (vd.isKnownPastPointer()) {
3762 size_t minInt = 0xFFFFFFFF;
3763 for (const auto &pair : vd.getMapping()) {
3764 if (pair.first.size() != 1)
3765 continue;
3766 if (minInt < (size_t)pair.first[0])
3767 continue;
3768 minInt = pair.first[0];
3769 mv = pair.second;
3770 }
3771 assert(mv != BaseType::Unknown);
3772 vd.insert({0}, mv);
3773 goto known;
3774 }
3775 }
3776 }
3777 if (errorIfNoType)
3778 EmitWarning("CannotDeduceType", MTI,
3779 "failed to deduce type of copy ", MTI);
3780 vd = TypeTree(BaseType::Pointer).Only(0, &MTI);
3781 goto known;
3782 }
3783 if (errorIfNoType) {
3784 std::string str;
3785 raw_string_ostream ss(str);
3786 ss << "Cannot deduce type of copy " << MTI;
3787 EmitNoTypeError(str, MTI, gutils, BuilderZ);
3788 vd = TypeTree(BaseType::Integer).Only(0, &MTI);
3789 } else {
3790 vd = TypeTree(BaseType::Pointer).Only(0, &MTI);
3791 }
3792 }
3793
3794 known:;
3795 {
3796
3797 unsigned start = 0;
3798 while (1) {
3799 unsigned nextStart = size;
3800
3801 auto dt = vd[{-1}];
3802 for (size_t i = start; i < size; ++i) {
3803 bool Legal = true;
3804 auto tmp = dt;
3805 auto next = vd[{(int)i}];
3806 tmp.checkedOrIn(next, /*PointerIntSame*/ true, Legal);
3807 // Prevent fusion of {Anything, Float} since anything is an int rule
3808 // but float requires zeroing.
3809 if ((dt == BaseType::Anything &&
3810 (next != BaseType::Anything && next.isKnown())) ||
3811 (next == BaseType::Anything &&
3812 (dt != BaseType::Anything && dt.isKnown())))
3813 Legal = false;
3814 if (!Legal) {
3815 if (Mode == DerivativeMode::ForwardMode ||
3817 // if both are floats (of any type), forward mode is the same.
3818 // + [potentially zero if const, otherwise copy]
3819 // if both are int/pointer (of any type), also the same
3820 // + copy
3821 // if known non-constant, also the same
3822 // + copy
3823 if ((dt.isFloat() == nullptr) ==
3824 (vd[{(int)i}].isFloat() == nullptr)) {
3825 Legal = true;
3826 }
3827 if (!gutils->isConstantValue(orig_src) &&
3828 !gutils->runtimeActivity) {
3829 Legal = true;
3830 }
3831 }
3832 if (!Legal) {
3833 nextStart = i;
3834 break;
3835 }
3836 } else
3837 dt = tmp;
3838 }
3839 if (!dt.isKnown()) {
3840 TR.dump();
3841 llvm::errs() << " vd:" << vd.str() << " start:" << start
3842 << " size: " << size << " dt:" << dt.str() << "\n";
3843 }
3844 assert(dt.isKnown());
3845 toIterate.emplace_back(dt.isFloat(), start, nextStart - start);
3846
3847 if (nextStart == size)
3848 break;
3849 start = nextStart;
3850 }
3851 }
3852 }
3853
3854 // llvm::errs() << "MIT: " << MTI << "|size: " << size << " vd: " <<
3855 // vd.str() << "\n";
3856
3857 unsigned dstalign = dstAlign.valueOrOne().value();
3858 unsigned srcalign = srcAlign.valueOrOne().value();
3859
3860 bool forwardsShadow, backwardsShadow;
3861 shadowStoreSchedule(MTI, orig_dst, forwardsShadow, backwardsShadow);
3862
3863 for (auto &&[floatTy_ref, seg_start_ref, seg_size_ref] : toIterate) {
3864 auto floatTy = floatTy_ref;
3865 auto seg_start = seg_start_ref;
3866 auto seg_size = seg_size_ref;
3867
3868 Value *length = new_size;
3869 if (seg_start != std::get<1>(toIterate.back())) {
3870 length = ConstantInt::get(new_size->getType(), seg_start + seg_size);
3871 }
3872 if (seg_start != 0)
3873 length = BuilderZ.CreateSub(
3874 length, ConstantInt::get(new_size->getType(), seg_start));
3875
3876 unsigned subdstalign = dstalign;
3877 // todo make better alignment calculation
3878 if (dstalign != 0) {
3879 if (seg_start % dstalign != 0) {
3880 dstalign = 1;
3881 }
3882 }
3883 unsigned subsrcalign = srcalign;
3884 // todo make better alignment calculation
3885 if (srcalign != 0) {
3886 if (seg_start % srcalign != 0) {
3887 srcalign = 1;
3888 }
3889 }
3890 IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&MTI));
3891 Value *shadow_dst = gutils->isConstantValue(orig_dst)
3892 ? nullptr
3893 : gutils->invertPointerM(orig_dst, BuilderZ);
3894 Value *shadow_src = gutils->isConstantValue(orig_src)
3895 ? nullptr
3896 : gutils->invertPointerM(orig_src, BuilderZ);
3897
3898 auto rev_rule = [&](Value *shadow_dst, Value *shadow_src) {
3899 if (shadow_dst == nullptr)
3900 shadow_dst = gutils->getNewFromOriginal(orig_dst);
3901 if (shadow_src == nullptr)
3902 shadow_src = gutils->getNewFromOriginal(orig_src);
3904 gutils, Mode, floatTy, ID, subdstalign, subsrcalign,
3905 /*offset*/ seg_start, gutils->isConstantValue(orig_dst), shadow_dst,
3906 gutils->getNewFromOriginal(orig_dst),
3907 gutils->isConstantValue(orig_src), shadow_src,
3908 gutils->getNewFromOriginal(orig_src),
3909 /*length*/ length, /*volatile*/ isVolatile, &MTI,
3910 /*allowForward*/ forwardsShadow, /*shadowsLookedup*/ false,
3911 /*backwardsShadow*/ backwardsShadow);
3912 };
3913
3914 auto fwd_rule = [&](Value *ddst, Value *dsrc) {
3915 if (ddst == nullptr)
3916 ddst = gutils->getNewFromOriginal(orig_dst);
3917 if (dsrc == nullptr)
3918 dsrc = gutils->getNewFromOriginal(orig_src);
3919 MaybeAlign dalign;
3920 if (subdstalign)
3921 dalign = MaybeAlign(subdstalign);
3922 MaybeAlign salign;
3923 if (subsrcalign)
3924 salign = MaybeAlign(subsrcalign);
3925 if (ddst->getType()->isIntegerTy())
3926 ddst =
3927 BuilderZ.CreateIntToPtr(ddst, getInt8PtrTy(ddst->getContext()));
3928 if (seg_start != 0) {
3929 ddst = BuilderZ.CreateConstInBoundsGEP1_64(
3930 Type::getInt8Ty(ddst->getContext()), ddst, seg_start);
3931 }
3932 CallInst *call;
3933 if (floatTy && gutils->isConstantValue(orig_src)) {
3934 call = BuilderZ.CreateMemSet(
3935 ddst, ConstantInt::get(Type::getInt8Ty(ddst->getContext()), 0),
3936 length, dalign, cast<ConstantInt>(isVolatile)->isOne());
3937 } else {
3938 auto toPtr = [&](Value *ptr) {
3939 if (ptr->getType()->isIntegerTy())
3940 ptr =
3941 BuilderZ.CreateIntToPtr(ptr, getInt8PtrTy(ptr->getContext()));
3942 if (seg_start != 0) {
3943 ptr = BuilderZ.CreateConstInBoundsGEP1_64(
3944 Type::getInt8Ty(ptr->getContext()), ptr, seg_start);
3945 }
3946 return ptr;
3947 };
3948 dsrc = toPtr(dsrc);
3949
3950 // With runtime activity, a source that is inactive at runtime has
3951 // its primal as its shadow. Its float data has a zero derivative, so
3952 // zero the shadow instead of copying the primal into it.
3953 Value *copyLength = length;
3954 Value *zeroLength = nullptr;
3955 if (floatTy && gutils->runtimeActivity) {
3956 Value *src = toPtr(gutils->getNewFromOriginal(orig_src));
3957 if (src->getType() != dsrc->getType())
3958 src = BuilderZ.CreatePointerBitCastOrAddrSpaceCast(
3959 src, dsrc->getType());
3960 Value *inactive = BuilderZ.CreateICmpEQ(dsrc, src);
3961 Value *zero = ConstantInt::get(length->getType(), 0);
3962 copyLength = BuilderZ.CreateSelect(inactive, zero, length);
3963 zeroLength = BuilderZ.CreateSelect(inactive, length, zero);
3964 }
3965
3966 if (ID == Intrinsic::memmove) {
3967 call =
3968 BuilderZ.CreateMemMove(ddst, dalign, dsrc, salign, copyLength);
3969 } else {
3970 call =
3971 BuilderZ.CreateMemCpy(ddst, dalign, dsrc, salign, copyLength);
3972 }
3973 call->setAttributes(MTI.getAttributes());
3974
3975 if (zeroLength)
3976 BuilderZ.CreateMemSet(
3977 ddst, ConstantInt::get(Type::getInt8Ty(ddst->getContext()), 0),
3978 zeroLength, dalign);
3979 }
3980 // TODO shadow scope/noalias (performance)
3981 call->setMetadata(LLVMContext::MD_alias_scope,
3982 MTI.getMetadata(LLVMContext::MD_alias_scope));
3983 call->setMetadata(LLVMContext::MD_noalias,
3984 MTI.getMetadata(LLVMContext::MD_noalias));
3985 call->setMetadata(LLVMContext::MD_tbaa,
3986 MTI.getMetadata(LLVMContext::MD_tbaa));
3987 call->setMetadata(LLVMContext::MD_tbaa_struct,
3988 MTI.getMetadata(LLVMContext::MD_tbaa_struct));
3989 call->setMetadata(LLVMContext::MD_invariant_group,
3990 MTI.getMetadata(LLVMContext::MD_invariant_group));
3991 if (auto m = hasMetadata(&MTI, "enzyme_truetype"))
3992 if (auto sliced = sliceTrueType(m, seg_start, seg_size))
3993 call->setMetadata("enzyme_truetype", sliced);
3994 call->setTailCallKind(MTI.getTailCallKind());
3995 };
3996
3997 if (Mode == DerivativeMode::ForwardMode ||
3999 applyChainRule(BuilderZ, fwd_rule, shadow_dst, shadow_src);
4000 else
4001 applyChainRule(BuilderZ, rev_rule, shadow_dst, shadow_src);
4002 }
4003
4004 eraseIfUnused(MTI);
4005 }
4006
4007 void visitFenceInst(llvm::FenceInst &FI) {
4008 using namespace llvm;
4009
4010 switch (Mode) {
4011 default:
4012 break;
4015 bool emitReverse = true;
4016 if (EnzymeJuliaAddrLoad) {
4017 if (auto prev = dyn_cast_or_null<CallBase>(FI.getPrevNode())) {
4018 if (auto F = prev->getCalledFunction())
4019 if (F->getName() == "julia.safepoint")
4020 emitReverse = false;
4021 }
4022 if (auto prev = dyn_cast_or_null<CallBase>(FI.getNextNode())) {
4023 if (auto F = prev->getCalledFunction())
4024 if (F->getName() == "julia.safepoint")
4025 emitReverse = false;
4026 }
4027 }
4028 if (emitReverse) {
4029 IRBuilder<> Builder2(&FI);
4030 getReverseBuilder(Builder2);
4031 auto order = FI.getOrdering();
4032 switch (order) {
4033 case AtomicOrdering::Acquire:
4034 order = AtomicOrdering::Release;
4035 break;
4036 case AtomicOrdering::Release:
4037 order = AtomicOrdering::Acquire;
4038 break;
4039 default:
4040 break;
4041 }
4042 Builder2.CreateFence(order, FI.getSyncScopeID());
4043 }
4044 }
4045 }
4046 eraseIfUnused(FI);
4047 }
4048
4049 void visitIntrinsicInst(llvm::IntrinsicInst &II) {
4050 using namespace llvm;
4051
4052 if (II.getIntrinsicID() == Intrinsic::stacksave) {
4053 eraseIfUnused(II, /*erase*/ true, /*check*/ false);
4054 return;
4055 }
4056 if (II.getIntrinsicID() == Intrinsic::stackrestore ||
4057 II.getIntrinsicID() == Intrinsic::lifetime_end ||
4058 II.getCalledFunction()->getName() == "llvm.enzyme.lifetime_end") {
4059 eraseIfUnused(II, /*erase*/ true, /*check*/ false);
4060 return;
4061 }
4062#if LLVM_VERSION_MAJOR >= 20
4063 if (II.getIntrinsicID() == Intrinsic::experimental_memset_pattern) {
4065 return;
4066 }
4067#endif
4068
4069 // When compiling Enzyme against standard LLVM, and not Intel's
4070 // modified version of LLVM, the intrinsic `llvm.intel.subscript` is
4071 // not fully understood by LLVM. One of the results of this is that the ID
4072 // of the intrinsic is set to Intrinsic::not_intrinsic - hence we are
4073 // handling the intrinsic here.
4074 if (isIntelSubscriptIntrinsic(II)) {
4079 }
4080 } else {
4081 SmallVector<Value *, 2> orig_ops(II.getNumOperands());
4082
4083 for (unsigned i = 0; i < II.getNumOperands(); ++i) {
4084 orig_ops[i] = II.getOperand(i);
4085 }
4086 if (handleAdjointForIntrinsic(II.getIntrinsicID(), II, orig_ops))
4087 return;
4088 }
4089 if (gutils->knownRecomputeHeuristic.find(&II) !=
4090 gutils->knownRecomputeHeuristic.end()) {
4091 if (!gutils->knownRecomputeHeuristic[&II]) {
4092 CallInst *const newCall =
4093 cast<CallInst>(gutils->getNewFromOriginal(&II));
4094 IRBuilder<> BuilderZ(newCall);
4095 BuilderZ.setFastMathFlags(getFast());
4096
4097 gutils->cacheForReverse(BuilderZ, newCall,
4098 getIndex(&II, CacheType::Self, BuilderZ));
4099 }
4100 }
4101 eraseIfUnused(II);
4102 }
4103
4104 bool
4105 handleAdjointForIntrinsic(llvm::Intrinsic::ID ID, llvm::Instruction &I,
4106 llvm::SmallVectorImpl<llvm::Value *> &orig_ops) {
4107 using namespace llvm;
4108
4109 Module *M = I.getParent()->getParent()->getParent();
4110
4111 switch (ID) {
4112#if LLVM_VERSION_MAJOR < 20
4113 case Intrinsic::nvvm_ldg_global_i:
4114 case Intrinsic::nvvm_ldg_global_p:
4115 case Intrinsic::nvvm_ldg_global_f:
4116#endif
4117 case Intrinsic::nvvm_ldu_global_i:
4118 case Intrinsic::nvvm_ldu_global_p:
4119 case Intrinsic::nvvm_ldu_global_f: {
4120 auto CI = cast<ConstantInt>(I.getOperand(1));
4121 visitLoadLike(I, /*Align*/ MaybeAlign(CI->getZExtValue()),
4122 /*constantval*/ false);
4123 return false;
4124 }
4125 default:
4126 break;
4127 }
4128
4129 if (ID == Intrinsic::masked_store) {
4130 auto align0 = cast<ConstantInt>(I.getOperand(2))->getZExtValue();
4131 auto align = MaybeAlign(align0);
4132 visitCommonStore(I, /*orig_ptr*/ I.getOperand(1),
4133 /*orig_val*/ I.getOperand(0), align,
4134 /*isVolatile*/ false, llvm::AtomicOrdering::NotAtomic,
4135 SyncScope::SingleThread,
4136 /*mask*/ gutils->getNewFromOriginal(I.getOperand(3)));
4137 return false;
4138 }
4139 if (ID == Intrinsic::masked_load) {
4140 auto align0 = cast<ConstantInt>(I.getOperand(1))->getZExtValue();
4141 auto align = MaybeAlign(align0);
4142 auto &DL = gutils->newFunc->getParent()->getDataLayout();
4143 bool constantval = parseTBAA(I, DL, nullptr)[{-1}].isIntegral();
4144 visitLoadLike(I, align, constantval,
4145 /*mask*/ gutils->getNewFromOriginal(I.getOperand(2)),
4146 /*orig_maskInit*/ I.getOperand(3));
4147 return false;
4148 }
4149
4150 auto mod = I.getParent()->getParent()->getParent();
4151 auto called = cast<CallInst>(&I)->getCalledFunction();
4152 (void)called;
4153 switch (ID) {
4154#include "IntrinsicDerivatives.inc"
4155 default:
4156 break;
4157 }
4158
4159 switch (Mode) {
4161 switch (ID) {
4162#if LLVM_VERSION_MAJOR <= 20
4163 case Intrinsic::nvvm_barrier0:
4164#else
4165 case Intrinsic::nvvm_barrier_cta_sync_aligned_all:
4166 case Intrinsic::nvvm_barrier_cta_sync_aligned_count:
4167#endif
4168#if LLVM_VERSION_MAJOR < 22
4169 case Intrinsic::nvvm_barrier0_popc:
4170 case Intrinsic::nvvm_barrier0_and:
4171 case Intrinsic::nvvm_barrier0_or:
4172#else
4173 case Intrinsic::nvvm_barrier_cta_red_and_aligned_all:
4174 case Intrinsic::nvvm_barrier_cta_red_and_aligned_count:
4175 case Intrinsic::nvvm_barrier_cta_red_or_aligned_all:
4176 case Intrinsic::nvvm_barrier_cta_red_or_aligned_count:
4177 case Intrinsic::nvvm_barrier_cta_red_popc_aligned_all:
4178 case Intrinsic::nvvm_barrier_cta_red_popc_aligned_count:
4179#endif
4180 case Intrinsic::nvvm_membar_cta:
4181 case Intrinsic::nvvm_membar_gl:
4182 case Intrinsic::nvvm_membar_sys:
4183 case Intrinsic::amdgcn_s_barrier:
4184 return false;
4185 default:
4186 if (gutils->isConstantInstruction(&I))
4187 return false;
4188 if (ID == Intrinsic::umax || ID == Intrinsic::smax ||
4189 ID == Intrinsic::abs || ID == Intrinsic::sadd_with_overflow ||
4190 ID == Intrinsic::uadd_with_overflow ||
4191 ID == Intrinsic::smul_with_overflow ||
4192 ID == Intrinsic::umul_with_overflow ||
4193 ID == Intrinsic::ssub_with_overflow ||
4194 ID == Intrinsic::usub_with_overflow)
4195 if (looseTypeAnalysis) {
4196 EmitWarning("CannotDeduceType", I,
4197 "failed to deduce type of intrinsic ", I);
4198 return false;
4199 }
4200 std::string s;
4201 llvm::raw_string_ostream ss(s);
4202 ss << *gutils->oldFunc << "\n";
4203 ss << *gutils->newFunc << "\n";
4204 ss << "cannot handle (augmented) unknown intrinsic\n" << I;
4205 IRBuilder<> BuilderZ(&I);
4206 getForwardBuilder(BuilderZ);
4207 EmitNoDerivativeError(ss.str(), I, gutils, BuilderZ);
4208 return false;
4209 }
4210 return false;
4211 }
4212
4215
4216 IRBuilder<> Builder2(&I);
4217 getReverseBuilder(Builder2);
4218
4219 Value *vdiff = nullptr;
4220 if (!gutils->isConstantValue(&I)) {
4221 vdiff = diffe(&I, Builder2);
4222 setDiffe(&I, Constant::getNullValue(gutils->getShadowType(I.getType())),
4223 Builder2);
4224 }
4225 (void)vdiff;
4226
4227 switch (ID) {
4228#if LLVM_VERSION_MAJOR < 22
4229 case Intrinsic::nvvm_barrier0_popc:
4230 case Intrinsic::nvvm_barrier0_and:
4231 case Intrinsic::nvvm_barrier0_or:
4232#else
4233 case Intrinsic::nvvm_barrier_cta_red_and_aligned_all:
4234 case Intrinsic::nvvm_barrier_cta_red_and_aligned_count:
4235 case Intrinsic::nvvm_barrier_cta_red_or_aligned_all:
4236 case Intrinsic::nvvm_barrier_cta_red_or_aligned_count:
4237 case Intrinsic::nvvm_barrier_cta_red_popc_aligned_all:
4238 case Intrinsic::nvvm_barrier_cta_red_popc_aligned_count:
4239#endif
4240 {
4241 SmallVector<Value *, 1> args = {};
4242#if LLVM_VERSION_MAJOR > 20
4243 args.push_back(ConstantInt::get(Type::getInt32Ty(M->getContext()), 0));
4244 auto cal = cast<CallInst>(Builder2.CreateCall(
4246 M, Intrinsic::nvvm_barrier_cta_sync_aligned_all),
4247 args));
4248 cal->setCallingConv(getIntrinsicDeclaration(
4249 M, Intrinsic::nvvm_barrier_cta_sync_aligned_all)
4250 ->getCallingConv());
4251#else
4252 auto cal = cast<CallInst>(Builder2.CreateCall(
4253 getIntrinsicDeclaration(M, Intrinsic::nvvm_barrier0), args));
4254 cal->setCallingConv(getIntrinsicDeclaration(M, Intrinsic::nvvm_barrier0)
4255 ->getCallingConv());
4256#endif
4257 cal->setDebugLoc(gutils->getNewFromOriginal(I.getDebugLoc()));
4258 return false;
4259 }
4260
4261#if LLVM_VERSION_MAJOR <= 20
4262 case Intrinsic::nvvm_barrier0:
4263#else
4264 case Intrinsic::nvvm_barrier_cta_sync_aligned_all:
4265 case Intrinsic::nvvm_barrier_cta_sync_aligned_count:
4266#endif
4267 case Intrinsic::amdgcn_s_barrier:
4268 case Intrinsic::nvvm_membar_cta:
4269 case Intrinsic::nvvm_membar_gl:
4270 case Intrinsic::nvvm_membar_sys: {
4271 SmallVector<Value *, 2> args = {};
4272#if LLVM_VERSION_MAJOR > 20
4273 if (ID == Intrinsic::nvvm_barrier_cta_sync_aligned_all ||
4274 ID == Intrinsic::nvvm_barrier_cta_sync_aligned_count) {
4275 auto *CB = cast<CallBase>(&I);
4276 for (Use &arg : CB->args())
4277 args.push_back(
4278 lookup(gutils->getNewFromOriginal(arg.get()), Builder2));
4279 }
4280#endif
4281 auto cal = cast<CallInst>(
4282 Builder2.CreateCall(getIntrinsicDeclaration(M, ID), args));
4283 cal->setCallingConv(getIntrinsicDeclaration(M, ID)->getCallingConv());
4284 cal->setDebugLoc(gutils->getNewFromOriginal(I.getDebugLoc()));
4285 return false;
4286 }
4287
4288 case Intrinsic::lifetime_start: {
4289 if (gutils->isConstantInstruction(&I))
4290 return false;
4291 SmallVector<Value *, 2> args = {
4292 lookup(gutils->getNewFromOriginal(orig_ops[0]), Builder2),
4293 lookup(gutils->getNewFromOriginal(orig_ops[1]), Builder2)};
4294 Type *tys[] = {args[1]->getType()};
4295 auto cal = Builder2.CreateCall(
4296 getIntrinsicDeclaration(M, Intrinsic::lifetime_end, tys), args);
4297 cal->setCallingConv(
4298 getIntrinsicDeclaration(M, Intrinsic::lifetime_end, tys)
4299 ->getCallingConv());
4300 return false;
4301 }
4302
4303 case Intrinsic::vector_reduce_fmax: {
4304 if (vdiff && !gutils->isConstantValue(orig_ops[0])) {
4305 auto prev = lookup(gutils->getNewFromOriginal(orig_ops[0]), Builder2);
4306 auto VT = cast<VectorType>(orig_ops[0]->getType());
4307
4308 assert(!VT->getElementCount().isScalable());
4309 size_t numElems = VT->getElementCount().getKnownMinValue();
4310 SmallVector<Value *> elems;
4311 SmallVector<Value *> cmps;
4312
4313 for (size_t i = 0; i < numElems; ++i)
4314 elems.push_back(Builder2.CreateExtractElement(prev, (uint64_t)i));
4315
4316 Value *curmax = elems[0];
4317 for (size_t i = 0; i < numElems - 1; ++i) {
4318 cmps.push_back(Builder2.CreateFCmpOLT(curmax, elems[i + 1]));
4319 if (i + 2 != numElems)
4320 curmax = CreateSelect(Builder2, cmps[i], elems[i + 1], curmax);
4321 }
4322
4323 auto rule = [&](Value *vdiff) {
4324 auto nv = Constant::getNullValue(orig_ops[0]->getType());
4325 Value *res = Builder2.CreateInsertElement(nv, vdiff, (uint64_t)0);
4326
4327 for (size_t i = 0; i < numElems - 1; ++i) {
4328 auto rhs_v = Builder2.CreateInsertElement(nv, vdiff, i + 1);
4329 res = CreateSelect(Builder2, cmps[i], rhs_v, res);
4330 }
4331 return res;
4332 };
4333 Value *dif0 =
4334 applyChainRule(orig_ops[0]->getType(), Builder2, rule, vdiff);
4335 addToDiffe(orig_ops[0], dif0, Builder2, I.getType());
4336 }
4337 return false;
4338 }
4339 default:
4340 if (gutils->isConstantInstruction(&I))
4341 return false;
4342 if (ID == Intrinsic::umax || ID == Intrinsic::smax ||
4343 ID == Intrinsic::abs || ID == Intrinsic::sadd_with_overflow ||
4344 ID == Intrinsic::uadd_with_overflow ||
4345 ID == Intrinsic::smul_with_overflow ||
4346 ID == Intrinsic::umul_with_overflow ||
4347 ID == Intrinsic::ssub_with_overflow ||
4348 ID == Intrinsic::usub_with_overflow)
4349 if (looseTypeAnalysis) {
4350 EmitWarning("CannotDeduceType", I,
4351 "failed to deduce type of intrinsic ", I);
4352 return false;
4353 }
4354 std::string s;
4355 llvm::raw_string_ostream ss(s);
4356 ss << *gutils->oldFunc << "\n";
4357 ss << *gutils->newFunc << "\n";
4358 if (Intrinsic::isOverloaded(ID))
4359 ss << "cannot handle (reverse) unknown intrinsic\n"
4360 << Intrinsic::getName(ID, ArrayRef<Type *>(),
4361 gutils->oldFunc->getParent(), nullptr)
4362 << "\n"
4363 << I;
4364 else
4365 ss << "cannot handle (reverse) unknown intrinsic\n"
4366 << Intrinsic::getName(ID) << "\n"
4367 << I;
4368 EmitNoDerivativeError(ss.str(), I, gutils, Builder2);
4369 return false;
4370 }
4371 return false;
4372 }
4376
4377 IRBuilder<> Builder2(&I);
4378 getForwardBuilder(Builder2);
4379
4380 switch (ID) {
4381
4382 case Intrinsic::vector_reduce_fmax: {
4383 if (gutils->isConstantInstruction(&I))
4384 return false;
4385 auto prev = gutils->getNewFromOriginal(orig_ops[0]);
4386 auto VT = cast<VectorType>(orig_ops[0]->getType());
4387
4388 assert(!VT->getElementCount().isScalable());
4389 size_t numElems = VT->getElementCount().getKnownMinValue();
4390 SmallVector<Value *> elems;
4391 SmallVector<Value *> cmps;
4392
4393 for (size_t i = 0; i < numElems; ++i)
4394 elems.push_back(Builder2.CreateExtractElement(prev, (uint64_t)i));
4395
4396 Value *curmax = elems[0];
4397 for (size_t i = 0; i < numElems - 1; ++i) {
4398 cmps.push_back(Builder2.CreateFCmpOLT(curmax, elems[i + 1]));
4399 if (i + 2 != numElems)
4400 curmax = CreateSelect(Builder2, cmps[i], elems[i + 1], curmax);
4401 }
4402
4403 auto rule = [&](Value *vdiff) {
4404 Value *res = Builder2.CreateExtractElement(vdiff, (uint64_t)0);
4405
4406 for (size_t i = 0; i < numElems - 1; ++i) {
4407 auto rhs_v = Builder2.CreateExtractElement(vdiff, i + 1);
4408 res = CreateSelect(Builder2, cmps[i], rhs_v, res);
4409 }
4410 return res;
4411 };
4412 auto vdiff = diffe(orig_ops[0], Builder2);
4413
4414 Value *dif = applyChainRule(I.getType(), Builder2, rule, vdiff);
4415 setDiffe(&I, dif, Builder2);
4416 return false;
4417 }
4418 default:
4419 if (!gutils->isConstantValue(&I)) {
4420 auto toset =
4421 Constant::getNullValue(gutils->getShadowType(I.getType()));
4422 setDiffe(&I, toset, Builder2);
4423 }
4424 if (gutils->isConstantInstruction(&I))
4425 return false;
4426 if (ID == Intrinsic::umax || ID == Intrinsic::smax ||
4427 ID == Intrinsic::abs || ID == Intrinsic::sadd_with_overflow ||
4428 ID == Intrinsic::uadd_with_overflow ||
4429 ID == Intrinsic::smul_with_overflow ||
4430 ID == Intrinsic::umul_with_overflow ||
4431 ID == Intrinsic::ssub_with_overflow ||
4432 ID == Intrinsic::usub_with_overflow)
4433 if (looseTypeAnalysis) {
4434 EmitWarning("CannotDeduceType", I,
4435 "failed to deduce type of intrinsic ", I);
4436 return false;
4437 }
4438 std::string s;
4439 llvm::raw_string_ostream ss(s);
4440 if (Intrinsic::isOverloaded(ID))
4441 ss << "cannot handle (forward) unknown intrinsic\n"
4442 << Intrinsic::getName(ID, ArrayRef<Type *>(),
4443 gutils->oldFunc->getParent(), nullptr)
4444 << "\n"
4445 << I;
4446 else
4447 ss << "cannot handle (forward) unknown intrinsic\n"
4448 << Intrinsic::getName(ID) << "\n"
4449 << I;
4450 EmitNoDerivativeError(ss.str(), I, gutils, Builder2);
4451 return false;
4452 }
4453 return false;
4454 }
4455 }
4456
4457 return false;
4458 }
4459
4460// first one allows adding attributes to blas functions declared in the second
4461#include "BlasAttributor.inc"
4462#include "BlasDerivatives.inc"
4463
4464 void visitOMPCall(llvm::CallInst &call) {
4465 using namespace llvm;
4466
4467 Function *kmpc = call.getCalledFunction();
4468
4469 if (overwritten_args_map.find(&call) == overwritten_args_map.end()) {
4470 llvm::errs() << " call: " << call << "\n";
4471 for (auto &pair : overwritten_args_map) {
4472 llvm::errs() << " + " << *pair.first << "\n";
4473 }
4474 }
4475
4476 auto found_ow = overwritten_args_map.find(&call);
4477 assert(found_ow != overwritten_args_map.end());
4478 const bool subsequent_calls_may_write = found_ow->second.first;
4479 const std::vector<bool> &overwritten_args = found_ow->second.second;
4480
4481 IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&call));
4482 BuilderZ.setFastMathFlags(getFast());
4483
4484 Function *task = dyn_cast<Function>(call.getArgOperand(2));
4485 if (task == nullptr && isa<ConstantExpr>(call.getArgOperand(2))) {
4486 task = dyn_cast<Function>(
4487 cast<ConstantExpr>(call.getArgOperand(2))->getOperand(0));
4488 }
4489 if (task == nullptr) {
4490 llvm::errs() << "could not derive underlying task from omp call: " << call
4491 << "\n";
4492 llvm_unreachable("could not derive underlying task from omp call");
4493 }
4494 if (task->empty()) {
4495 llvm::errs()
4496 << "could not derive underlying task contents from omp call: " << call
4497 << "\n";
4498 llvm_unreachable(
4499 "could not derive underlying task contents from omp call");
4500 }
4501
4502 auto called = task;
4503 // bool modifyPrimal = true;
4504
4505 bool foreignFunction = called == nullptr;
4506
4507 SmallVector<Value *, 8> args = {0, 0, 0};
4508 SmallVector<Value *, 8> pre_args = {0, 0, 0};
4509 std::vector<DIFFE_TYPE> argsInverted = {DIFFE_TYPE::CONSTANT,
4511 SmallVector<Instruction *, 4> postCreate;
4512 SmallVector<Instruction *, 4> userReplace;
4513
4514 SmallVector<Value *, 4> OutTypes;
4515 SmallVector<Type *, 4> OutFPTypes;
4516
4517 std::vector<bool> nowrite_shadows = {false, false};
4518
4519 for (unsigned i = 3; i < call.arg_size(); ++i) {
4520
4521 auto argi = gutils->getNewFromOriginal(call.getArgOperand(i));
4522
4523 pre_args.push_back(argi);
4524 nowrite_shadows.push_back(false);
4525
4527 IRBuilder<> Builder2(&call);
4528 getReverseBuilder(Builder2);
4529 args.push_back(lookup(argi, Builder2));
4530 }
4531
4532 auto argTy = gutils->getDiffeType(call.getArgOperand(i), foreignFunction);
4533 argsInverted.push_back(argTy);
4534
4535 if (argTy == DIFFE_TYPE::CONSTANT) {
4536 continue;
4537 }
4538
4539 auto argType = argi->getType();
4540
4541 if (argTy == DIFFE_TYPE::DUP_ARG || argTy == DIFFE_TYPE::DUP_NONEED) {
4543 IRBuilder<> Builder2(&call);
4544 getReverseBuilder(Builder2);
4545 args.push_back(
4546 lookup(gutils->invertPointerM(call.getArgOperand(i), Builder2),
4547 Builder2));
4548 }
4549
4550 auto baseOp = getBaseObject(call.getArgOperand(i));
4551 if (auto arg = dyn_cast<Argument>(baseOp)) {
4552 if (arg->getArgNo() < gutils->nowrite_shadows.size() &&
4553 gutils->nowrite_shadows[arg->getArgNo()]) {
4554 nowrite_shadows.back() = true;
4555 }
4556 }
4557 if (isAllocationCall(baseOp, gutils->TLI)) {
4558 assert(!gutils->isConstantValue(baseOp));
4563 bool forwardsShadow = true;
4564 {
4565 auto found = gutils->backwardsOnlyShadows.find(baseOp);
4566 if (found != gutils->backwardsOnlyShadows.end()) {
4567 forwardsShadow = found->second.primalInitialize;
4568 }
4569 }
4570 if (!forwardsShadow) {
4571 nowrite_shadows.back() = true;
4572 }
4573 }
4574 }
4575
4576 pre_args.push_back(
4577 gutils->invertPointerM(call.getArgOperand(i), BuilderZ));
4578
4579 // Note sometimes whattype mistakenly says something should be constant
4580 // [because composed of integer pointers alone]
4581 assert(whatType(argType, Mode) == DIFFE_TYPE::DUP_ARG ||
4582 whatType(argType, Mode) == DIFFE_TYPE::CONSTANT);
4583 } else {
4584 assert(TR.query(call.getArgOperand(i))[{-1}].isFloat());
4585 OutTypes.push_back(call.getArgOperand(i));
4586 OutFPTypes.push_back(argType);
4587 assert(whatType(argType, Mode) == DIFFE_TYPE::OUT_DIFF ||
4588 whatType(argType, Mode) == DIFFE_TYPE::CONSTANT);
4589 }
4590 }
4591
4592 DIFFE_TYPE subretType = DIFFE_TYPE::CONSTANT;
4593
4594 Value *tape = nullptr;
4595 CallInst *augmentcall = nullptr;
4596 // Value *cachereplace = nullptr;
4597
4598 // TODO consider reduction of int 0 args
4599 FnTypeInfo nextTypeInfo(called);
4600
4601 if (called) {
4602 std::map<Value *, std::set<int64_t>> intseen;
4603
4604 TypeTree IntPtr;
4605 IntPtr.insert({-1, -1}, BaseType::Integer);
4606 IntPtr.insert({-1}, BaseType::Pointer);
4607
4608 int argnum = 0;
4609 for (auto &arg : called->args()) {
4610 if (argnum <= 1) {
4611 nextTypeInfo.Arguments.insert(
4612 std::pair<Argument *, TypeTree>(&arg, IntPtr));
4613 nextTypeInfo.KnownValues.insert(
4614 std::pair<Argument *, std::set<int64_t>>(&arg, {0}));
4615 } else {
4616 nextTypeInfo.Arguments.insert(std::pair<Argument *, TypeTree>(
4617 &arg, TR.query(call.getArgOperand(argnum - 2 + 3))));
4618 nextTypeInfo.KnownValues.insert(
4619 std::pair<Argument *, std::set<int64_t>>(
4620 &arg,
4621 TR.knownIntegralValues(call.getArgOperand(argnum - 2 + 3))));
4622 }
4623
4624 ++argnum;
4625 }
4626 nextTypeInfo.Return = TR.query(&call);
4627 }
4628
4629 // std::optional<std::map<std::pair<Instruction*, std::string>, unsigned>>
4630 // sub_index_map;
4631 // Optional<int> tapeIdx;
4632 // Optional<int> returnIdx;
4633 // Optional<int> differetIdx;
4634
4635 const AugmentedReturn *subdata = nullptr;
4637 assert(augmentedReturn);
4638 if (augmentedReturn) {
4639 auto fd = augmentedReturn->subaugmentations.find(&call);
4640 if (fd != augmentedReturn->subaugmentations.end()) {
4641 subdata = fd->second;
4642 }
4643 }
4644 }
4645
4648 if (called) {
4649 subdata = &gutils->Logic.CreateAugmentedPrimal(
4650 RequestContext(&call, &BuilderZ), cast<Function>(called),
4651 subretType, argsInverted, TR.analyzer->interprocedural,
4652 /*return is used*/ false,
4653 /*shadowReturnUsed*/ false, nextTypeInfo,
4654 subsequent_calls_may_write, overwritten_args, nowrite_shadows,
4655 false, gutils->runtimeActivity, gutils->strongZero,
4656 gutils->getWidth(), /*AtomicAdd*/ true,
4657 /*OpenMP*/ true);
4659 assert(augmentedReturn);
4660 auto subaugmentations =
4661 (std::map<const llvm::CallInst *, AugmentedReturn *>
4662 *)&augmentedReturn->subaugmentations;
4664 *subaugmentations, &call, (AugmentedReturn *)subdata);
4665 }
4666
4667 assert(subdata);
4668 auto newcalled = subdata->fn;
4669
4670 if (subdata->returns.find(AugmentedStruct::Tape) !=
4671 subdata->returns.end()) {
4672 ValueToValueMapTy VMap;
4673 newcalled = CloneFunction(newcalled, VMap);
4674 auto tapeArg = newcalled->arg_end();
4675 tapeArg--;
4676 Type *tapeElemType = subdata->tapeType;
4677 SmallVector<std::pair<ssize_t, Value *>, 4> geps;
4678 SmallPtrSet<Instruction *, 4> gepsToErase;
4679 for (auto a : tapeArg->users()) {
4680 if (auto gep = dyn_cast<GetElementPtrInst>(a)) {
4681 auto idx = gep->idx_begin();
4682 idx++;
4683 auto cidx = cast<ConstantInt>(idx->get());
4684 assert(gep->getNumIndices() == 2);
4685 SmallPtrSet<StoreInst *, 1> storesToErase;
4686 for (auto st : gep->users()) {
4687 auto SI = cast<StoreInst>(st);
4688 Value *op = SI->getValueOperand();
4689 storesToErase.insert(SI);
4690 geps.emplace_back(cidx->getLimitedValue(), op);
4691 }
4692 for (auto SI : storesToErase)
4693 SI->eraseFromParent();
4694 gepsToErase.insert(gep);
4695 } else if (auto SI = dyn_cast<StoreInst>(a)) {
4696 Value *op = SI->getValueOperand();
4697 gepsToErase.insert(SI);
4698 geps.emplace_back(-1, op);
4699 } else {
4700 llvm::errs() << "unknown tape user: " << a << "\n";
4701 assert(0 && "unknown tape user");
4702 llvm_unreachable("unknown tape user");
4703 }
4704 }
4705 for (auto gep : gepsToErase)
4706 gep->eraseFromParent();
4707 IRBuilder<> ph(&*newcalled->getEntryBlock().begin());
4708 tape = UndefValue::get(tapeElemType);
4709 ValueToValueMapTy available;
4710 auto subarg = newcalled->arg_begin();
4711 subarg++;
4712 subarg++;
4713 for (size_t i = 3; i < pre_args.size(); ++i) {
4714 available[&*subarg] = pre_args[i];
4715 subarg++;
4716 }
4717 for (auto pair : geps) {
4718 Value *op = pair.second;
4719 Value *alloc = op;
4720 Value *replacement = gutils->unwrapM(op, BuilderZ, available,
4722 tape =
4723 pair.first == -1
4724 ? replacement
4725 : BuilderZ.CreateInsertValue(tape, replacement, pair.first);
4726 if (auto ci = dyn_cast<CastInst>(alloc)) {
4727 alloc = ci->getOperand(0);
4728 }
4729 if (auto uload = dyn_cast<Instruction>(replacement)) {
4730 gutils->unwrappedLoads.erase(uload);
4731 if (auto ci = dyn_cast<CastInst>(replacement)) {
4732 if (auto ucast = dyn_cast<Instruction>(ci->getOperand(0)))
4733 gutils->unwrappedLoads.erase(ucast);
4734 }
4735 }
4736 if (auto ci = dyn_cast<CallInst>(alloc)) {
4737 if (auto F = ci->getCalledFunction()) {
4738 // Store cached values
4739 if (F->getName() == "malloc") {
4740 const_cast<AugmentedReturn *>(subdata)
4741 ->tapeIndiciesToFree.emplace(pair.first);
4742 Value *toload = tapeArg;
4743 if (pair.first != -1) {
4744 Value *Idxs[] = {
4745 ConstantInt::get(
4746 Type::getInt64Ty(tapeArg->getContext()), 0),
4747 ConstantInt::get(
4748 Type::getInt32Ty(tapeArg->getContext()),
4749 pair.first)};
4750 toload = ph.CreateInBoundsGEP(tapeElemType, toload, Idxs);
4751 }
4752 op->replaceAllUsesWith(ph.CreateLoad(op->getType(), toload));
4753 cast<Instruction>(op)->eraseFromParent();
4754 if (op != alloc)
4755 ci->eraseFromParent();
4756 continue;
4757 }
4758 }
4759 }
4760 Value *Idxs[] = {
4761 ConstantInt::get(Type::getInt64Ty(tapeArg->getContext()), 0),
4762 ConstantInt::get(Type::getInt32Ty(tapeArg->getContext()),
4763 pair.first)};
4764 op->replaceAllUsesWith(ph.CreateLoad(
4765 op->getType(),
4766 pair.first == -1
4767 ? tapeArg
4768 : ph.CreateInBoundsGEP(tapeElemType, tapeArg, Idxs)));
4769 cast<Instruction>(op)->eraseFromParent();
4770 }
4771 assert(tape);
4772 auto alloc =
4773 IRBuilder<>(gutils->inversionAllocs).CreateAlloca(tapeElemType);
4774 BuilderZ.CreateStore(tape, alloc);
4775 pre_args.push_back(alloc);
4776 assert(tape);
4777 gutils->cacheForReverse(BuilderZ, tape,
4778 getIndex(&call, CacheType::Tape, BuilderZ));
4779 }
4780
4781 auto numargs = ConstantInt::get(Type::getInt32Ty(call.getContext()),
4782 pre_args.size() - 3);
4783 pre_args[0] = gutils->getNewFromOriginal(call.getArgOperand(0));
4784 pre_args[1] = numargs;
4785 pre_args[2] = BuilderZ.CreatePointerCast(
4786 newcalled, kmpc->getFunctionType()->getParamType(2));
4787 augmentcall =
4788 BuilderZ.CreateCall(kmpc->getFunctionType(), kmpc, pre_args);
4789 augmentcall->setCallingConv(call.getCallingConv());
4790 augmentcall->setDebugLoc(
4791 gutils->getNewFromOriginal(call.getDebugLoc()));
4792 BuilderZ.SetInsertPoint(
4793 gutils->getNewFromOriginal(&call)->getNextNode());
4794 gutils->erase(gutils->getNewFromOriginal(&call));
4795 } else {
4796 assert(0 && "unhandled unknown outline");
4797 }
4798 }
4799
4800 {
4801 Intrinsic::ID ID = Intrinsic::not_intrinsic;
4802 if (!subdata && !isMemFreeLibMFunction(getFuncNameFromCall(&call), &ID)) {
4803 llvm::errs() << *gutils->oldFunc->getParent() << "\n";
4804 llvm::errs() << *gutils->oldFunc << "\n";
4805 llvm::errs() << *gutils->newFunc << "\n";
4806 llvm::errs() << *called << "\n";
4807 llvm_unreachable("no subdata");
4808 }
4809 }
4810
4811 if (subdata) {
4812 auto found = subdata->returns.find(AugmentedStruct::DifferentialReturn);
4813 assert(found == subdata->returns.end());
4814 }
4815 if (subdata) {
4816 auto found = subdata->returns.find(AugmentedStruct::Return);
4817 assert(found == subdata->returns.end());
4818 }
4819
4822 IRBuilder<> Builder2(&call);
4823 getReverseBuilder(Builder2);
4824
4826 BuilderZ.SetInsertPoint(
4827 gutils->getNewFromOriginal(&call)->getNextNode());
4828 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
4829 }
4830
4831 Function *newcalled = nullptr;
4832 if (called) {
4833 if (subdata && subdata->returns.find(AugmentedStruct::Tape) !=
4834 subdata->returns.end()) {
4836 if (tape == nullptr) {
4837#if LLVM_VERSION_MAJOR >= 18
4838 auto It = BuilderZ.GetInsertPoint();
4839 It.setHeadBit(true);
4840 BuilderZ.SetInsertPoint(It);
4841#endif
4842 tape = BuilderZ.CreatePHI(subdata->tapeType, 0, "tapeArg");
4843 }
4844 tape = gutils->cacheForReverse(
4845 BuilderZ, tape, getIndex(&call, CacheType::Tape, BuilderZ));
4846 }
4847 tape = lookup(tape, Builder2);
4848 auto alloc = IRBuilder<>(gutils->inversionAllocs)
4849 .CreateAlloca(tape->getType());
4850 Builder2.CreateStore(tape, alloc);
4851 args.push_back(alloc);
4852 }
4853
4854 if (Mode == DerivativeMode::ReverseModeGradient && subdata) {
4855 for (size_t i = 0; i < argsInverted.size(); i++) {
4856 if (subdata->constant_args[i] == argsInverted[i])
4857 continue;
4858 assert(subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG);
4859 assert(argsInverted[i] == DIFFE_TYPE::DUP_NONEED);
4860 argsInverted[i] = DIFFE_TYPE::DUP_ARG;
4861 }
4862 }
4863
4864 newcalled = gutils->Logic.CreatePrimalAndGradient(
4865 RequestContext(&call, &Builder2),
4867 .todiff = cast<Function>(called),
4868 .retType = subretType,
4869 .constant_args = argsInverted,
4870 .subsequent_calls_may_write = subsequent_calls_may_write,
4871 .overwritten_args = overwritten_args,
4872 .returnUsed = false,
4873 .shadowReturnUsed = false,
4875 .width = gutils->getWidth(),
4876 .freeMemory = true,
4877 .AtomicAdd = true,
4878 .additionalType = tape ? getUnqual(tape->getType()) : nullptr,
4879 .forceAnonymousTape = false,
4880 .typeInfo = nextTypeInfo,
4881 .runtimeActivity = gutils->runtimeActivity,
4882 .strongZero = gutils->strongZero},
4883 TR.analyzer->interprocedural, subdata,
4884 /*omp*/ true);
4885
4886 if (subdata && subdata->returns.find(AugmentedStruct::Tape) !=
4887 subdata->returns.end()) {
4888 auto tapeArg = newcalled->arg_end();
4889 tapeArg--;
4890 LoadInst *tape = nullptr;
4891 for (auto u : tapeArg->users()) {
4892 assert(!tape);
4893 if (!isa<LoadInst>(u)) {
4894 llvm::errs() << " newcalled: " << *newcalled << "\n";
4895 llvm::errs() << " u: " << *u << "\n";
4896 }
4897 tape = cast<LoadInst>(u);
4898 }
4899 assert(tape);
4900 SmallVector<Value *, 4> extracts;
4901 if (subdata->tapeIndices.size() == 1) {
4902 assert(subdata->tapeIndices.begin()->second == -1);
4903 extracts.push_back(tape);
4904 } else {
4905 for (auto a : tape->users()) {
4906 extracts.push_back(a);
4907 }
4908 }
4909 SmallVector<LoadInst *, 4> geps;
4910 for (auto E : extracts) {
4911 AllocaInst *AI = nullptr;
4912 for (auto U : E->users()) {
4913 if (auto SI = dyn_cast<StoreInst>(U)) {
4914 assert(SI->getValueOperand() == E);
4915 AI = cast<AllocaInst>(SI->getPointerOperand());
4916 }
4917 }
4918 if (AI) {
4919 for (auto U : AI->users()) {
4920 if (auto LI = dyn_cast<LoadInst>(U)) {
4921 geps.push_back(LI);
4922 }
4923 }
4924 }
4925 }
4926 for (auto LI : geps) {
4927 CallInst *freeCall = nullptr;
4928 for (auto LU : LI->users()) {
4929 if (auto CI = dyn_cast<CallInst>(LU)) {
4930 if (auto F = CI->getCalledFunction()) {
4931 if (F->getName() == "free") {
4932 freeCall = CI;
4933 break;
4934 }
4935 }
4936 } else if (auto BC = dyn_cast<CastInst>(LU)) {
4937 for (auto CU : BC->users()) {
4938 if (auto CI = dyn_cast<CallInst>(CU)) {
4939 if (auto F = CI->getCalledFunction()) {
4940 if (F->getName() == "free") {
4941 freeCall = CI;
4942 break;
4943 }
4944 }
4945 }
4946 }
4947 if (freeCall)
4948 break;
4949 }
4950 }
4951 if (freeCall) {
4952 freeCall->eraseFromParent();
4953 }
4954 }
4955 }
4956
4957 Value *OutAlloc = nullptr;
4958 auto ST = StructType::get(newcalled->getContext(), OutFPTypes);
4959 if (OutTypes.size()) {
4960 OutAlloc = IRBuilder<>(gutils->inversionAllocs).CreateAlloca(ST);
4961 args.push_back(OutAlloc);
4962
4963 SmallVector<Type *, 3> MetaTypes;
4964 for (auto P :
4965 cast<Function>(newcalled)->getFunctionType()->params()) {
4966 MetaTypes.push_back(P);
4967 }
4968 MetaTypes.push_back(getUnqual(ST));
4969 auto FT = FunctionType::get(Type::getVoidTy(newcalled->getContext()),
4970 MetaTypes, false);
4971 Function *F =
4972 Function::Create(FT, GlobalVariable::InternalLinkage,
4973 cast<Function>(newcalled)->getName() + "#out",
4974 *task->getParent());
4975 BasicBlock *entry =
4976 BasicBlock::Create(newcalled->getContext(), "entry", F);
4977 IRBuilder<> B(entry);
4978 SmallVector<Value *, 2> SubArgs;
4979 for (auto &arg : F->args())
4980 SubArgs.push_back(&arg);
4981 Value *cacheArg = SubArgs.back();
4982 SubArgs.pop_back();
4983 Value *outdiff = B.CreateCall(newcalled, SubArgs);
4984 for (size_t ee = 0; ee < OutTypes.size(); ee++) {
4985 Value *dif = B.CreateExtractValue(outdiff, ee);
4986 Value *Idxs[] = {
4987 ConstantInt::get(Type::getInt64Ty(ST->getContext()), 0),
4988 ConstantInt::get(Type::getInt32Ty(ST->getContext()), ee)};
4989 Value *ptr = B.CreateInBoundsGEP(ST, cacheArg, Idxs);
4990
4991 if (dif->getType()->isIntOrIntVectorTy()) {
4992
4993 ptr = B.CreateBitCast(
4994 ptr,
4996 IntToFloatTy(dif->getType()),
4997 cast<PointerType>(ptr->getType())->getAddressSpace()));
4998 dif = B.CreateBitCast(dif, IntToFloatTy(dif->getType()));
4999 }
5000
5001 MaybeAlign align;
5002 AtomicRMWInst::BinOp op = AtomicRMWInst::FAdd;
5003 if (auto vt = dyn_cast<VectorType>(dif->getType())) {
5004 assert(!vt->getElementCount().isScalable());
5005 size_t numElems = vt->getElementCount().getKnownMinValue();
5006 for (size_t i = 0; i < numElems; ++i) {
5007 auto vdif = B.CreateExtractElement(dif, i);
5008 Value *Idxs[] = {
5009 ConstantInt::get(Type::getInt64Ty(vt->getContext()), 0),
5010 ConstantInt::get(Type::getInt32Ty(vt->getContext()), i)};
5011 auto vptr = B.CreateInBoundsGEP(vt, ptr, Idxs);
5012 B.CreateAtomicRMW(op, vptr, vdif, align,
5013 AtomicOrdering::Monotonic, SyncScope::System);
5014 }
5015 } else {
5016 B.CreateAtomicRMW(op, ptr, dif, align, AtomicOrdering::Monotonic,
5017 SyncScope::System);
5018 }
5019 }
5020 B.CreateRetVoid();
5021 newcalled = F;
5022 }
5023
5024 auto numargs = ConstantInt::get(Type::getInt32Ty(call.getContext()),
5025 args.size() - 3);
5026 args[0] =
5027 lookup(gutils->getNewFromOriginal(call.getArgOperand(0)), Builder2);
5028 args[1] = numargs;
5029 args[2] = Builder2.CreatePointerCast(
5030 newcalled, kmpc->getFunctionType()->getParamType(2));
5031
5032 CallInst *diffes =
5033 Builder2.CreateCall(kmpc->getFunctionType(), kmpc, args);
5034 diffes->setCallingConv(call.getCallingConv());
5035 diffes->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
5036
5037 for (size_t i = 0; i < OutTypes.size(); i++) {
5038
5039 size_t size = 1;
5040 if (OutTypes[i]->getType()->isSized())
5041 size = (gutils->newFunc->getParent()
5042 ->getDataLayout()
5043 .getTypeSizeInBits(OutTypes[i]->getType()) +
5044 7) /
5045 8;
5046 Value *Idxs[] = {
5047 ConstantInt::get(Type::getInt64Ty(call.getContext()), 0),
5048 ConstantInt::get(Type::getInt32Ty(call.getContext()), i)};
5049 ((DiffeGradientUtils *)gutils)
5050 ->addToDiffe(OutTypes[i],
5051 Builder2.CreateLoad(
5052 OutFPTypes[i],
5053 Builder2.CreateInBoundsGEP(ST, OutAlloc, Idxs)),
5054 Builder2, TR.addingType(size, OutTypes[i]));
5055 }
5056
5057 if (tape && shouldFree()) {
5058 for (auto idx : subdata->tapeIndiciesToFree) {
5059 CreateDealloc(Builder2,
5060 idx == -1 ? tape
5061 : Builder2.CreateExtractValue(tape, idx));
5062 }
5063 }
5064 } else {
5065 assert(0 && "openmp indirect unhandled");
5066 }
5067 }
5068 }
5069
5071 llvm::CallInst &call, llvm::Value *origArg, llvm::Value *dsto,
5072 llvm::Value *srco, llvm::Value *len_arg, llvm::IRBuilder<> &Builder2,
5073 llvm::ArrayRef<llvm::OperandBundleDef> ReverseDefs) {
5074 using namespace llvm;
5075
5076 size_t size = 1;
5077 if (auto ci = dyn_cast<ConstantInt>(len_arg)) {
5078 size = ci->getLimitedValue();
5079 }
5080 auto &DL = gutils->newFunc->getParent()->getDataLayout();
5081 auto vd = TR.query(origArg).Data0().ShiftIndices(DL, 0, size, 0);
5082 if (!vd.isKnownPastPointer()) {
5083#if LLVM_VERSION_MAJOR < 17
5084 if (looseTypeAnalysis) {
5085 if (isa<CastInst>(origArg) &&
5086 cast<CastInst>(origArg)->getSrcTy()->isPointerTy() &&
5087 cast<CastInst>(origArg)
5088 ->getSrcTy()
5089 ->getPointerElementType()
5090 ->isFPOrFPVectorTy()) {
5091 vd = TypeTree(ConcreteType(cast<CastInst>(origArg)
5092 ->getSrcTy()
5093 ->getPointerElementType()
5094 ->getScalarType()))
5095 .Only(0, &call);
5096 goto knownF;
5097 }
5098 // Handle LoadInst (common for the Fortran ABI where arguments are
5099 // passed by reference: the buffer argument is a load of the array's
5100 // address from an alloca'd variable)
5101 if (auto LI = dyn_cast<LoadInst>(origArg)) {
5102 auto *PT = dyn_cast<PointerType>(LI->getType());
5103 if (PT && PT->getPointerElementType()->isFPOrFPVectorTy()) {
5104 vd = TypeTree(
5105 ConcreteType(PT->getPointerElementType()->getScalarType()))
5106 .Only(0, &call);
5107 goto knownF;
5108 }
5109 }
5110 // Handle pointer arguments directly (Fortran MPI routines pass all
5111 // arguments by reference, so the buffer argument may be a pointer to
5112 // real data)
5113 if (origArg->getType()->isPointerTy() &&
5114 origArg->getType()->getPointerElementType()->isFPOrFPVectorTy()) {
5115 vd = TypeTree(ConcreteType(origArg->getType()
5116 ->getPointerElementType()
5117 ->getScalarType()))
5118 .Only(0, &call);
5119 goto knownF;
5120 }
5121 }
5122#endif
5123 // With opaque pointers (LLVM >= 17) the pointee type of a pointer is
5124 // not available from its type. Fall back to the element type when the
5125 // buffer argument is an alloca or a GEP thereof, which is ground truth
5126 // even with opaque pointers. This covers the Fortran ABI,
5127 // where every argument is passed by reference and the buffer may be
5128 // the address of a local variable that is only ever accessed by MPI
5129 // routines (so TypeAnalysis never observes its contents).
5130 if (auto AI = dyn_cast<AllocaInst>(origArg)) {
5131 if (AI->getAllocatedType()->isFPOrFPVectorTy()) {
5132 vd = TypeTree(ConcreteType(AI->getAllocatedType()->getScalarType()))
5133 .Only(0, &call);
5134 goto knownF;
5135 }
5136 }
5137 if (auto GEP = dyn_cast<GetElementPtrInst>(origArg)) {
5138 if (GEP->getSourceElementType()->isFPOrFPVectorTy()) {
5139 vd = TypeTree(
5140 ConcreteType(GEP->getSourceElementType()->getScalarType()))
5141 .Only(0, &call);
5142 goto knownF;
5143 }
5144 }
5145 TR.dump();
5146 EmitFailure("CannotDeduceType", call.getDebugLoc(), &call,
5147 "failed to deduce type of copy ", call);
5148 }
5149 knownF:
5150 unsigned start = 0;
5151 while (1) {
5152 unsigned nextStart = size;
5153
5154 auto dt = vd[{-1}];
5155 for (size_t i = start; i < size; ++i) {
5156 bool Legal = true;
5157 dt.checkedOrIn(vd[{(int)i}], /*PointerIntSame*/ true, Legal);
5158 if (!Legal) {
5159 nextStart = i;
5160 break;
5161 }
5162 }
5163 if (!dt.isKnown()) {
5164 TR.dump();
5165 llvm::errs() << " vd:" << vd.str() << " start:" << start
5166 << " size: " << size << " dt:" << dt.str() << "\n";
5167 }
5168 assert(dt.isKnown());
5169
5170 Value *length = len_arg;
5171 if (nextStart != size) {
5172 length = ConstantInt::get(len_arg->getType(), nextStart);
5173 }
5174 if (start != 0)
5175 length = Builder2.CreateSub(
5176 length, ConstantInt::get(len_arg->getType(), start));
5177
5178 if (auto secretty = dt.isFloat()) {
5179 auto offset = start;
5180 if (dsto->getType()->isIntegerTy())
5181 dsto =
5182 Builder2.CreateIntToPtr(dsto, getInt8PtrTy(dsto->getContext()));
5183 unsigned dstaddr =
5184 cast<PointerType>(dsto->getType())->getAddressSpace();
5185 auto secretpt = getPointerType(secretty, dstaddr);
5186 if (offset != 0) {
5187 dsto = Builder2.CreateConstInBoundsGEP1_64(
5188 Type::getInt8Ty(dsto->getContext()), dsto, offset);
5189 }
5190 if (srco->getType()->isIntegerTy())
5191 srco =
5192 Builder2.CreateIntToPtr(srco, getInt8PtrTy(dsto->getContext()));
5193 unsigned srcaddr =
5194 cast<PointerType>(srco->getType())->getAddressSpace();
5195 secretpt = getPointerType(secretty, srcaddr);
5196
5197 if (offset != 0) {
5198 srco = Builder2.CreateConstInBoundsGEP1_64(
5199 Type::getInt8Ty(srco->getContext()), srco, offset);
5200 }
5201 Value *args[3] = {
5202 Builder2.CreatePointerCast(dsto, secretpt),
5203 Builder2.CreatePointerCast(srco, secretpt),
5204 Builder2.CreateUDiv(
5205 length,
5206
5207 ConstantInt::get(length->getType(),
5208 Builder2.GetInsertBlock()
5209 ->getParent()
5210 ->getParent()
5211 ->getDataLayout()
5212 .getTypeAllocSizeInBits(secretty) /
5213 8))};
5214
5216 *Builder2.GetInsertBlock()->getParent()->getParent(), secretty,
5217 /*dstalign*/ 1, /*srcalign*/ 1, dstaddr, srcaddr,
5218 cast<IntegerType>(length->getType())->getBitWidth(),
5219 /*runtimeActivity*/ false, gutils->isAtomic(srco));
5220
5221 Builder2.CreateCall(dmemcpy, args, ReverseDefs);
5222 }
5223
5224 if (nextStart == size)
5225 break;
5226 start = nextStart;
5227 }
5228 }
5229
5230 void recursivelyHandleSubfunction(llvm::CallInst &call,
5231 llvm::Function *called,
5232 bool subsequent_calls_may_write,
5233 const std::vector<bool> &overwritten_args,
5234 bool shadowReturnUsed,
5235 DIFFE_TYPE subretType, bool subretused) {
5236 using namespace llvm;
5237
5238 IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&call));
5239 BuilderZ.setFastMathFlags(getFast());
5240
5241 CallInst *newCall = cast<CallInst>(gutils->getNewFromOriginal(&call));
5242 Module &M = *call.getParent()->getParent()->getParent();
5243
5244 bool foreignFunction = called == nullptr;
5245
5246 FnTypeInfo nextTypeInfo(called);
5247
5248 if (called) {
5249 nextTypeInfo = TR.getCallInfo(call, *called);
5250 }
5251
5252 const AugmentedReturn *subdata = nullptr;
5255 assert(augmentedReturn);
5256 if (augmentedReturn) {
5257 auto fd = augmentedReturn->subaugmentations.find(&call);
5258 if (fd != augmentedReturn->subaugmentations.end()) {
5259 subdata = fd->second;
5260 }
5261 }
5262 }
5263
5264 if (Mode == DerivativeMode::ForwardMode ||
5267 IRBuilder<> Builder2(&call);
5268 getForwardBuilder(Builder2);
5269
5270 SmallVector<Value *, 8> args;
5271 std::vector<DIFFE_TYPE> argsInverted;
5272 std::map<int, Type *> gradByVal;
5273 std::map<int, std::vector<Attribute>> structAttrs;
5274
5275 for (unsigned i = 0; i < call.arg_size(); ++i) {
5276
5277 if (call.paramHasAttr(i, Attribute::StructRet)) {
5278 structAttrs[args.size()].push_back(Attribute::get(
5279 call.getContext(), "enzyme_sret",
5280 convertSRetTypeToString(call.getParamAttr(i, Attribute::StructRet)
5281 .getValueAsType())));
5282 }
5283 for (auto attr : {"enzymejl_returnRoots", "enzymejl_parmtype",
5284 "enzymejl_parmtype_ref", "enzyme_type",
5285 "enzymejl_sret_union_bytes", "enzymejl_rooted_typ"})
5286 if (call.getAttributes().hasParamAttr(i, attr)) {
5287 structAttrs[args.size()].push_back(call.getParamAttr(i, attr));
5288 }
5289 for (auto ty : PrimalParamAttrsToPreserve)
5290 if (call.getAttributes().hasParamAttr(i, ty)) {
5291 auto attr = call.getAttributes().getParamAttr(i, ty);
5292 structAttrs[args.size()].push_back(attr);
5293 }
5294
5295 auto argi = gutils->getNewFromOriginal(call.getArgOperand(i));
5296
5297 if (call.isByValArgument(i)) {
5298 gradByVal[args.size()] = call.getParamByValType(i);
5299 }
5300
5301 bool writeOnlyNoCapture = true;
5302 bool readOnly = true;
5303 if (!isNoCapture(&call, i)) {
5304 writeOnlyNoCapture = false;
5305 }
5306 if (!isWriteOnly(&call, i)) {
5307 writeOnlyNoCapture = false;
5308 }
5309 if (!isReadOnly(&call, i)) {
5310 readOnly = false;
5311 }
5312
5313 if (shouldDisableNoWrite(&call))
5314 writeOnlyNoCapture = false;
5315
5316 auto argTy =
5317 gutils->getDiffeType(call.getArgOperand(i), foreignFunction);
5318
5319 bool replace =
5320 (argTy == DIFFE_TYPE::DUP_NONEED &&
5321 (writeOnlyNoCapture ||
5322 !isa<Argument>(getBaseObject(call.getArgOperand(i))))) ||
5323 (writeOnlyNoCapture && Mode == DerivativeMode::ForwardModeSplit) ||
5324 (writeOnlyNoCapture && readOnly);
5325
5326 if (replace) {
5327 argi = getUndefinedValueForType(M, argi->getType());
5328 }
5329 argsInverted.push_back(argTy);
5330 args.push_back(argi);
5331
5332 if (argTy == DIFFE_TYPE::CONSTANT) {
5333 continue;
5334 }
5335
5336 if (gutils->getWidth() == 1)
5337 for (auto ty : ShadowParamAttrsToPreserve)
5338 if (call.getAttributes().hasParamAttr(i, ty)) {
5339 auto attr = call.getAttributes().getParamAttr(i, ty);
5340 structAttrs[args.size()].push_back(attr);
5341 }
5342
5343 for (auto attr : {"enzymejl_returnRoots", "enzymejl_parmtype",
5344 "enzymejl_parmtype_ref", "enzyme_type",
5345 "enzymejl_sret_union_bytes", "enzymejl_rooted_typ"})
5346 if (call.getAttributes().hasParamAttr(i, attr)) {
5347 if (gutils->getWidth() == 1) {
5348 structAttrs[args.size()].push_back(call.getParamAttr(i, attr));
5349 } else if (attr == std::string("enzymejl_returnRoots")) {
5350 structAttrs[args.size()].push_back(
5351 Attribute::get(call.getContext(), "enzymejl_returnRoots_v",
5352 call.getAttributes()
5353 .getParamAttr(i, "enzymejl_returnRoots")
5354 .getValueAsString()));
5355 } else if (attr == std::string("enzymejl_sret_union_bytes")) {
5356 structAttrs[args.size()].push_back(Attribute::get(
5357 call.getContext(), "enzymejl_sret_union_bytes_v",
5358 call.getAttributes()
5359 .getParamAttr(i, "enzymejl_sret_union_bytes")
5360 .getValueAsString()));
5361 } else if (attr == std::string("enzymejl_rooted_typ")) {
5362 structAttrs[args.size()].push_back(
5363 Attribute::get(call.getContext(), "enzymejl_rooted_typ_v",
5364 call.getAttributes()
5365 .getParamAttr(i, "enzymejl_rooted_typ")
5366 .getValueAsString()));
5367 }
5368 }
5369 if (call.paramHasAttr(i, Attribute::StructRet)) {
5370 if (gutils->getWidth() == 1) {
5371 structAttrs[args.size()].push_back(
5372 Attribute::get(call.getContext(), "enzyme_sret",
5374 call.getParamAttr(i, Attribute::StructRet)
5375 .getValueAsType())));
5376 } else {
5377 structAttrs[args.size()].push_back(
5378 Attribute::get(call.getContext(), "enzyme_sret_v",
5380 call.getParamAttr(i, Attribute::StructRet)
5381 .getValueAsType())));
5382 }
5383 }
5384
5385 assert(argTy == DIFFE_TYPE::DUP_ARG || argTy == DIFFE_TYPE::DUP_NONEED);
5386
5387 args.push_back(gutils->invertPointerM(call.getArgOperand(i), Builder2));
5388 }
5389#if LLVM_VERSION_MAJOR >= 16
5390 std::optional<int> tapeIdx;
5391#else
5392 Optional<int> tapeIdx;
5393#endif
5394 if (subdata) {
5395 auto found = subdata->returns.find(AugmentedStruct::Tape);
5396 if (found != subdata->returns.end()) {
5397 tapeIdx = found->second;
5398 }
5399 }
5400 Value *tape = nullptr;
5401 if (tapeIdx) {
5402
5403 auto idx = *tapeIdx;
5404 FunctionType *FT = subdata->fn->getFunctionType();
5405#if LLVM_VERSION_MAJOR >= 18
5406 auto It = BuilderZ.GetInsertPoint();
5407 It.setHeadBit(true);
5408 BuilderZ.SetInsertPoint(It);
5409#endif
5410 tape = BuilderZ.CreatePHI(
5411 (tapeIdx == -1)
5412 ? FT->getReturnType()
5413 : cast<StructType>(FT->getReturnType())->getElementType(idx),
5414 1, "tapeArg");
5415
5416 assert(!tape->getType()->isEmptyTy());
5417 gutils->TapesToPreventRecomputation.insert(cast<Instruction>(tape));
5418 tape = gutils->cacheForReverse(
5419 BuilderZ, tape, getIndex(&call, CacheType::Tape, BuilderZ));
5420 args.push_back(tape);
5421 }
5422
5423 Value *newcalled = nullptr;
5424 FunctionType *FT = nullptr;
5425
5426 if (called) {
5427 newcalled = gutils->Logic.CreateForwardDiff(
5428 RequestContext(&call, &BuilderZ), cast<Function>(called),
5429 subretType, argsInverted, TR.analyzer->interprocedural,
5430 /*returnValue*/ subretused, Mode,
5431 ((DiffeGradientUtils *)gutils)->FreeMemory, gutils->runtimeActivity,
5432 gutils->strongZero, gutils->getWidth(),
5433 tape ? tape->getType() : nullptr, nextTypeInfo,
5434 subsequent_calls_may_write, overwritten_args,
5435 /*augmented*/ subdata);
5436 FT = cast<Function>(newcalled)->getFunctionType();
5437 } else {
5438 auto callval = call.getCalledOperand();
5439 newcalled = gutils->invertPointerM(callval, BuilderZ);
5440
5441 if (gutils->getWidth() > 1) {
5442 newcalled = BuilderZ.CreateExtractValue(newcalled, {0});
5443 }
5444
5446 BuilderZ, gutils->getNewFromOriginal(callval), newcalled,
5447 "Attempting to call an indirect active function "
5448 "whose runtime value is inactive",
5449 gutils->getNewFromOriginal(call.getDebugLoc()), &call);
5450
5451 auto ft = call.getFunctionType();
5452 bool retActive = subretType != DIFFE_TYPE::CONSTANT;
5453
5455 ft, Mode, gutils->getWidth(), tape ? tape->getType() : nullptr,
5456 argsInverted, false, /*returnTape*/ false,
5457 /*returnPrimal*/ subretused, /*returnShadow*/ retActive);
5458 PointerType *fptype = getUnqual(FT);
5459 newcalled = BuilderZ.CreatePointerCast(newcalled, getUnqual(fptype));
5460 newcalled = BuilderZ.CreateLoad(fptype, newcalled);
5461 }
5462
5463 assert(newcalled);
5464 assert(FT);
5465
5466 SmallVector<ValueType, 2> BundleTypes;
5467 for (auto A : argsInverted)
5468 if (A == DIFFE_TYPE::CONSTANT)
5469 BundleTypes.push_back(ValueType::Primal);
5470 else
5471 BundleTypes.push_back(ValueType::Both);
5472
5473 auto Defs = gutils->getInvertedBundles(&call, BundleTypes, Builder2,
5474 /*lookup*/ false);
5475
5476 CallInst *diffes = Builder2.CreateCall(FT, newcalled, args, Defs);
5477 diffes->setCallingConv(call.getCallingConv());
5478 diffes->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
5479
5480 for (auto pair : gradByVal) {
5481 diffes->addParamAttr(
5482 pair.first,
5483 Attribute::getWithByValType(diffes->getContext(), pair.second));
5484 }
5485
5486 for (auto &pair : structAttrs) {
5487 for (auto val : pair.second)
5488 diffes->addParamAttr(pair.first, val);
5489 }
5490
5491 auto newcall = gutils->getNewFromOriginal(&call);
5492 auto ifound = gutils->invertedPointers.find(&call);
5493 Value *primal = nullptr;
5494 Value *diffe = nullptr;
5495
5496 if (subretused && subretType != DIFFE_TYPE::CONSTANT) {
5497 primal = Builder2.CreateExtractValue(diffes, 0);
5498 diffe = Builder2.CreateExtractValue(diffes, 1);
5499 } else if (subretType != DIFFE_TYPE::CONSTANT) {
5500 diffe = diffes;
5501 } else if (!FT->getReturnType()->isVoidTy()) {
5502 primal = diffes;
5503 }
5504
5505 if (ifound != gutils->invertedPointers.end()) {
5506 auto placeholder = cast<PHINode>(&*ifound->second);
5507 if (primal) {
5508 gutils->replaceAWithB(newcall, primal);
5509 gutils->erase(newcall);
5510 } else {
5511 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
5512 }
5513 if (diffe) {
5514 gutils->replaceAWithB(placeholder, diffe);
5515 } else {
5516 gutils->invertedPointers.erase(ifound);
5517 }
5518 gutils->erase(placeholder);
5519 } else {
5520 if (primal && diffe) {
5521 gutils->replaceAWithB(newcall, primal);
5522 if (!gutils->isConstantValue(&call)) {
5523 setDiffe(&call, diffe, Builder2);
5524 }
5525 gutils->erase(newcall);
5526 } else if (diffe) {
5527 setDiffe(&call, diffe, Builder2);
5528 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
5529 } else if (primal) {
5530 gutils->replaceAWithB(newcall, primal);
5531 gutils->erase(newcall);
5532 } else {
5533 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
5534 }
5535 }
5536
5537 return;
5538 }
5539
5540 bool modifyPrimal = shouldAugmentCall(&call, gutils);
5541
5542 SmallVector<Value *, 8> args;
5543 SmallVector<Value *, 8> pre_args;
5544 std::vector<DIFFE_TYPE> argsInverted;
5545 SmallVector<Instruction *, 4> postCreate;
5546 SmallVector<Instruction *, 4> userReplace;
5547 std::map<int, Type *> preByVal;
5548 std::map<int, Type *> gradByVal;
5549 std::map<int, std::vector<Attribute>> structAttrs;
5550
5551 bool replaceFunction = false;
5552
5553 if (Mode == DerivativeMode::ReverseModeCombined && !foreignFunction) {
5554 replaceFunction = legalCombinedForwardReverse(
5555 &call, *replacedReturns, postCreate, userReplace, gutils,
5556 unnecessaryInstructions, oldUnreachable, subretused);
5557 if (replaceFunction) {
5558 modifyPrimal = false;
5559 }
5560 }
5561
5562 SmallVector<ValueType, 2> PreBundleTypes;
5563 SmallVector<ValueType, 2> BundleTypes;
5564 std::vector<bool> nowrite_shadows;
5565
5566 for (unsigned i = 0; i < call.arg_size(); ++i) {
5567
5568 auto argi = gutils->getNewFromOriginal(call.getArgOperand(i));
5569
5570 if (call.isByValArgument(i)) {
5571 preByVal[pre_args.size()] = call.getParamByValType(i);
5572 }
5573 for (auto attr : {"enzymejl_returnRoots", "enzymejl_parmtype",
5574 "enzymejl_parmtype_ref", "enzyme_type",
5575 "enzymejl_sret_union_bytes", "enzymejl_rooted_typ"})
5576 if (call.getAttributes().hasParamAttr(i, attr)) {
5577 structAttrs[pre_args.size()].push_back(call.getParamAttr(i, attr));
5578 }
5579 if (call.paramHasAttr(i, Attribute::StructRet)) {
5580 structAttrs[pre_args.size()].push_back(Attribute::get(
5581 call.getContext(), "enzyme_sret",
5583 call.getParamAttr(i, Attribute::StructRet).getValueAsType())));
5584 }
5585 for (auto ty : PrimalParamAttrsToPreserve)
5586 if (call.getAttributes().hasParamAttr(i, ty)) {
5587 auto attr = call.getAttributes().getParamAttr(i, ty);
5588 structAttrs[pre_args.size()].push_back(attr);
5589 }
5590
5591 auto argTy = gutils->getDiffeType(call.getArgOperand(i), foreignFunction);
5592
5593 bool writeOnlyNoCapture = true;
5594 bool readNoneNoCapture = false;
5595 if (!isNoCapture(&call, i)) {
5596 writeOnlyNoCapture = false;
5597 readNoneNoCapture = false;
5598 }
5599 if (!isWriteOnly(&call, i)) {
5600 writeOnlyNoCapture = false;
5601 }
5602 if (!(isReadOnly(&call, i) && isWriteOnly(&call, i))) {
5603 readNoneNoCapture = false;
5604 }
5605
5606 if (shouldDisableNoWrite(&call)) {
5607 writeOnlyNoCapture = false;
5608 readNoneNoCapture = false;
5609 }
5610
5611 Value *prearg = argi;
5612
5613 ValueType preType = ValueType::Primal;
5614 ValueType revType = ValueType::Primal;
5615
5616 nowrite_shadows.push_back(false);
5617
5618 // Keep the existing passed value if coming from outside.
5619 if (readNoneNoCapture ||
5620 (argTy == DIFFE_TYPE::DUP_NONEED &&
5621 (writeOnlyNoCapture ||
5622 !isa<Argument>(getBaseObject(call.getArgOperand(i)))))) {
5623 prearg = getUndefinedValueForType(M, argi->getType());
5624 preType = ValueType::None;
5625 }
5626 pre_args.push_back(prearg);
5627
5629 IRBuilder<> Builder2(&call);
5630 getReverseBuilder(Builder2);
5631
5632 if (call.isByValArgument(i)) {
5633 gradByVal[args.size()] = call.getParamByValType(i);
5634 }
5635
5636 if ((writeOnlyNoCapture && !replaceFunction) ||
5637 (readNoneNoCapture ||
5638 (argTy == DIFFE_TYPE::DUP_NONEED &&
5639 (writeOnlyNoCapture ||
5640 !isa<Argument>(getBaseObject(call.getOperand(i))))))) {
5641 argi = getUndefinedValueForType(M, argi->getType());
5642 revType = ValueType::None;
5643 }
5644 args.push_back(lookup(argi, Builder2));
5645 }
5646
5647 argsInverted.push_back(argTy);
5648
5649 if (argTy == DIFFE_TYPE::CONSTANT) {
5650 PreBundleTypes.push_back(preType);
5651 BundleTypes.push_back(revType);
5652 continue;
5653 }
5654
5655 auto argType = argi->getType();
5656
5657 if (argTy == DIFFE_TYPE::DUP_ARG || argTy == DIFFE_TYPE::DUP_NONEED) {
5658 if (gutils->getWidth() == 1)
5659 for (auto ty : ShadowParamAttrsToPreserve)
5660 if (call.getAttributes().hasParamAttr(i, ty)) {
5661 auto attr = call.getAttributes().getParamAttr(i, ty);
5662 structAttrs[pre_args.size()].push_back(attr);
5663 }
5664
5665 for (auto attr : {"enzymejl_returnRoots", "enzymejl_parmtype",
5666 "enzymejl_parmtype_ref", "enzyme_type",
5667 "enzymejl_sret_union_bytes", "enzymejl_rooted_typ"})
5668 if (call.getAttributes().hasParamAttr(i, attr)) {
5669 if (gutils->getWidth() == 1) {
5670 structAttrs[pre_args.size()].push_back(
5671 call.getParamAttr(i, attr));
5672 } else if (attr == std::string("enzymejl_returnRoots")) {
5673 structAttrs[pre_args.size()].push_back(
5674 Attribute::get(call.getContext(), "enzymejl_returnRoots_v",
5675 call.getAttributes()
5676 .getParamAttr(i, attr)
5677 .getValueAsString()));
5678 } else if (attr == std::string("enzymejl_sret_union_bytes")) {
5679 structAttrs[pre_args.size()].push_back(Attribute::get(
5680 call.getContext(), "enzymejl_sret_union_bytes_v",
5681 call.getAttributes()
5682 .getParamAttr(i, attr)
5683 .getValueAsString()));
5684 } else if (attr == std::string("enzymejl_rooted_typ")) {
5685 structAttrs[pre_args.size()].push_back(
5686 Attribute::get(call.getContext(), "enzymejl_rooted_typ_v",
5687 call.getAttributes()
5688 .getParamAttr(i, attr)
5689 .getValueAsString()));
5690 }
5691 }
5692 if (call.paramHasAttr(i, Attribute::StructRet)) {
5693 if (gutils->getWidth() == 1) {
5694 structAttrs[pre_args.size()].push_back(
5695 Attribute::get(call.getContext(), "enzyme_sret",
5697 call.getParamAttr(i, Attribute::StructRet)
5698 .getValueAsType())));
5699 } else {
5700 structAttrs[pre_args.size()].push_back(
5701 Attribute::get(call.getContext(), "enzyme_sret_v",
5703 call.getParamAttr(i, Attribute::StructRet)
5704 .getValueAsType())));
5705 }
5706 }
5708 IRBuilder<> Builder2(&call);
5709 getReverseBuilder(Builder2);
5710
5711 Value *darg = nullptr;
5712
5713 if (((writeOnlyNoCapture && TR.query(call.getArgOperand(
5714 i))[{-1, -1}] == BaseType::Pointer) ||
5715 gutils->isConstantInstruction(&call)) &&
5716 !replaceFunction) {
5718 M, gutils->getShadowType(argi->getType()));
5719 } else {
5720 darg = gutils->invertPointerM(call.getArgOperand(i), Builder2);
5721 revType = (revType == ValueType::None) ? ValueType::Shadow
5723 }
5724 args.push_back(lookup(darg, Builder2));
5725 }
5726
5727 auto baseOp = getBaseObject(call.getArgOperand(i));
5728 if (auto arg = dyn_cast<Argument>(baseOp)) {
5729 if (arg->getArgNo() < gutils->nowrite_shadows.size() &&
5730 gutils->nowrite_shadows[arg->getArgNo()]) {
5731 nowrite_shadows.back() = true;
5732 }
5733 }
5734 if (isAllocationCall(baseOp, gutils->TLI)) {
5735 assert(!gutils->isConstantValue(baseOp));
5740 bool forwardsShadow = true;
5741 {
5742 auto found = gutils->backwardsOnlyShadows.find(baseOp);
5743 if (found != gutils->backwardsOnlyShadows.end()) {
5744 forwardsShadow = found->second.primalInitialize;
5745 }
5746 }
5747 if (!forwardsShadow) {
5748 nowrite_shadows.back() = true;
5749 }
5750 }
5751 }
5752
5753 if (Mode == DerivativeMode::ReverseModeGradient && !replaceFunction) {
5754 nowrite_shadows.back() = true;
5755 pre_args.push_back(getUndefinedValueForType(M, argi->getType()));
5756 } else {
5757 pre_args.push_back(
5758 gutils->invertPointerM(call.getArgOperand(i), BuilderZ));
5759 }
5760 preType =
5762
5763 // Note sometimes whattype mistakenly says something should be
5764 // constant [because composed of integer pointers alone]
5765 auto wt = whatType(argType, Mode);
5766 if (wt != DIFFE_TYPE::DUP_ARG && wt != DIFFE_TYPE::CONSTANT) {
5767 std::string str;
5768 raw_string_ostream ss(str);
5769 ss << "Mismatched estimated activity type for " << *argType
5770 << " expected DUP_ARG or CONSTANT found " << wt
5771 << ", call = " << call << "\n";
5772 if (CustomErrorHandler) {
5773 CustomErrorHandler(str.c_str(), wrap(&call),
5774 ErrorType::InternalError, nullptr, nullptr,
5775 nullptr);
5776 } else {
5777 EmitFailure("MismatchArgType", call.getDebugLoc(), &call, ss.str());
5778 }
5779 }
5780 } else {
5781 if (foreignFunction)
5782 assert(!argType->isIntOrIntVectorTy());
5783 assert(whatType(argType, Mode) == DIFFE_TYPE::OUT_DIFF ||
5784 whatType(argType, Mode) == DIFFE_TYPE::CONSTANT);
5785 }
5786 PreBundleTypes.push_back(preType);
5787 BundleTypes.push_back(revType);
5788 }
5789 if (called) {
5790 if (call.arg_size() !=
5791 cast<Function>(called)->getFunctionType()->getNumParams()) {
5792 llvm::errs() << *gutils->oldFunc->getParent() << "\n";
5793 llvm::errs() << *gutils->oldFunc << "\n";
5794 llvm::errs() << call << "\n";
5795 llvm::errs() << " number of arg operands != function parameters\n";
5796 EmitFailure("MismatchArgs", call.getDebugLoc(), &call,
5797 "Number of arg operands != function parameters\n", call);
5798 }
5799 }
5800
5801 Value *tape = nullptr;
5802 CallInst *augmentcall = nullptr;
5803 Value *cachereplace = nullptr;
5804
5805 // std::optional<std::map<std::pair<Instruction*, std::string>,
5806 // unsigned>> sub_index_map;
5807#if LLVM_VERSION_MAJOR >= 16
5808 std::optional<int> tapeIdx;
5809 std::optional<int> returnIdx;
5810 std::optional<int> differetIdx;
5811#else
5812 Optional<int> tapeIdx;
5813 Optional<int> returnIdx;
5814 Optional<int> differetIdx;
5815#endif
5816 if (modifyPrimal) {
5817
5818 Value *newcalled = nullptr;
5819 FunctionType *FT = nullptr;
5820 const AugmentedReturn *fnandtapetype = nullptr;
5821
5822 if (!called) {
5823 auto callval = call.getCalledOperand();
5824 Value *uncast = callval;
5825 while (auto CE = dyn_cast<ConstantExpr>(uncast)) {
5826 if (CE->isCast()) {
5827 uncast = CE->getOperand(0);
5828 continue;
5829 }
5830 break;
5831 }
5832 if (isa<ConstantInt>(uncast)) {
5833 std::string str;
5834 raw_string_ostream ss(str);
5835 ss << "cannot find shadow for " << *callval
5836 << " for use as function in " << call;
5837 EmitNoDerivativeError(ss.str(), call, gutils, BuilderZ);
5838 }
5839 newcalled = gutils->invertPointerM(callval, BuilderZ);
5840
5843 BuilderZ, gutils->getNewFromOriginal(callval), newcalled,
5844 "Attempting to call an indirect active function "
5845 "whose runtime value is inactive",
5846 gutils->getNewFromOriginal(call.getDebugLoc()), &call);
5847
5848 FunctionType *ft = call.getFunctionType();
5849
5850 std::set<llvm::Type *> seen;
5851 DIFFE_TYPE subretType = whatType(call.getType(), Mode,
5852 /*intAreConstant*/ false, seen);
5854 ft, /*returnUsed*/ true, /*subretType*/ subretType);
5855 FT = FunctionType::get(
5856 StructType::get(newcalled->getContext(), res.second), res.first,
5857 ft->isVarArg());
5858 auto fptype = getUnqual(FT);
5859 newcalled = BuilderZ.CreatePointerCast(newcalled, getUnqual(fptype));
5860 newcalled = BuilderZ.CreateLoad(fptype, newcalled);
5861 tapeIdx = 0;
5862
5863 if (!call.getType()->isVoidTy()) {
5864 returnIdx = 1;
5865 if (subretType == DIFFE_TYPE::DUP_ARG ||
5866 subretType == DIFFE_TYPE::DUP_NONEED) {
5867 differetIdx = 2;
5868 }
5869 }
5870 } else {
5873 subdata = &gutils->Logic.CreateAugmentedPrimal(
5874 RequestContext(&call, &BuilderZ), cast<Function>(called),
5875 subretType, argsInverted, TR.analyzer->interprocedural,
5876 /*return is used*/ subretused, shadowReturnUsed, nextTypeInfo,
5877 subsequent_calls_may_write, overwritten_args, nowrite_shadows,
5878 false, gutils->runtimeActivity, gutils->strongZero,
5879 gutils->getWidth(), gutils->AtomicAdd);
5881 assert(augmentedReturn);
5882 auto subaugmentations =
5883 (std::map<const llvm::CallInst *, AugmentedReturn *>
5884 *)&augmentedReturn->subaugmentations;
5886 *subaugmentations, &call, (AugmentedReturn *)subdata);
5887 }
5888 }
5889 {
5890 Intrinsic::ID ID = Intrinsic::not_intrinsic;
5891 if (!subdata &&
5893 llvm::errs() << *gutils->oldFunc->getParent() << "\n";
5894 llvm::errs() << *gutils->oldFunc << "\n";
5895 llvm::errs() << *gutils->newFunc << "\n";
5896 llvm::errs() << *called << "\n";
5897 assert(subdata);
5898 }
5899 }
5900
5901 if (subdata) {
5902 fnandtapetype = subdata;
5903 newcalled = subdata->fn;
5904 FT = cast<Function>(newcalled)->getFunctionType();
5905
5906 auto found =
5908 if (found != subdata->returns.end()) {
5909 differetIdx = found->second;
5910 } else {
5911 assert(!shadowReturnUsed);
5912 }
5913
5914 found = subdata->returns.find(AugmentedStruct::Return);
5915 if (found != subdata->returns.end()) {
5916 returnIdx = found->second;
5917 } else {
5918 assert(!subretused);
5919 }
5920
5921 found = subdata->returns.find(AugmentedStruct::Tape);
5922 if (found != subdata->returns.end()) {
5923 tapeIdx = found->second;
5924 }
5925 }
5926 }
5927 // sub_index_map = fnandtapetype.tapeIndices;
5928
5929 // llvm::errs() << "seeing sub_index_map of " << sub_index_map->size()
5930 // << " in ap " << cast<Function>(called)->getName() << "\n";
5933
5934 assert(newcalled);
5935 assert(FT);
5936
5937 if (false) {
5938 badaugmentedfn:;
5939 auto NC = dyn_cast<Function>(newcalled);
5940 llvm::errs() << *gutils->oldFunc << "\n";
5941 llvm::errs() << *gutils->newFunc << "\n";
5942 if (NC)
5943 llvm::errs() << " trying to call " << NC->getName() << " " << *FT
5944 << "\n";
5945 else
5946 llvm::errs() << " trying to call " << *newcalled << " " << *FT
5947 << "\n";
5948
5949 for (unsigned i = 0; i < pre_args.size(); ++i) {
5950 assert(pre_args[i]);
5951 assert(pre_args[i]->getType());
5952 llvm::errs() << "args[" << i << "] = " << *pre_args[i]
5953 << " FT:" << *FT->getParamType(i) << "\n";
5954 }
5955 assert(0 && "calling with wrong number of arguments");
5956 exit(1);
5957 }
5958
5959 if (pre_args.size() != FT->getNumParams())
5960 goto badaugmentedfn;
5961
5962 for (unsigned i = 0; i < pre_args.size(); ++i) {
5963 if (pre_args[i]->getType() == FT->getParamType(i))
5964 continue;
5965 else if (!call.getCalledFunction())
5966 pre_args[i] =
5967 BuilderZ.CreateBitCast(pre_args[i], FT->getParamType(i));
5968 else
5969 goto badaugmentedfn;
5970 }
5971
5972 augmentcall = BuilderZ.CreateCall(
5973 FT, newcalled, pre_args,
5974 gutils->getInvertedBundles(&call, PreBundleTypes, BuilderZ,
5975 /*lookup*/ false));
5976 augmentcall->setCallingConv(call.getCallingConv());
5977 augmentcall->setDebugLoc(
5978 gutils->getNewFromOriginal(call.getDebugLoc()));
5979
5980 for (auto pair : preByVal) {
5981 augmentcall->addParamAttr(
5982 pair.first, Attribute::getWithByValType(augmentcall->getContext(),
5983 pair.second));
5984 }
5985
5986 for (auto &pair : structAttrs) {
5987 for (auto val : pair.second)
5988 augmentcall->addParamAttr(pair.first, val);
5989 }
5990
5991 if (!augmentcall->getType()->isVoidTy())
5992 augmentcall->setName(call.getName() + "_augmented");
5993
5994 if (tapeIdx) {
5995 auto tval = *tapeIdx;
5996 tape = (tval == -1) ? augmentcall
5997 : BuilderZ.CreateExtractValue(
5998 augmentcall, {(unsigned)tval}, "subcache");
5999 if (tape->getType()->isEmptyTy()) {
6000 auto tt = tape->getType();
6001 // If the tape is returned directly, it is the augmented call
6002 // itself, which must be kept for its side effects.
6003 if (tape != augmentcall)
6004 gutils->erase(cast<Instruction>(tape));
6005 tape = UndefValue::get(tt);
6006 } else {
6007 gutils->TapesToPreventRecomputation.insert(cast<Instruction>(tape));
6008 }
6009 tape = gutils->cacheForReverse(
6010 BuilderZ, tape, getIndex(&call, CacheType::Tape, BuilderZ));
6011 }
6012
6013 if (subretused) {
6014 Value *dcall = nullptr;
6015 assert(returnIdx);
6016 assert(augmentcall);
6017 auto rval = *returnIdx;
6018 dcall = (rval < 0) ? augmentcall
6019 : BuilderZ.CreateExtractValue(augmentcall,
6020 {(unsigned)rval});
6021 gutils->originalToNewFn[&call] = dcall;
6022 gutils->newToOriginalFn.erase(newCall);
6023 gutils->newToOriginalFn[dcall] = &call;
6024
6025 assert(dcall->getType() == call.getType());
6026 assert(dcall);
6027
6028 if (!gutils->isConstantValue(&call)) {
6029 if (!call.getType()->isFPOrFPVectorTy() && TR.anyPointer(&call)) {
6030 } else if (Mode != DerivativeMode::ReverseModePrimal) {
6031 ((DiffeGradientUtils *)gutils)->differentials[dcall] =
6032 ((DiffeGradientUtils *)gutils)->differentials[newCall];
6033 ((DiffeGradientUtils *)gutils)->differentials.erase(newCall);
6034 }
6035 }
6036 assert(dcall->getType() == call.getType());
6037 gutils->replaceAWithB(newCall, dcall);
6038
6039 if (isa<Instruction>(dcall) && !isa<PHINode>(dcall)) {
6040 cast<Instruction>(dcall)->takeName(newCall);
6041 }
6042
6044 !gutils->unnecessaryIntermediates.count(&call)) {
6045
6046 std::map<UsageKey, bool> Seen =
6048 bool primalNeededInReverse = false;
6049 {
6050 auto found = gutils->knownRecomputeHeuristic.find(&call);
6051 if (found != gutils->knownRecomputeHeuristic.end()) {
6052 if (!found->second) {
6053 primalNeededInReverse = true;
6054 }
6055 }
6056 }
6057 if (!primalNeededInReverse) {
6058
6059 auto minCutMode = (Mode == DerivativeMode::ReverseModePrimal)
6061 : Mode;
6062 primalNeededInReverse =
6064 QueryType::Primal>(gutils, &call, minCutMode, Seen,
6065 oldUnreachable);
6066 }
6067 if (primalNeededInReverse)
6068 gutils->cacheForReverse(
6069 BuilderZ, dcall, getIndex(&call, CacheType::Self, BuilderZ));
6070 }
6071 BuilderZ.SetInsertPoint(newCall->getNextNode());
6072 gutils->erase(newCall);
6073 } else {
6074 BuilderZ.SetInsertPoint(BuilderZ.GetInsertPoint()->getNextNode());
6075 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
6076 gutils->originalToNewFn[&call] = augmentcall;
6077 gutils->newToOriginalFn[augmentcall] = &call;
6078 }
6079
6080 } else {
6081 if (subdata && subdata->returns.find(AugmentedStruct::Tape) ==
6082 subdata->returns.end()) {
6083 } else {
6084 // assert(!tape);
6085 // assert(subdata);
6086 if (FT) {
6087 if (!tape) {
6088 assert(tapeIdx);
6089 auto tval = *tapeIdx;
6090#if LLVM_VERSION_MAJOR >= 18
6091 auto It = BuilderZ.GetInsertPoint();
6092 It.setHeadBit(true);
6093 BuilderZ.SetInsertPoint(It);
6094#endif
6095 tape = BuilderZ.CreatePHI(
6096 (tapeIdx == -1) ? FT->getReturnType()
6097 : cast<StructType>(FT->getReturnType())
6098 ->getElementType(tval),
6099 1, "tapeArg");
6100 }
6101 tape = gutils->cacheForReverse(
6102 BuilderZ, tape, getIndex(&call, CacheType::Tape, BuilderZ));
6103 }
6104 }
6105
6106 if (subretused) {
6107 Intrinsic::ID ID = Intrinsic::not_intrinsic;
6109 QueryType::Primal>(gutils, &call, Mode, oldUnreachable) &&
6110 !gutils->unnecessaryIntermediates.count(&call)) {
6111
6112 if (!isMemFreeLibMFunction(getFuncNameFromCall(&call), &ID)) {
6113
6114#if LLVM_VERSION_MAJOR >= 18
6115 auto It = BuilderZ.GetInsertPoint();
6116 It.setHeadBit(true);
6117 BuilderZ.SetInsertPoint(It);
6118#endif
6119 auto idx = getIndex(&call, CacheType::Self, BuilderZ);
6120 if (idx == IndexMappingError) {
6121 std::string str;
6122 raw_string_ostream ss(str);
6123 ss << "Failed to compute consistent cache index for operation: "
6124 << call << "\n";
6125 if (CustomErrorHandler) {
6126 CustomErrorHandler(str.c_str(), wrap(&call),
6127 ErrorType::InternalError, nullptr, nullptr,
6128 nullptr);
6129 } else {
6130 EmitFailure("GetIndexError", call.getDebugLoc(), &call,
6131 ss.str());
6132 }
6133 } else {
6135 cachereplace = newCall;
6136 else
6137 cachereplace = BuilderZ.CreatePHI(
6138 call.getType(), 1, call.getName() + "_tmpcacheB");
6139 cachereplace =
6140 gutils->cacheForReverse(BuilderZ, cachereplace, idx);
6141 }
6142 }
6143 } else {
6144#if LLVM_VERSION_MAJOR >= 18
6145 auto It = BuilderZ.GetInsertPoint();
6146 It.setHeadBit(true);
6147 BuilderZ.SetInsertPoint(It);
6148#endif
6149 auto pn = BuilderZ.CreatePHI(
6150 call.getType(), 1, (call.getName() + "_replacementE").str());
6151 gutils->fictiousPHIs[pn] = &call;
6152 cachereplace = pn;
6153 }
6154 } else {
6155 // TODO move right after newCall for the insertion point of BuilderZ
6156
6157 BuilderZ.SetInsertPoint(BuilderZ.GetInsertPoint()->getNextNode());
6158 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
6159 }
6160 }
6161
6162 auto ifound = gutils->invertedPointers.find(&call);
6163 if (ifound != gutils->invertedPointers.end()) {
6164 auto placeholder = cast<PHINode>(&*ifound->second);
6165
6166 bool subcheck = (subretType == DIFFE_TYPE::DUP_ARG ||
6167 subretType == DIFFE_TYPE::DUP_NONEED);
6168
6169 //! We only need the shadow pointer for non-forward Mode if it is used
6170 //! in a non return setting
6171 bool hasNonReturnUse = false;
6172 for (auto use : call.users()) {
6174 !isa<ReturnInst>(use)) {
6175 hasNonReturnUse = true;
6176 }
6177 }
6178
6179 if (subcheck && hasNonReturnUse) {
6180
6181 Value *newip = nullptr;
6184
6185 if (!differetIdx) {
6186 std::string str;
6187 raw_string_ostream ss(str);
6188 ss << "Did not have return index set when differentiating "
6189 "function\n";
6190 ss << " call" << call << "\n";
6191 ss << " augmentcall" << *augmentcall << "\n";
6192 if (CustomErrorHandler) {
6193 CustomErrorHandler(str.c_str(), wrap(&call),
6194 ErrorType::InternalError, nullptr, nullptr,
6195 nullptr);
6196 } else {
6197 EmitFailure("GetIndexError", call.getDebugLoc(), &call,
6198 ss.str());
6199 }
6200 placeholder->replaceAllUsesWith(
6201 UndefValue::get(placeholder->getType()));
6202 if (placeholder == &*BuilderZ.GetInsertPoint()) {
6203 BuilderZ.SetInsertPoint(placeholder->getNextNode());
6204 }
6205 gutils->erase(placeholder);
6206 } else {
6207 auto drval = *differetIdx;
6208 newip = (drval < 0)
6209 ? augmentcall
6210 : BuilderZ.CreateExtractValue(augmentcall,
6211 {(unsigned)drval},
6212 call.getName() + "'ac");
6213 assert(newip->getType() == placeholder->getType());
6214 placeholder->replaceAllUsesWith(newip);
6215 if (placeholder == &*BuilderZ.GetInsertPoint()) {
6216 BuilderZ.SetInsertPoint(placeholder->getNextNode());
6217 }
6218 gutils->erase(placeholder);
6219 }
6220 } else {
6221 newip = placeholder;
6222 }
6223
6224 newip = gutils->cacheForReverse(
6225 BuilderZ, newip, getIndex(&call, CacheType::Shadow, BuilderZ));
6226
6227 gutils->invertedPointers.insert(std::make_pair(
6228 (const Value *)&call, InvertedPointerVH(gutils, newip)));
6229 } else {
6230 gutils->invertedPointers.erase(ifound);
6231 if (placeholder == &*BuilderZ.GetInsertPoint()) {
6232 BuilderZ.SetInsertPoint(placeholder->getNextNode());
6233 }
6234 gutils->erase(placeholder);
6235 }
6236 }
6237
6238 if (fnandtapetype && fnandtapetype->tapeType &&
6242 shouldFree()) {
6243 assert(tape);
6244 auto tapep = BuilderZ.CreatePointerCast(
6245 tape, getPointerType(
6246 fnandtapetype->tapeType,
6247 cast<PointerType>(tape->getType())->getAddressSpace()));
6248 auto truetape =
6249 BuilderZ.CreateLoad(fnandtapetype->tapeType, tapep, "tapeld");
6250 truetape->setMetadata("enzyme_mustcache",
6251 MDNode::get(truetape->getContext(), {}));
6252
6253 CreateDealloc(BuilderZ, tape);
6254 tape = truetape;
6255 }
6256 } else {
6257 auto ifound = gutils->invertedPointers.find(&call);
6258 if (ifound != gutils->invertedPointers.end()) {
6259 auto placeholder = cast<PHINode>(&*ifound->second);
6260 gutils->invertedPointers.erase(ifound);
6261 gutils->erase(placeholder);
6262 }
6263 if (/*!topLevel*/ Mode != DerivativeMode::ReverseModeCombined &&
6264 subretused && !call.doesNotAccessMemory()) {
6266 QueryType::Primal>(gutils, &call, Mode, oldUnreachable) &&
6267 !gutils->unnecessaryIntermediates.count(&call)) {
6268 assert(!replaceFunction);
6269#if LLVM_VERSION_MAJOR >= 18
6270 auto It = BuilderZ.GetInsertPoint();
6271 It.setHeadBit(true);
6272 BuilderZ.SetInsertPoint(It);
6273#endif
6274 cachereplace = BuilderZ.CreatePHI(call.getType(), 1,
6275 call.getName() + "_cachereplace2");
6276 cachereplace = gutils->cacheForReverse(
6277 BuilderZ, cachereplace,
6278 getIndex(&call, CacheType::Self, BuilderZ));
6279 } else {
6280#if LLVM_VERSION_MAJOR >= 18
6281 auto It = BuilderZ.GetInsertPoint();
6282 It.setHeadBit(true);
6283 BuilderZ.SetInsertPoint(It);
6284#endif
6285 auto pn = BuilderZ.CreatePHI(call.getType(), 1,
6286 call.getName() + "_replacementC");
6287 gutils->fictiousPHIs[pn] = &call;
6288 cachereplace = pn;
6289 }
6290 }
6291
6292 if (!subretused && !replaceFunction)
6293 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
6294 }
6295
6296 // Note here down only contains the reverse bits
6298 return;
6299 }
6300
6301 IRBuilder<> Builder2(&call);
6302 getReverseBuilder(Builder2);
6303
6304 Value *newcalled = nullptr;
6305 FunctionType *FT = nullptr;
6306
6307 DerivativeMode subMode = (replaceFunction || !modifyPrimal)
6310 if (called) {
6311 if (Mode == DerivativeMode::ReverseModeGradient && subdata) {
6312 for (size_t i = 0; i < argsInverted.size(); i++) {
6313 if (subdata->constant_args[i] == argsInverted[i])
6314 continue;
6315 assert(subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG);
6316 assert(argsInverted[i] == DIFFE_TYPE::DUP_NONEED);
6317 argsInverted[i] = DIFFE_TYPE::DUP_ARG;
6318 }
6319 }
6320
6321 newcalled = gutils->Logic.CreatePrimalAndGradient(
6322 RequestContext(&call, &Builder2),
6324 .todiff = cast<Function>(called),
6325 .retType = subretType,
6326 .constant_args = argsInverted,
6327 .subsequent_calls_may_write = subsequent_calls_may_write,
6328 .overwritten_args = overwritten_args,
6329 .returnUsed = replaceFunction && subretused,
6330 .shadowReturnUsed = shadowReturnUsed && replaceFunction,
6331 .mode = subMode,
6332 .width = gutils->getWidth(),
6333 .freeMemory = true,
6334 .AtomicAdd = gutils->AtomicAdd,
6335 .additionalType = tape ? tape->getType() : nullptr,
6336 .forceAnonymousTape = false,
6337 .typeInfo = nextTypeInfo,
6338 .runtimeActivity = gutils->runtimeActivity,
6339 .strongZero = gutils->strongZero},
6340 TR.analyzer->interprocedural, subdata);
6341 if (!newcalled)
6342 return;
6343 FT = cast<Function>(newcalled)->getFunctionType();
6344 } else {
6345
6346 assert(subMode != DerivativeMode::ReverseModeCombined);
6347
6348 auto callval = call.getCalledOperand();
6349
6350 if (gutils->isConstantValue(callval)) {
6351 std::string s;
6352 llvm::raw_string_ostream ss(s);
6353 ss << *gutils->oldFunc << "\n";
6354 ss << "in Mode: " << to_string(Mode) << "\n";
6355 ss << " orig: " << call << " callval: " << *callval << "\n";
6356 ss << " constant function being called, but active call instruction\n";
6357 auto val = EmitNoDerivativeError(ss.str(), call, gutils, Builder2);
6358 if (val)
6359 newcalled = val;
6360 else
6361 newcalled =
6362 UndefValue::get(gutils->getShadowType(callval->getType()));
6363 } else {
6364 newcalled = lookup(gutils->invertPointerM(callval, Builder2), Builder2);
6365 }
6366
6367 auto ft = call.getFunctionType();
6368
6369 auto res =
6370 getDefaultFunctionTypeForGradient(ft, /*subretType*/ subretType);
6371 // TODO Note there is empty tape added here, replace with generic
6372 res.first.push_back(getInt8PtrTy(newcalled->getContext()));
6373 FT = FunctionType::get(
6374 StructType::get(newcalled->getContext(), res.second), res.first,
6375 ft->isVarArg());
6376 auto fptype = getUnqual(FT);
6377 newcalled = Builder2.CreatePointerCast(newcalled, getUnqual(fptype));
6378 newcalled = Builder2.CreateLoad(
6379 fptype, Builder2.CreateConstGEP1_64(fptype, newcalled, 1));
6380 }
6381
6382 if (subretType == DIFFE_TYPE::OUT_DIFF) {
6383 args.push_back(diffe(&call, Builder2));
6384 }
6385
6386 if (tape) {
6387 auto ntape = gutils->lookupM(tape, Builder2);
6388 assert(ntape);
6389 assert(ntape->getType());
6390 args.push_back(ntape);
6391 }
6392
6393 assert(newcalled);
6394 assert(FT);
6395
6396 if (false) {
6397 badfn:;
6398 auto NC = dyn_cast<Function>(newcalled);
6399 llvm::errs() << *gutils->oldFunc << "\n";
6400 llvm::errs() << *gutils->newFunc << "\n";
6401 if (NC)
6402 llvm::errs() << " trying to call " << NC->getName() << " " << *FT
6403 << "\n";
6404 else
6405 llvm::errs() << " trying to call " << *newcalled << " " << *FT << "\n";
6406
6407 for (unsigned i = 0; i < args.size(); ++i) {
6408 assert(args[i]);
6409 assert(args[i]->getType());
6410 llvm::errs() << "args[" << i << "] = " << *args[i]
6411 << " FT:" << *FT->getParamType(i) << "\n";
6412 }
6413 assert(0 && "calling with wrong number of arguments");
6414 exit(1);
6415 }
6416
6417 if (args.size() != FT->getNumParams())
6418 goto badfn;
6419
6420 for (unsigned i = 0; i < args.size(); ++i) {
6421 if (args[i]->getType() == FT->getParamType(i))
6422 continue;
6423 else if (!call.getCalledFunction())
6424 args[i] = Builder2.CreateBitCast(args[i], FT->getParamType(i));
6425 else
6426 goto badfn;
6427 }
6428
6429 CallInst *diffes =
6430 Builder2.CreateCall(FT, newcalled, args,
6431 gutils->getInvertedBundles(
6432 &call, BundleTypes, Builder2, /*lookup*/ true));
6433 diffes->setCallingConv(call.getCallingConv());
6434 diffes->setDebugLoc(gutils->getNewFromOriginal(call.getDebugLoc()));
6435
6436 for (auto pair : gradByVal) {
6437 diffes->addParamAttr(pair.first, Attribute::getWithByValType(
6438 diffes->getContext(), pair.second));
6439 }
6440
6441 for (auto &pair : structAttrs) {
6442 for (auto val : pair.second)
6443 diffes->addParamAttr(pair.first, val);
6444 }
6445
6446 unsigned structidx = 0;
6447 if (replaceFunction) {
6448 if (subretused)
6449 structidx++;
6450 if (shadowReturnUsed)
6451 structidx++;
6452 }
6453
6454 for (unsigned i = 0; i < call.arg_size(); ++i) {
6455 if (argsInverted[i] == DIFFE_TYPE::OUT_DIFF) {
6456 Value *diffeadd = Builder2.CreateExtractValue(diffes, {structidx});
6457 ++structidx;
6458
6459 if (!gutils->isConstantValue(call.getArgOperand(i))) {
6460 size_t size = 1;
6461 if (call.getArgOperand(i)->getType()->isSized())
6462 size = (gutils->newFunc->getParent()
6463 ->getDataLayout()
6464 .getTypeSizeInBits(call.getArgOperand(i)->getType()) +
6465 7) /
6466 8;
6467
6468 addToDiffe(call.getArgOperand(i), diffeadd, Builder2,
6469 TR.addingType(size, call.getArgOperand(i)));
6470 }
6471 }
6472 }
6473
6474 if (diffes->getType()->isVoidTy()) {
6475 if (structidx != 0) {
6476 llvm::errs() << *gutils->oldFunc->getParent() << "\n";
6477 llvm::errs() << "diffes: " << *diffes << " structidx=" << structidx
6478 << " subretused=" << subretused
6479 << " shadowReturnUsed=" << shadowReturnUsed << "\n";
6480 }
6481 assert(structidx == 0);
6482 } else {
6483 assert(cast<StructType>(diffes->getType())->getNumElements() ==
6484 structidx);
6485 }
6486
6487 if (subretType == DIFFE_TYPE::OUT_DIFF)
6488 setDiffe(&call,
6489 Constant::getNullValue(gutils->getShadowType(call.getType())),
6490 Builder2);
6491
6492 if (replaceFunction) {
6493
6494 // if a function is replaced for joint forward/reverse, handle inverted
6495 // pointers
6496 auto ifound = gutils->invertedPointers.find(&call);
6497 if (ifound != gutils->invertedPointers.end()) {
6498 auto placeholder = cast<PHINode>(&*ifound->second);
6499 gutils->invertedPointers.erase(ifound);
6500 if (shadowReturnUsed) {
6501 dumpMap(gutils->invertedPointers);
6502 auto dretval = cast<Instruction>(
6503 Builder2.CreateExtractValue(diffes, {subretused ? 1U : 0U}));
6504 /* todo handle this case later */
6505 assert(!subretused);
6506 gutils->invertedPointers.insert(std::make_pair(
6507 (const Value *)&call, InvertedPointerVH(gutils, dretval)));
6508 }
6509 gutils->erase(placeholder);
6510 }
6511
6512 Instruction *retval = nullptr;
6513
6514 if (subretused) {
6515 retval = cast<Instruction>(Builder2.CreateExtractValue(diffes, {0}));
6516 if (retval) {
6517 gutils->replaceAndRemoveUnwrapCacheFor(newCall, retval);
6518 }
6519 gutils->replaceAWithB(newCall, retval, /*storeInCache*/ true);
6520 } else {
6521 eraseIfUnused(call, /*erase*/ false, /*check*/ false);
6522 }
6523
6524 SmallPtrSet<Value *, 2> postCreateSet(postCreate.begin(),
6525 postCreate.end());
6526 for (auto a : postCreate) {
6527 a->moveBefore(*Builder2.GetInsertBlock(), Builder2.GetInsertPoint());
6528 for (size_t i = 0; i < a->getNumOperands(); i++) {
6529 auto op = dyn_cast<Instruction>(a->getOperand(i));
6530 if (!op || postCreateSet.count(op))
6531 continue;
6532 if (gutils->isOriginal(op->getParent())) {
6533 IRBuilder<> BuilderA(a);
6534 a->setOperand(i, gutils->lookupM(op, BuilderA));
6535 }
6536 }
6537 }
6538
6539 gutils->originalToNewFn[&call] = retval ? retval : diffes;
6540 gutils->newToOriginalFn.erase(newCall);
6541 gutils->newToOriginalFn[retval ? retval : diffes] = &call;
6542
6543 gutils->erase(newCall);
6544
6545 return;
6546 }
6547
6548 if (cachereplace) {
6549 if (subretused) {
6550 Value *dcall = nullptr;
6551 assert(cachereplace->getType() == call.getType());
6552 assert(dcall == nullptr);
6553 dcall = cachereplace;
6554 assert(dcall);
6555
6556 if (!gutils->isConstantValue(&call)) {
6557 gutils->originalToNewFn[&call] = dcall;
6558 gutils->newToOriginalFn.erase(newCall);
6559 gutils->newToOriginalFn[dcall] = &call;
6560 if (!call.getType()->isFPOrFPVectorTy() && TR.anyPointer(&call)) {
6561 } else {
6562 ((DiffeGradientUtils *)gutils)->differentials[dcall] =
6563 ((DiffeGradientUtils *)gutils)->differentials[newCall];
6564 ((DiffeGradientUtils *)gutils)->differentials.erase(newCall);
6565 }
6566 }
6567 assert(dcall->getType() == call.getType());
6568 newCall->replaceAllUsesWith(dcall);
6569 if (isa<Instruction>(dcall) && !isa<PHINode>(dcall)) {
6570 cast<Instruction>(dcall)->takeName(&call);
6571 }
6572 gutils->erase(newCall);
6573 } else {
6574 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
6575 if (augmentcall) {
6576 gutils->originalToNewFn[&call] = augmentcall;
6577 gutils->newToOriginalFn.erase(newCall);
6578 gutils->newToOriginalFn[augmentcall] = &call;
6579 }
6580 }
6581 }
6582 return;
6583 }
6584
6585 void handleMPI(llvm::CallInst &call, llvm::Function *called,
6586 llvm::StringRef funcName);
6587
6588 bool handleKnownCallDerivatives(llvm::CallInst &call, llvm::Function *called,
6589 llvm::StringRef funcName,
6590 bool subsequent_calls_may_write,
6591 const std::vector<bool> &overwritten_args,
6592 llvm::CallInst *const newCall);
6593
6594 // Return
6595 void visitCallInst(llvm::CallInst &call) {
6596 using namespace llvm;
6597
6598 StringRef funcName = getFuncNameFromCall(&call);
6599
6600 // When compiling Enzyme against standard LLVM, and not Intel's
6601 // modified version of LLVM, the intrinsic `llvm.intel.subscript` is
6602 // not fully understood by LLVM. One of the results of this is that the
6603 // visitor dispatches to visitCallInst, rather than visitIntrinsicInst, when
6604 // presented with the intrinsic - hence why we are handling it here.
6605 if (startsWith(funcName, ("llvm.intel.subscript"))) {
6606 assert(isa<IntrinsicInst>(call));
6607 visitIntrinsicInst(cast<IntrinsicInst>(call));
6608 return;
6609 }
6610
6611 if (funcName == "llvm.enzyme.lifetime_start") {
6612 visitIntrinsicInst(cast<IntrinsicInst>(call));
6613 return;
6614 }
6615 if (funcName == "llvm.enzyme.lifetime_end") {
6616 SmallVector<Value *, 2> orig_ops(call.getNumOperands());
6617 for (unsigned i = 0; i < call.getNumOperands(); ++i) {
6618 orig_ops[i] = call.getOperand(i);
6619 }
6620 handleAdjointForIntrinsic(Intrinsic::lifetime_end, call, orig_ops);
6621 eraseIfUnused(call);
6622 return;
6623 }
6624
6625 CallInst *const newCall = cast<CallInst>(gutils->getNewFromOriginal(&call));
6626 IRBuilder<> BuilderZ(newCall);
6627 BuilderZ.setFastMathFlags(getFast());
6628
6629 if (overwritten_args_map.find(&call) == overwritten_args_map.end() &&
6632 llvm::errs() << " call: " << call << "\n";
6633 for (auto &pair : overwritten_args_map) {
6634 llvm::errs() << " + " << *pair.first << "\n";
6635 }
6636 }
6637
6638 assert(overwritten_args_map.find(&call) != overwritten_args_map.end() ||
6641 const bool subsequent_calls_may_write =
6642 (Mode == DerivativeMode::ForwardMode ||
6644 ? false
6645 : overwritten_args_map.find(&call)->second.first;
6646 const std::vector<bool> &overwritten_args =
6647 (Mode == DerivativeMode::ForwardMode ||
6649 ? std::vector<bool>()
6650 : overwritten_args_map.find(&call)->second.second;
6651
6652 auto called = getFunctionFromCall(&call);
6653
6654 bool subretused = false;
6655 bool shadowReturnUsed = false;
6656 auto smode = Mode;
6659 DIFFE_TYPE subretType = gutils->getReturnDiffeType(
6660 &call, &subretused, &shadowReturnUsed, smode);
6661
6662 if (Mode == DerivativeMode::ForwardMode ||
6664 auto found = customFwdCallHandlers.find(funcName);
6665 if (found != customFwdCallHandlers.end()) {
6666 Value *invertedReturn = nullptr;
6667 auto ifound = gutils->invertedPointers.find(&call);
6668 if (ifound != gutils->invertedPointers.end()) {
6669 invertedReturn = cast<PHINode>(&*ifound->second);
6670 }
6671
6672 Value *normalReturn = subretused ? newCall : nullptr;
6673
6674 bool noMod = found->second(BuilderZ, &call, *gutils, normalReturn,
6675 invertedReturn);
6676 if (noMod) {
6677 if (subretused)
6678 assert(normalReturn == newCall);
6679 eraseIfUnused(call);
6680 }
6681
6682 ifound = gutils->invertedPointers.find(&call);
6683 if (ifound != gutils->invertedPointers.end()) {
6684 auto placeholder = cast<PHINode>(&*ifound->second);
6685 if (invertedReturn && invertedReturn != placeholder) {
6686 if (invertedReturn->getType() !=
6687 gutils->getShadowType(call.getType())) {
6688 llvm::errs() << " o: " << call << "\n";
6689 llvm::errs() << " ot: " << *call.getType() << "\n";
6690 llvm::errs() << " ir: " << *invertedReturn << "\n";
6691 llvm::errs() << " irt: " << *invertedReturn->getType() << "\n";
6692 llvm::errs() << " p: " << *placeholder << "\n";
6693 llvm::errs() << " PT: " << *placeholder->getType() << "\n";
6694 llvm::errs() << " newCall: " << *newCall << "\n";
6695 llvm::errs() << " newCallT: " << *newCall->getType() << "\n";
6696 }
6697 assert(invertedReturn->getType() ==
6698 gutils->getShadowType(call.getType()));
6699 placeholder->replaceAllUsesWith(invertedReturn);
6700 gutils->erase(placeholder);
6701 gutils->invertedPointers.insert(
6702 std::make_pair((const Value *)&call,
6703 InvertedPointerVH(gutils, invertedReturn)));
6704 } else {
6705 gutils->invertedPointers.erase(&call);
6706 gutils->erase(placeholder);
6707 }
6708 }
6709
6710 if (normalReturn && normalReturn != newCall) {
6711 assert(normalReturn->getType() == newCall->getType());
6712 gutils->replaceAWithB(newCall, normalReturn);
6713 gutils->erase(newCall);
6714 }
6715 return;
6716 }
6717 }
6718
6722 auto found = customCallHandlers.find(funcName);
6723 if (found != customCallHandlers.end()) {
6724 IRBuilder<> Builder2(&call);
6727 getReverseBuilder(Builder2);
6728
6729 Value *invertedReturn = nullptr;
6730 auto ifound = gutils->invertedPointers.find(&call);
6731 PHINode *placeholder = nullptr;
6732 if (ifound != gutils->invertedPointers.end()) {
6733 placeholder = cast<PHINode>(&*ifound->second);
6734 if (shadowReturnUsed)
6735 invertedReturn = placeholder;
6736 }
6737
6738 Value *normalReturn = subretused ? newCall : nullptr;
6739
6740 Value *tape = nullptr;
6741
6742 Type *tapeType = nullptr;
6743
6746 bool noMod = found->second.first(BuilderZ, &call, *gutils,
6747 normalReturn, invertedReturn, tape);
6748 if (noMod) {
6749 if (subretused)
6750 assert(normalReturn == newCall);
6751 eraseIfUnused(call);
6752 }
6753 if (tape) {
6754 tapeType = tape->getType();
6755 gutils->cacheForReverse(BuilderZ, tape,
6756 getIndex(&call, CacheType::Tape, BuilderZ));
6757 }
6759 assert(augmentedReturn);
6760 auto subaugmentations =
6761 (std::map<const llvm::CallInst *, AugmentedReturn *>
6762 *)&augmentedReturn->subaugmentations;
6764 *subaugmentations, &call, (AugmentedReturn *)tapeType);
6765 }
6766 }
6767
6771 augmentedReturn->tapeIndices.find(
6772 std::make_pair(&call, CacheType::Tape)) !=
6773 augmentedReturn->tapeIndices.end()) {
6774 assert(augmentedReturn);
6775 auto subaugmentations =
6776 (std::map<const llvm::CallInst *, AugmentedReturn *>
6777 *)&augmentedReturn->subaugmentations;
6778 auto fd = subaugmentations->find(&call);
6779 assert(fd != subaugmentations->end());
6780 // Note we are using the storage space here to persist
6781 // the LLVM type, as storing a new augmentedReturn has issues
6782 // regarding persisting the data structure, and when it will
6783 // be freed, since it will no longer live in the map in
6784 // EnzymeLogic.
6785 tapeType = (llvm::Type *)fd->second;
6786
6787#if LLVM_VERSION_MAJOR >= 18
6788 auto It = BuilderZ.GetInsertPoint();
6789 It.setHeadBit(true);
6790 BuilderZ.SetInsertPoint(It);
6791#endif
6792 tape = BuilderZ.CreatePHI(tapeType, 0);
6793 tape = gutils->cacheForReverse(
6794 BuilderZ, tape, getIndex(&call, CacheType::Tape, BuilderZ),
6795 /*ignoreType*/ true);
6796 }
6797 if (tape)
6798 tape = gutils->lookupM(tape, Builder2);
6799 found->second.second(Builder2, &call, *(DiffeGradientUtils *)gutils,
6800 tape);
6801 }
6802
6803 if (placeholder) {
6804 if (!shadowReturnUsed) {
6805 gutils->invertedPointers.erase(&call);
6806 gutils->erase(placeholder);
6807 } else {
6808 if (invertedReturn && invertedReturn != placeholder) {
6809 if (invertedReturn->getType() !=
6810 gutils->getShadowType(call.getType())) {
6811 llvm::errs() << " o: " << call << "\n";
6812 llvm::errs() << " ot: " << *call.getType() << "\n";
6813 llvm::errs() << " ir: " << *invertedReturn << "\n";
6814 llvm::errs() << " irt: " << *invertedReturn->getType() << "\n";
6815 llvm::errs() << " p: " << *placeholder << "\n";
6816 llvm::errs() << " PT: " << *placeholder->getType() << "\n";
6817 llvm::errs() << " newCall: " << *newCall << "\n";
6818 llvm::errs() << " newCallT: " << *newCall->getType() << "\n";
6819 }
6820 assert(invertedReturn->getType() ==
6821 gutils->getShadowType(call.getType()));
6822 placeholder->replaceAllUsesWith(invertedReturn);
6823 gutils->erase(placeholder);
6824 invertedReturn = gutils->cacheForReverse(
6825 BuilderZ, invertedReturn,
6826 getIndex(&call, CacheType::Shadow, BuilderZ));
6827 } else {
6828 auto idx = getIndex(&call, CacheType::Shadow, BuilderZ);
6829 invertedReturn =
6830 gutils->cacheForReverse(BuilderZ, placeholder, idx);
6831 if (idx == IndexMappingError) {
6832 if (placeholder->getType() != invertedReturn->getType())
6833 llvm::errs() << " place: " << *placeholder
6834 << " invRet: " << *invertedReturn;
6835 placeholder->replaceAllUsesWith(invertedReturn);
6836 gutils->erase(placeholder);
6837 }
6838 }
6839
6840 gutils->invertedPointers.insert(
6841 std::make_pair((const Value *)&call,
6842 InvertedPointerVH(gutils, invertedReturn)));
6843 }
6844 }
6845
6846 bool primalNeededInReverse;
6847
6848 if (gutils->knownRecomputeHeuristic.count(&call)) {
6849 primalNeededInReverse = !gutils->knownRecomputeHeuristic[&call];
6850 } else {
6851 std::map<UsageKey, bool> Seen =
6853 primalNeededInReverse =
6855 QueryType::Primal>(gutils, &call, Mode, Seen, oldUnreachable);
6856 }
6857 if (subretused && primalNeededInReverse) {
6858 if (normalReturn != newCall) {
6859 assert(normalReturn->getType() == newCall->getType());
6860 gutils->replaceAWithB(newCall, normalReturn);
6861 BuilderZ.SetInsertPoint(newCall->getNextNode());
6862 gutils->erase(newCall);
6863 }
6864 normalReturn = gutils->cacheForReverse(
6865 BuilderZ, normalReturn,
6866 getIndex(&call, CacheType::Self, BuilderZ));
6867 } else {
6868 if (normalReturn && normalReturn != newCall) {
6869 assert(normalReturn->getType() == newCall->getType());
6871 gutils->replaceAWithB(newCall, normalReturn);
6872 BuilderZ.SetInsertPoint(newCall->getNextNode());
6873 gutils->erase(newCall);
6874 } else if (Mode == DerivativeMode::ReverseModeGradient)
6875 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
6876 }
6877 return;
6878 }
6879 }
6880
6881 if (called) {
6882 if (funcName == "__kmpc_fork_call") {
6883 visitOMPCall(call);
6884 return;
6885 }
6886 }
6887
6888 if (handleKnownCallDerivatives(call, called, funcName,
6889 subsequent_calls_may_write, overwritten_args,
6890 newCall))
6891 return;
6892
6893 bool useConstantFallback =
6895 gutils, call, QueryType::Primal, nullptr);
6896 if (!useConstantFallback) {
6897 if (gutils->isConstantInstruction(&call) &&
6898 gutils->isConstantValue(&call)) {
6899 EmitWarning("ConstnatFallback", call,
6900 "Call was deduced inactive but still doing differential "
6901 "rewrite as it may escape an allocation",
6902 call);
6903 }
6904 }
6905 if (useConstantFallback) {
6906 if (!gutils->isConstantValue(&call)) {
6907 auto found = gutils->invertedPointers.find(&call);
6908 if (found != gutils->invertedPointers.end()) {
6909 PHINode *placeholder = cast<PHINode>(&*found->second);
6910 gutils->invertedPointers.erase(found);
6911 gutils->erase(placeholder);
6912 }
6913 }
6914 bool noFree = Mode == DerivativeMode::ForwardMode ||
6916 noFree |= call.hasFnAttr(Attribute::NoFree);
6917 if (!noFree && called) {
6918 noFree |= called->hasFnAttribute(Attribute::NoFree);
6919 }
6920
6921 std::map<UsageKey, bool> CacheResults =
6923
6924 if (!noFree && !EnzymeGlobalActivity) {
6925 bool mayActiveFree = false;
6926 for (unsigned i = 0; i < call.arg_size(); ++i) {
6927 Value *a = call.getOperand(i);
6928
6929 if (EnzymeJuliaAddrLoad && isSpecialPtr(a->getType()))
6930 continue;
6931 // if could not be a pointer, it cannot be freed
6932 if (!TR.query(a)[{-1}].isPossiblePointer())
6933 continue;
6934 // if active value, we need to do memory preservation
6935 if (!gutils->isConstantValue(a)) {
6936 mayActiveFree = true;
6937 break;
6938 }
6939 // if used in reverse (even if just primal), need to do
6940 // memory preservation
6941 const auto obj = getBaseObject(a);
6942 // If not allocation/allocainst, it is possible this aliases
6943 // a pointer needed in the reverse pass
6944 bool isAllocation = false;
6945 for (auto objv = obj;;) {
6946 if (isAllocationCall(objv, gutils->TLI)) {
6947 isAllocation = true;
6948 break;
6949 }
6950 if (auto objC = dyn_cast<CallBase>(objv))
6951 if (auto F = getFunctionFromCall(objC))
6952 if (!F->empty()) {
6953 SmallPtrSet<Value *, 1> set;
6954 for (auto &B : *F) {
6955 if (auto RI = dyn_cast<ReturnInst>(B.getTerminator())) {
6956 auto v = getBaseObject(RI->getOperand(0));
6957 if (isa<ConstantPointerNull>(v))
6958 continue;
6959 set.insert(v);
6960 }
6961 }
6962 if (set.size() == 1) {
6963 objv = *set.begin();
6964 continue;
6965 }
6966 }
6967 break;
6968 }
6969 if (!isAllocation) {
6970 mayActiveFree = true;
6971 break;
6972 }
6973 {
6974 auto found = gutils->knownRecomputeHeuristic.find(obj);
6975 if (found != gutils->knownRecomputeHeuristic.end()) {
6976 if (!found->second) {
6977 auto CacheResults2(CacheResults);
6978 CacheResults2.erase(UsageKey(obj, QueryType::Primal));
6980 QueryType::Primal>(gutils, obj,
6982 CacheResults2, oldUnreachable)) {
6983 mayActiveFree = true;
6984 break;
6985 }
6986 }
6987 continue;
6988 }
6989 }
6990 auto CacheResults2(CacheResults);
6992 QueryType::Primal>(gutils, obj,
6994 CacheResults2, oldUnreachable)) {
6995 mayActiveFree = true;
6996 break;
6997 }
6998 }
6999 if (!mayActiveFree)
7000 noFree = true;
7001 }
7002 if (!noFree) {
7003 auto callval = call.getCalledOperand();
7004 if (!isa<Constant>(callval))
7005 callval = gutils->getNewFromOriginal(callval);
7006 newCall->setCalledOperand(gutils->Logic.CreateNoFree(
7007 RequestContext(&call, &BuilderZ), callval));
7008 }
7009 if (gutils->knownRecomputeHeuristic.find(&call) !=
7010 gutils->knownRecomputeHeuristic.end()) {
7011 if (!gutils->knownRecomputeHeuristic[&call]) {
7012 gutils->cacheForReverse(BuilderZ, newCall,
7013 getIndex(&call, CacheType::Self, BuilderZ));
7014 eraseIfUnused(call);
7015 return;
7016 }
7017 }
7018
7019 // If we need this value and it is illegal to recompute it (it writes or
7020 // may load overwritten data)
7021 // Store and reload it
7024 Mode != DerivativeMode::ForwardModeError && subretused &&
7025 (call.mayWriteToMemory() ||
7026 !gutils->legalRecompute(&call, ValueToValueMapTy(), nullptr))) {
7027 if (!gutils->unnecessaryIntermediates.count(&call)) {
7028
7029 std::map<UsageKey, bool> Seen =
7031 bool primalNeededInReverse = false;
7032 {
7033 auto found = gutils->knownRecomputeHeuristic.find(&call);
7034 if (found != gutils->knownRecomputeHeuristic.end()) {
7035 if (!found->second) {
7036 primalNeededInReverse = true;
7037 }
7038 }
7039 }
7040 if (!primalNeededInReverse) {
7041
7042 auto minCutMode = (Mode == DerivativeMode::ReverseModePrimal)
7044 : Mode;
7045 primalNeededInReverse =
7047 QueryType::Primal>(gutils, &call, minCutMode, Seen,
7048 oldUnreachable);
7049 }
7050 if (primalNeededInReverse) {
7051 gutils->cacheForReverse(BuilderZ, newCall,
7052 getIndex(&call, CacheType::Self, BuilderZ));
7053 eraseIfUnused(call);
7054 return;
7055 }
7056 }
7057 // Force erasure in reverse pass, since cached if needed
7060 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
7061 else
7062 eraseIfUnused(call);
7063 return;
7064 }
7065
7066 // If this call may write to memory and is a copy (in the just reverse
7067 // pass), erase it
7068 // Any uses of it should be handled by the case above so it is safe to
7069 // RAUW
7070 if (call.mayWriteToMemory() &&
7073 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
7074 return;
7075 }
7076
7077 // if call does not write memory and isn't used, we can erase it
7078 if (!call.mayWriteToMemory() && !subretused) {
7079 eraseIfUnused(call, /*erase*/ true, /*check*/ false);
7080 return;
7081 }
7082
7083 return;
7084 }
7085
7087 call, called, subsequent_calls_may_write, overwritten_args,
7088 shadowReturnUsed, subretType, subretused);
7089 }
7090};
7091
7092#endif // ENZYME_ADJOINT_GENERATOR_H
static bool isZero(llvm::Constant *cst)
DIFFE_TYPE
Potential differentiable argument classifications.
FunctionType * getFunctionTypeForClone(llvm::FunctionType *FTy, DerivativeMode mode, unsigned width, llvm::Type *additionalArg, llvm::ArrayRef< DIFFE_TYPE > constant_args, bool diffeReturnArg, bool returnTape, bool returnPrimal, bool returnShadow)
static bool isAllocationCall(const llvm::Value *TmpOrig, llvm::TargetLibraryInfo &TLI)
static bool isReadOnly(Operation *op)
static Operation * getFunctionFromCall(CallOpInterface iface)
llvm::cl::opt< bool > EnzymeGlobalActivity
constexpr const char * to_string(ActivityAnalyzer::UseActivity UA)
std::pair< SmallVector< Type *, 4 >, SmallVector< Type *, 4 > > getDefaultFunctionTypeForAugmentation(FunctionType *called, bool returnUsed, DIFFE_TYPE retType)
assuming not top level
bool legalCombinedForwardReverse(CallInst *origop, const std::map< ReturnInst *, StoreInst * > &replacedReturns, SmallVectorImpl< Instruction * > &postCreate, SmallVectorImpl< Instruction * > &userReplace, const GradientUtils *gutils, const SmallPtrSetImpl< const Instruction * > &unnecessaryInstructions, const SmallPtrSetImpl< BasicBlock * > &oldUnreachable, const bool subretused)
bool shouldAugmentCall(CallInst *op, const GradientUtils *gutils)
std::pair< SmallVector< Type *, 4 >, SmallVector< Type *, 4 > > getDefaultFunctionTypeForGradient(FunctionType *called, DIFFE_TYPE retType, ArrayRef< DIFFE_TYPE > tys)
assuming not top level
CacheType
Definition EnzymeLogic.h:80
static std::string str(AugmentedStruct c)
Definition EnzymeLogic.h:62
llvm::cl::opt< bool > looseTypeAnalysis
llvm::cl::opt< bool > nonmarkedglobals_inactiveloads
StringMap< std::pair< std::function< bool(IRBuilder<> &, CallInst *, GradientUtils &, Value *&, Value *&, Value *&)>, std::function< void(IRBuilder<> &, CallInst *, DiffeGradientUtils &, Value *)> > > customCallHandlers
void SubTransferHelper(GradientUtils *gutils, DerivativeMode mode, Type *secretty, Intrinsic::ID intrinsic, unsigned dstalign, unsigned srcalign, unsigned offset, bool dstConstant, Value *shadow_dst, Value *primal_dst, bool srcConstant, Value *shadow_src, Value *primal_src, Value *length, Value *isVolatile, llvm::CallInst *MTI, bool allowForward, bool shadowsLookedUp, bool backwardsShadow)
SmallVector< unsigned int, 9 > MD_ToCopy
StringMap< std::function< bool(IRBuilder<> &, CallInst *, GradientUtils &, Value *&, Value *&)> > customFwdCallHandlers
constexpr int IndexMappingError
std::pair< const llvm::Value *, QueryType > UsageKey
static TypeTree parseTBAA(TBAAStructTypeNode AccessType, llvm::Instruction &I, const llvm::DataLayout &DL, std::shared_ptr< llvm::ModuleSlotTracker > MST)
Given a TBAA access node return the corresponding TypeTree This includes recursively parsing the acce...
Definition TBAA.h:439
TypeTree defaultTypeTreeForLLVM(llvm::Type *ET, llvm::Instruction *I, bool intIsPointer)
static bool isMemFreeLibMFunction(llvm::StringRef str, llvm::Intrinsic::ID *ID=nullptr)
CallInst * CreateDealloc(llvm::IRBuilder<> &Builder, llvm::Value *ToFree)
Definition Utils.cpp:871
void EmitNoTypeError(const std::string &message, llvm::Instruction &inst, GradientUtils *gutils, llvm::IRBuilder<> &Builder2)
Definition Utils.cpp:4675
LLVMValueRef(* CustomErrorHandler)(const char *, LLVMValueRef, ErrorType, const void *, LLVMValueRef, LLVMBuilderRef)
Definition Utils.cpp:62
llvm::Value * EmitNoDerivativeError(const std::string &message, llvm::Instruction &inst, GradientUtils *gutils, llvm::IRBuilder<> &Builder2, llvm::Value *condition)
Definition Utils.cpp:4593
llvm::MDNode * sliceTrueType(const llvm::MDNode *md, size_t start, size_t length)
Restrict an enzyme_truetype annotation describing a memory region to the byte range [start,...
Definition Utils.cpp:4713
void ErrorIfRuntimeInactive(llvm::IRBuilder<> &B, llvm::Value *primal, llvm::Value *shadow, const char *Message, llvm::DebugLoc &&loc, llvm::Instruction *orig)
Definition Utils.cpp:1031
Function * getOrInsertDifferentialFloatMemcpy(Module &M, Type *elementType, unsigned dstalign, unsigned srcalign, unsigned dstaddr, unsigned srcaddr, unsigned bitwidth, bool runtimeActivity, bool atomic)
Create function for type that is equivalent to memcpy but adds to destination rather than a direct co...
Definition Utils.cpp:1160
llvm::FastMathFlags getFast()
Get LLVM fast math flags.
Definition Utils.cpp:4026
llvm::Constant * getUndefinedValueForType(llvm::Module &M, llvm::Type *T, bool forceZero)
Definition Utils.cpp:3918
std::vector< std::tuple< llvm::Type *, size_t, size_t > > parseTrueType(const llvm::MDNode *md, DerivativeMode Mode, bool const_src)
Definition Utils.cpp:4749
static llvm::Value * checkedMul(bool strongZero, llvm::IRBuilder<> &Builder2, llvm::Value *idiff, llvm::Value *pres, const llvm::Twine &Name="")
Definition Utils.h:2296
static bool isIntelSubscriptIntrinsic(const llvm::IntrinsicInst &II)
Definition Utils.h:1649
static std::string getRenamedPerCallingConv(llvm::StringRef caller, llvm::StringRef callee)
Definition Utils.h:2812
static llvm::PointerType * getPointerType(llvm::Type *T, unsigned AddressSpace=0)
Pointer into the given address space.
Definition Utils.h:1313
static void dumpMap(const llvm::ValueMap< T, N > &o, llvm::function_ref< bool(const llvm::Value *)> shouldPrint=[](T) { return true;})
Print a map, optionally with a shouldPrint function to decide to print a given value.
Definition Utils.h:430
@ Args
Return is a struct of all args.
static bool startsWith(llvm::StringRef string, llvm::StringRef prefix)
Definition Utils.h:847
static llvm::PointerType * getUnqual(llvm::Type *T)
Definition Utils.h:1329
void EmitFailure(llvm::StringRef RemarkName, const llvm::DiagnosticLocation &Loc, const llvm::Instruction *CodeRegion, Args &...args)
Definition Utils.h:347
static llvm::PointerType * getInt8PtrTy(llvm::LLVMContext &Context, unsigned AddressSpace=0)
Definition Utils.h:1324
static std::map< K, V >::iterator insert_or_assign2(std::map< K, V > &map, K key, V val)
Insert into a map.
Definition Utils.h:990
static llvm::Function * getIntrinsicDeclaration(llvm::Module *M, llvm::Intrinsic::ID id, llvm::ArrayRef< llvm::Type * > Tys={})
Definition Utils.h:2541
static bool isNoCapture(const llvm::CallBase *call, size_t idx)
Definition Utils.h:2109
static void addFunctionNoCapture(llvm::Function *call, size_t idx)
Definition Utils.h:2609
static llvm::Value * getBaseObject(llvm::Value *V, bool offsetAllowed=true)
Definition Utils.h:1711
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
static llvm::MDNode * hasMetadata(const llvm::GlobalObject *O, llvm::StringRef kind)
Check if a global has metadata.
Definition Utils.h:483
static bool containsOnlyAtMostTopBit(const llvm::Value *V, llvm::Type *FT, const llvm::DataLayout &dl, llvm::Type **vFT=nullptr)
Definition Utils.h:2325
static bool shouldDisableNoWrite(const llvm::CallInst *CI)
Definition Utils.h:1633
void EmitWarning(llvm::StringRef RemarkName, const llvm::DiagnosticLocation &Loc, const llvm::BasicBlock *BB, const Args &...args)
Definition Utils.h:252
static bool isSpecialPtr(llvm::Type *Ty)
Definition Utils.h:2664
llvm::cl::opt< bool > EnzymeJuliaAddrLoad
@ MixedActivityError
static llvm::Type * IntToFloatTy(llvm::Type *T)
Convert a integer type to a floating point type of the same size.
Definition Utils.h:799
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
static bool EmitWarningEnabled(const llvm::LLVMContext &Ctx)
Definition Utils.h:246
static llvm::Value * CreateSelect(llvm::IRBuilder<> &Builder2, llvm::Value *cmp, llvm::Value *tval, llvm::Value *fval, const llvm::Twine &Name="")
Definition Utils.h:2283
static std::string convertSRetTypeToString(llvm::Type *T)
Definition Utils.h:2853
static llvm::StringRef getFuncNameFromCall(const llvm::CallBase *op)
Definition Utils.h:1470
static constexpr const char * MixedActivityHint
Remedy appended to mismatched-activity warnings, which are emitted when a value Enzyme proved inactiv...
Definition Utils.h:315
static bool isWriteOnly(const llvm::Function *F, ssize_t arg=-1)
Definition Utils.h:2033
ValueType
Classification of value as an original program variable, a derivative variable, neither,...
Definition Utils.h:543
static llvm::Attribute::AttrKind PrimalParamAttrsToPreserve[]
Definition Utils.h:2486
static llvm::Attribute::AttrKind ShadowParamAttrsToPreserve[]
Definition Utils.h:2518
DerivativeMode
Definition Utils.h:524
static DIFFE_TYPE whatType(llvm::Type *arg, DerivativeMode mode, bool integersAreConstant, std::set< llvm::Type * > &seen)
Attempt to automatically detect the differentiable classification based off of a given type.
Definition Utils.h:653
llvm::Value * lookup(llvm::Value *val, llvm::IRBuilder<> &Builder)
void visitExtractElementInst(llvm::ExtractElementInst &EEI)
void forwardModeInvertedPointerFallback(llvm::Instruction &I)
llvm::Value * diffe(llvm::Value *val, llvm::IRBuilder<> &Builder)
void visitShuffleVectorInst(llvm::ShuffleVectorInst &SVI)
void createSelectInstAdjoint(llvm::SelectInst &SI)
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)
void applyChainRule(llvm::ArrayRef< llvm::Value * > diffs, llvm::IRBuilder<> &Builder, Func rule)
Unwraps an collection of constant vector derivatives from their internal representations and applies ...
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)
void setDiffe(llvm::Value *val, llvm::Value *dif, llvm::IRBuilder<> &Builder)
void visitBinaryOperator(llvm::BinaryOperator &BO)
void visitInsertValueInst(llvm::InsertValueInst &IVI)
void visitCallInst(llvm::CallInst &call)
void visitGetElementPtrInst(llvm::GetElementPtrInst &gep)
llvm::Value * MPI_TYPE_SIZE(llvm::Value *DT, llvm::IRBuilder<> &B, llvm::Type *intType, llvm::Function *caller)
void visitFCmpInst(llvm::FCmpInst &I)
void visitExtractValueInst(llvm::ExtractValueInst &EVI)
void visitLoadLike(llvm::Instruction &I, llvm::MaybeAlign alignment, bool constantval, llvm::Value *mask=nullptr, llvm::Value *orig_maskInit=nullptr)
static llvm::AtomicOrdering reverseLoadOrdering(llvm::AtomicOrdering order)
The ordering with which the reverse pass may read a location that the primal modified atomically with...
void visitCastInst(llvm::CastInst &I)
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...
void visitInsertElementInst(llvm::InsertElementInst &IEI)
void visitCommonStore(llvm::Instruction &I, llvm::Value *orig_ptr, llvm::Value *orig_val, llvm::MaybeAlign prevalign, bool isVolatile, llvm::AtomicOrdering ordering, llvm::SyncScope::ID syncScope, llvm::Value *mask)
void visitICmpInst(llvm::ICmpInst &I)
void visitFenceInst(llvm::FenceInst &FI)
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...
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...
void visitIntrinsicInst(llvm::IntrinsicInst &II)
void createBinaryOperatorAdjoint(llvm::BinaryOperator &BO)
void visitPHINode(llvm::PHINode &phi)
llvm::SmallVector< llvm::SelectInst *, 4 > addToDiffe(llvm::Value *val, llvm::Value *dif, llvm::IRBuilder<> &Builder, llvm::Type *T, llvm::Value *mask=nullptr)
void eraseIfUnused(llvm::Instruction &I, bool erase=true, bool check=true)
AdjointGenerator(DerivativeMode Mode, GradientUtils *gutils, llvm::ArrayRef< DIFFE_TYPE > constant_args, DIFFE_TYPE retType, std::function< unsigned(llvm::Instruction *, CacheType, llvm::IRBuilder<> &)> getIndex, const std::map< llvm::CallInst *, std::pair< bool, const std::vector< bool > > > overwritten_args_map, const AugmentedReturn *augmentedReturn, const std::map< llvm::ReturnInst *, llvm::StoreInst * > *replacedReturns, const llvm::SmallPtrSetImpl< const llvm::Value * > &unnecessaryValues, const llvm::SmallPtrSetImpl< const llvm::Instruction * > &unnecessaryInstructions, const llvm::SmallPtrSetImpl< const llvm::Instruction * > &unnecessaryStores, const llvm::SmallPtrSetImpl< llvm::BasicBlock * > &oldUnreachable)
void visitInstruction(llvm::Instruction &inst)
void visitAtomicRMWInst(llvm::AtomicRMWInst &I)
void visitLoadInst(llvm::LoadInst &LI)
void visitMemSetInst(llvm::MemSetInst &MS)
void visitMemTransferInst(llvm::MemTransferInst &MTI)
void visitOMPCall(llvm::CallInst &call)
void visitSelectInst(llvm::SelectInst &SI)
void createBinaryOperatorDual(llvm::BinaryOperator &BO)
void getForwardBuilder(llvm::IRBuilder<> &Builder2)
void getReverseBuilder(llvm::IRBuilder<> &Builder2, bool original=true)
void applyChainRule(llvm::IRBuilder<> &Builder, Func rule, Args... args)
Unwraps a vector derivative from its internal representation and applies a function f to each element...
void visitAllocaInst(llvm::AllocaInst &I)
void recursivelyHandleSubfunction(llvm::CallInst &call, llvm::Function *called, bool subsequent_calls_may_write, const std::vector< bool > &overwritten_args, bool shadowReturnUsed, DIFFE_TYPE subretType, bool subretused)
llvm::Value * MPI_COMM_SIZE(llvm::Value *comm, llvm::IRBuilder<> &B, llvm::Type *rankTy, llvm::Function *caller)
void visitMemTransferCommon(llvm::Intrinsic::ID ID, llvm::MaybeAlign srcAlign, llvm::MaybeAlign dstAlign, llvm::CallInst &MTI, llvm::Value *orig_dst, llvm::Value *orig_src, llvm::Value *new_size, llvm::Value *isVolatile)
llvm::Value * MPI_COMM_RANK(llvm::Value *comm, llvm::IRBuilder<> &B, llvm::Type *rankTy, llvm::Function *caller)
bool handleAdjointForIntrinsic(llvm::Intrinsic::ID ID, llvm::Instruction &I, llvm::SmallVectorImpl< llvm::Value * > &orig_ops)
void handleMPI(llvm::CallInst &call, llvm::Function *called, llvm::StringRef funcName)
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...
void visitStoreInst(llvm::StoreInst &SI)
void visitMemSetCommon(llvm::CallInst &MS)
return structtype if recursive function
llvm::Function * fn
const std::vector< DIFFE_TYPE > constant_args
std::set< ssize_t > tapeIndiciesToFree
std::map< const llvm::CallInst *, const AugmentedReturn * > subaugmentations
Map from original call to sub augmentation data.
std::map< std::pair< llvm::Instruction *, CacheType >, int > tapeIndices
llvm::Type * tapeType
return structtype if recursive function
std::map< AugmentedStruct, int > returns
Map from information desired from a augmented return to its index in the returned struct.
llvm::Function *const newFunc
The function whose instructions we are caching.
llvm::TargetLibraryInfo & TLI
Various analysis results of newFunc.
llvm::BasicBlock * inversionAllocs
Concrete SubType of a given value.
bool isKnown() const
Whether this ConcreteType has information (is not unknown)
llvm::Type * isFloat() const
Return the floating point type, if this is a float.
llvm::Function * CreateNoFree(RequestContext context, llvm::Function *todiff)
const AugmentedReturn & CreateAugmentedPrimal(RequestContext context, llvm::Function *todiff, DIFFE_TYPE retType, llvm::ArrayRef< DIFFE_TYPE > constant_args, TypeAnalysis &TA, bool returnUsed, bool shadowReturnUsed, const FnTypeInfo &typeInfo, bool subsequent_calls_may_write, const std::vector< bool > _overwritten_args, const std::vector< bool > &nowrite_shadows, bool forceAnonymousTape, bool runtimeActivity, bool strongZero, unsigned width, bool AtomicAdd, bool omp=false)
Create an augmented forward pass.
llvm::Function * CreateForwardDiff(RequestContext context, llvm::Function *todiff, DIFFE_TYPE retType, llvm::ArrayRef< DIFFE_TYPE > constant_args, TypeAnalysis &TA, bool returnValue, DerivativeMode mode, bool freeMemory, bool runtimeActivity, bool strongZero, unsigned width, llvm::Type *additionalArg, const FnTypeInfo &typeInfo, bool subsequent_calls_may_write, const std::vector< bool > _overwritten_args, const AugmentedReturn *augmented, bool omp=false)
Create the forward (or forward split) mode derivative function.
llvm::Function * CreatePrimalAndGradient(RequestContext context, const ReverseCacheKey &&key, TypeAnalysis &TA, const AugmentedReturn *augmented, bool omp=false)
Create the reverse pass, or combined forward+reverse derivative function.
llvm::ValueMap< const llvm::Instruction *, AssertingReplacingVH > unwrappedLoads
llvm::SmallPtrSet< llvm::Instruction *, 4 > unnecessaryIntermediates
llvm::ValueMap< const llvm::Value *, InvertedPointerVH > invertedPointers
bool getContext(llvm::BasicBlock *BB, LoopContext &lc)
const std::map< llvm::Instruction *, bool > * can_modref_map
llvm::DebugLoc getNewFromOriginal(const llvm::DebugLoc L) const
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.
llvm::SmallPtrSet< llvm::Instruction *, 4 > TapesToPreventRecomputation
A set of tape extractions to enforce a cache of rather than attempting to recompute.
void eraseWithPlaceholder(llvm::Instruction *I, llvm::Instruction *orig, const llvm::Twine &suffix="_replacementA", bool erase=true)
bool legalRecompute(const llvm::Value *val, const llvm::ValueToValueMapTy &available, llvm::IRBuilder<> *BuilderM, bool reverse=false, bool legalRecomputeCache=true) const
TypeResults TR
std::map< const llvm::Value *, bool > knownRecomputeHeuristic
llvm::Value * getOrInsertTotalMultiplicativeProduct(llvm::Value *val, LoopContext &lc)
llvm::BasicBlock * addReverseBlock(llvm::BasicBlock *currentBlock, llvm::Twine const &name, bool forkCache=true, bool push=true)
llvm::SmallPtrSet< llvm::Value *, 4 > allocationsToBeRematerialized
llvm::ValueMap< llvm::Value *, Rematerializer > rematerializableAllocations
llvm::ValueMap< llvm::PHINode *, llvm::WeakTrackingVH > fictiousPHIs
llvm::BasicBlock * getOriginalFromNew(const llvm::BasicBlock *newinst) const
std::map< UsageKey, bool > populateSeenFromKnownRecompute() const
unsigned getWidth()
llvm::Function * oldFunc
llvm::SmallVector< llvm::OperandBundleDef, 2 > getInvertedBundles(llvm::CallInst *orig, llvm::ArrayRef< ValueType > types, llvm::IRBuilder<> &Builder2, bool lookup, const llvm::ValueToValueMapTy &available=llvm::ValueToValueMapTy())
llvm::LoopInfo * OrigLI
void replaceAWithB(llvm::Value *A, llvm::Value *B, bool storeInCache=false) override
Replace this instruction both in LLVM modules and any local data-structures.
EnzymeLogic & Logic
DIFFE_TYPE getDiffeType(llvm::Value *v, bool foreignFunction) const
void replaceAndRemoveUnwrapCacheFor(llvm::Value *A, llvm::Value *B)
llvm::Value * isOriginal(const llvm::Value *newinst) const
std::vector< bool > nowrite_shadows
llvm::Value * unwrapM(llvm::Value *const val, llvm::IRBuilder<> &BuilderM, const llvm::ValueToValueMapTy &available, UnwrapMode unwrapMode, llvm::BasicBlock *scope=nullptr, bool permitCache=true) override final
if full unwrap, don't just unwrap this instruction, but also its operands, etc
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.
bool isAtomic(llvm::Value *origptr) const
llvm::MDNode * getDerivativeAliasScope(const llvm::Value *origptr, ssize_t newptr)
llvm::ValueMap< const llvm::Value *, AssertingReplacingVH > originalToNewFn
llvm::Value * invertPointerM(llvm::Value *val, llvm::IRBuilder<> &BuilderM)
llvm::Value * getOrInsertConditionalIndex(llvm::Value *val, LoopContext &lc, bool pickTrue)
DIFFE_TYPE getReturnDiffeType(llvm::Value *orig, bool *primalReturnUsedP, bool *shadowReturnUsedP, DerivativeMode cmode) const
void setPtrDiffe(llvm::Instruction *orig, llvm::Value *ptr, llvm::Value *newval, llvm::IRBuilder<> &BuilderM, llvm::MaybeAlign align, unsigned start, unsigned size, bool isVolatile, llvm::AtomicOrdering ordering, llvm::SyncScope::ID syncScope, llvm::Value *mask, llvm::ArrayRef< llvm::Metadata * > noAlias, llvm::ArrayRef< llvm::Metadata * > scopes, bool needs_post_cache=false)
bool isConstantInstruction(const llvm::Instruction *inst) const
llvm::ValueMap< const llvm::Value *, AssertingReplacingVH > newToOriginalFn
static llvm::Type * getShadowType(llvm::Type *ty, unsigned width)
llvm::Value * cacheForReverse(llvm::IRBuilder<> &BuilderQ, llvm::Value *malloc, int idx, bool replace=true)
llvm::ValueMap< llvm::Value *, ShadowRematerializer > backwardsOnlyShadows
Only loaded from and stored to (not captured), mapped to the stores (and memset).
bool isConstantValue(llvm::Value *val) const
void erase(llvm::Instruction *I) override
Erase this instruction both from LLVM modules and any local data-structures.
TypeAnalysis & interprocedural
Calling TypeAnalysis to be used in the case of calls to other functions.
std::map< llvm::Value *, TypeTree > analysis
Intermediate conservative, but correct Type analysis results.
A holder class representing the results of running TypeAnalysis on a given function.
llvm::Type * addingType(size_t num, llvm::Value *val, size_t start=0) const
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...
void dump(llvm::raw_ostream &ss=llvm::errs()) const
Prints all known information.
TypeAnalyzer * analyzer
bool anyPointer(llvm::Value *val) const
Whether any part of the top level register can contain a pointer e.g.
std::set< int64_t > knownIntegralValues(llvm::Value *val) const
The set of values val will take on during this program.
const TypeTree & query(llvm::Value *val) const
The TypeTree of a particular Value.
ConcreteType intType(size_t num, llvm::Value *val, llvm::Instruction *I, GradientUtils *gutils=nullptr, llvm::IRBuilder<> *BuilderIfShouldErr=nullptr, bool pointerIntSame=false) const
llvm::Function * getFunction() const
FnTypeInfo getCallInfo(llvm::CallBase &CI, llvm::Function &fn) const
Class representing the underlying types of values as sequences of offsets to a ConcreteType.
Definition TypeTree.h:72
TypeTree Only(int Off, llvm::Instruction *orig) const
Prepend an offset to all mappings.
Definition TypeTree.h:473
TypeTree Data0() const
Peel off the outermost index at offset 0.
Definition TypeTree.h:515
static TypeTree parse(llvm::StringRef str, llvm::LLVMContext &ctx)
Definition TypeTree.h:86
TypeTree ShiftIndices(const llvm::DataLayout &dl, const int offset, const int maxSize, size_t addOffset=0) const
Replace mappings in the range in [offset, offset+maxSize] with those in.
Definition TypeTree.h:912
TypeTree Lookup(size_t len, const llvm::DataLayout &dl) const
Select all submappings whose first index is in range [0, len) and remove the first index.
Definition TypeTree.h:595
std::string str() const
Returns a string representation of this TypeTree.
Definition TypeTree.h:1455
bool insert(const std::vector< int > Seq, ConcreteType CT, bool PointerIntSame=false)
Return if changed.
Definition TypeTree.h:236
TypeTree PurgeAnything() const
Keep only mappings where the type is not an Anything
Definition TypeTree.h:1113
llvm::Type * allFloat(llvm::Value *val, const llvm::DataLayout &dl, bool anythingIsFloat=false) const
Definition TypeTree.h:816
bool callShouldNotUseDerivative(const GradientUtils *gutils, llvm::CallBase &orig, QueryType qtype, const llvm::Value *val)
Return whether or not this is a constant and should use reverse pass.
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)
Struct containing all contextual type information for a particular function call.
std::map< llvm::Argument *, TypeTree > Arguments
Types of arguments.
TypeTree Return
Type of return.
std::map< llvm::Argument *, std::set< int64_t > > KnownValues
The specific constant(s) known to represented by an argument, if constant.
Container for all loop information to synthesize gradients.
llvm::BasicBlock * header
Header of this loop.
llvm::AssertingVH< llvm::Instruction > incvar
Increment of the induction.
llvm::SmallPtrSet< llvm::BasicBlock *, 8 > exitBlocks
All blocks this loop exits too.
llvm::BasicBlock * preheader
Preheader of this loop.
todiff is the function to differentiate retType is the activity info of the return.