Skip to content

Commit f10b38d

Browse files
authored
handle indirect enzyme_const loads (rust-lang#753)
* add testcase for enzyme loads behind phis * handle load trough phis * adjust test for older llvms * adjust test for older llvms * adjust test for even older llvms
1 parent 8abdb9f commit f10b38d

File tree

2 files changed

+72
-0
lines changed

2 files changed

+72
-0
lines changed

enzyme/Enzyme/Enzyme.cpp

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -498,6 +498,38 @@ castToDiffeFunctionArgType(IRBuilder<> &Builder, llvm::CallInst *CI,
498498
return Builder.CreateBitCast(value, destType);
499499
}
500500

501+
static Optional<StringRef> getMetadataName(llvm::Value *res);
502+
503+
// if all phi arms are (recursively) based on the same metaString, use that
504+
static Optional<StringRef> recursePhiReads(PHINode *val) {
505+
Optional<StringRef> finalMetadata;
506+
SmallVector<PHINode *, 1> todo = {val};
507+
SmallSet<PHINode *, 1> done;
508+
while (todo.size()) {
509+
auto phiInst = todo.back();
510+
todo.pop_back();
511+
if (done.count(phiInst))
512+
continue;
513+
done.insert(phiInst);
514+
for (unsigned j = 0; j < phiInst->getNumIncomingValues(); ++j) {
515+
auto newVal = phiInst->getIncomingValue(j);
516+
if (auto phi = dyn_cast<PHINode>(newVal)) {
517+
todo.push_back(phi);
518+
} else {
519+
Optional<StringRef> metaString = getMetadataName(newVal);
520+
if (metaString) {
521+
if (!finalMetadata) {
522+
finalMetadata = metaString;
523+
} else if (finalMetadata != metaString) {
524+
return None;
525+
}
526+
}
527+
}
528+
}
529+
}
530+
return finalMetadata;
531+
}
532+
501533
static Optional<StringRef> getMetadataName(llvm::Value *res) {
502534
if (auto av = dyn_cast<MetadataAsValue>(res)) {
503535
return cast<MDString>(av->getMetadata())->getString();
@@ -528,6 +560,9 @@ static Optional<StringRef> getMetadataName(llvm::Value *res) {
528560
} else if (auto gv = dyn_cast<AllocaInst>(res)) {
529561
return gv->getName();
530562
} else {
563+
if (isa<PHINode>(res)) {
564+
return recursePhiReads(cast<PHINode>(res));
565+
}
531566
return Optional<StringRef>();
532567
}
533568
}
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -sroa -instsimplify -simplifycfg -S | FileCheck %s
2+
3+
; ModuleID = 'ld-temp.o'
4+
source_filename = "ld-temp.o"
5+
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
6+
target triple = "x86_64-unknown-linux-gnu"
7+
8+
@enzyme_const = internal global i32 0, align 4
9+
10+
define internal double @callable(double* %x) {
11+
ret double 1.000000e+00
12+
}
13+
14+
define internal void @_Z19testSensitivitiesADv(double* %primal, double* %grad, double* %a) {
15+
br label %bb1
16+
17+
bb1:
18+
%a1 = load i32, i32* @enzyme_const, align 4
19+
br label %bb3
20+
21+
bb2:
22+
%a2 = load i32, i32* @enzyme_const, align 4
23+
br label %bb3
24+
25+
bb3:
26+
%a3 = phi i32 [ %a1, %bb1 ], [ %a2, %bb2 ]
27+
;%a3 = load i32, i32* @enzyme_const, align 4
28+
29+
%c1 = call double (...) @__enzyme_autodiff(i8* bitcast (double (double*)* @callable to i8*), i32 %a3, double* %a)
30+
ret void
31+
}
32+
declare dso_local double @__enzyme_autodiff(...)
33+
34+
; CHECK: define internal void @diffecallable(double* %x, double %differeturn)
35+
; CHECK-NEXT: invert:
36+
; CHECK-NEXT: ret void
37+
; CHECK-NEXT: }

0 commit comments

Comments
 (0)