Skip to content

Commit 7be5247

Browse files
wsmoseswsmoses
authored
LLVM integration (#2942)
1 parent 93668c4 commit 7be5247

6 files changed

Lines changed: 27 additions & 20 deletions

File tree

‎.github/workflows/enzyme-mlir.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ jobs:
4545
- uses: actions/checkout@v4
4646
with:
4747
repository: 'llvm/llvm-project'
48-
ref: 'b713aaedb6eedd7d20f666038d945e50eafd4c27'
48+
ref: '615644763ffcfc939bda59e9f3e71d6e7d9fd264'
4949
path: 'llvm-project'
5050

5151
- name: Set BASE_DIR

‎enzyme/Enzyme/DiffeGradientUtils.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1098,7 +1098,7 @@ void DiffeGradientUtils::addToInvertedPtrDiffe(Instruction *orig,
10981098
// upon
10991099
if (isa<AllocaInst>(TmpOrig) &&
11001100
(Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
1101-
Arch == Triple::amdgcn)) {
1101+
Arch == Triple::amd_target)) {
11021102
Atomic = false;
11031103
}
11041104
// Moreover no need to do atomic on local shadows regardless since they are
@@ -1110,9 +1110,9 @@ void DiffeGradientUtils::addToInvertedPtrDiffe(Instruction *orig,
11101110
Atomic = false;
11111111

11121112
if (Atomic) {
1113-
// For amdgcn constant AS is 4 and if the primal is in it we need to cast
1114-
// the derivative value to AS 1
1115-
if (Arch == Triple::amdgcn &&
1113+
// For amd_target constant AS is 4 and if the primal is in it we need to
1114+
// cast the derivative value to AS 1
1115+
if (Arch == Triple::amd_target &&
11161116
cast<PointerType>(origptr->getType())->getAddressSpace() == 4) {
11171117
auto rule = [&](Value *ptr) {
11181118
return BuilderM.CreateAddrSpaceCast(ptr,

‎enzyme/Enzyme/Enzyme.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1459,7 +1459,7 @@ class EnzymeBase {
14591459

14601460
auto Arch = Triple(CI->getModule()->getTargetTriple()).getArch();
14611461
bool AtomicAdd = Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
1462-
Arch == Triple::amdgcn;
1462+
Arch == Triple::amd_target;
14631463

14641464
TypeAnalysis TA(Logic);
14651465
FnTypeInfo type_args = populate_type_args(TA, fn, mode);
@@ -2608,7 +2608,7 @@ class EnzymeBase {
26082608
.getArch();
26092609

26102610
bool AtomicAdd = Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
2611-
Arch == Triple::amdgcn;
2611+
Arch == Triple::amd_target;
26122612

26132613
IRBuilder<> Builder(CI);
26142614
auto val = GradientUtils::GetOrCreateShadowConstant(

‎enzyme/Enzyme/EnzymeLogic.cpp‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,7 @@ struct CacheAnalysis {
267267
assert(li.getParent()->getParent() == oldFunc);
268268

269269
auto Arch = llvm::Triple(oldFunc->getParent()->getTargetTriple()).getArch();
270-
if (Arch == Triple::amdgcn &&
270+
if (Arch == Triple::amd_target &&
271271
cast<PointerType>(li.getOperand(0)->getType())->getAddressSpace() ==
272272
4) {
273273
return false;
@@ -4545,8 +4545,9 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
45454545
auto Arch =
45464546
llvm::Triple(gutils->newFunc->getParent()->getTargetTriple()).getArch();
45474547
unsigned int SharedAddrSpace =
4548-
Arch == Triple::amdgcn ? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
4549-
: 3;
4548+
Arch == Triple::amd_target
4549+
? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
4550+
: 3;
45504551

45514552
if (key.mode == DerivativeMode::ReverseModeCombined) {
45524553
BasicBlock *sharedBlock = nullptr;
@@ -4556,7 +4557,7 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
45564557
gutils->inversionAllocs->begin());
45574558

45584559
if ((Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
4559-
Arch == Triple::amdgcn) &&
4560+
Arch == Triple::amd_target) &&
45604561
g.getType()->getAddressSpace() == SharedAddrSpace) {
45614562
if (sharedBlock == nullptr)
45624563
sharedBlock = BasicBlock::Create(entry->getContext(), "shblock",
@@ -4582,7 +4583,7 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
45824583
gutils->newFunc->getParent(), Intrinsic::nvvm_read_ptx_sreg_tid_y));
45834584
tz = ebuilder.CreateCall(getIntrinsicDeclaration(
45844585
gutils->newFunc->getParent(), Intrinsic::nvvm_read_ptx_sreg_tid_z));
4585-
} else if (Arch == Triple::amdgcn) {
4586+
} else if (Arch == Triple::amd_target) {
45864587
tx = ebuilder.CreateCall(getIntrinsicDeclaration(
45874588
gutils->newFunc->getParent(), Intrinsic::amdgcn_workitem_id_x));
45884589
ty = ebuilder.CreateCall(getIntrinsicDeclaration(
@@ -4601,12 +4602,12 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
46014602
IRBuilder<> instbuilder(OldEntryInsts, OldEntryInsts->begin());
46024603

46034604
#if LLVM_VERSION_MAJOR > 20
4604-
auto BarrierInst = Arch == Triple::amdgcn
4605+
auto BarrierInst = Arch == Triple::amd_target
46054606
? (llvm::Intrinsic::ID)Intrinsic::amdgcn_s_barrier
46064607
: (llvm::Intrinsic::ID)
46074608
Intrinsic::nvvm_barrier_cta_sync_aligned_all;
46084609
#else
4609-
auto BarrierInst = Arch == Triple::amdgcn
4610+
auto BarrierInst = Arch == Triple::amd_target
46104611
? (llvm::Intrinsic::ID)Intrinsic::amdgcn_s_barrier
46114612
: (llvm::Intrinsic::ID)Intrinsic::nvvm_barrier0;
46124613
#endif

‎enzyme/Enzyme/GradientUtils.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4676,12 +4676,12 @@ Constant *GradientUtils::GetOrCreateShadowConstant(
46764676
}
46774677

46784678
auto Arch = llvm::Triple(arg->getParent()->getTargetTriple()).getArch();
4679-
int SharedAddrSpace = Arch == Triple::amdgcn
4679+
int SharedAddrSpace = Arch == Triple::amd_target
46804680
? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
46814681
: 3;
46824682
int AddrSpace = cast<PointerType>(arg->getType())->getAddressSpace();
46834683
if ((Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
4684-
Arch == Triple::amdgcn) &&
4684+
Arch == Triple::amd_target) &&
46854685
AddrSpace == SharedAddrSpace) {
46864686
assert(0 && "shared memory not handled in meta global");
46874687
}
@@ -5686,12 +5686,12 @@ Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM,
56865686
auto Arch =
56875687
llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch();
56885688
int SharedAddrSpace =
5689-
Arch == Triple::amdgcn
5689+
Arch == Triple::amd_target
56905690
? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
56915691
: 3;
56925692
int AddrSpace = cast<PointerType>(arg->getType())->getAddressSpace();
56935693
if ((Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
5694-
Arch == Triple::amdgcn) &&
5694+
Arch == Triple::amd_target) &&
56955695
AddrSpace == SharedAddrSpace) {
56965696
llvm::errs() << "warning found shared memory\n";
56975697
Type *type = arg->getValueType();
@@ -6746,7 +6746,7 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
67466746
auto Arch =
67476747
llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch();
67486748
unsigned int SharedAddrSpace =
6749-
Arch == Triple::amdgcn
6749+
Arch == Triple::amd_target
67506750
? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
67516751
: 3;
67526752
if (cast<PointerType>(LI->getPointerOperand()->getType())
@@ -7101,7 +7101,7 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
71017101
auto Arch =
71027102
llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch();
71037103
unsigned int SharedAddrSpace =
7104-
Arch == Triple::amdgcn
7104+
Arch == Triple::amd_target
71057105
? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
71067106
: 3;
71077107
if (EnzymeSharedForward && scev1 != OrigSE->getCouldNotCompute() &&

‎enzyme/Enzyme/Utils.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,12 @@
6262
#include <optional>
6363
#endif
6464

65+
#if LLVM_VERSION_MAJOR >= 23
66+
#define amd_target amdgpu
67+
#else
68+
#define amd_target amdgcn
69+
#endif
70+
6571
#include "llvm/IR/DiagnosticInfo.h"
6672

6773
#include "llvm/Analysis/OptimizationRemarkEmitter.h"

0 commit comments

Comments
 (0)