diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index 86341a9fada..f25791c5293 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -4038,10 +4038,14 @@ class AdjointGenerator : public llvm::InstVisitor { case Intrinsic::nvvm_membar_cta: case Intrinsic::nvvm_membar_gl: case Intrinsic::nvvm_membar_sys: { - SmallVector args = {}; - auto cal = cast( - Builder2.CreateCall(getIntrinsicDeclaration(M, ID), args)); - cal->setCallingConv(getIntrinsicDeclaration(M, ID)->getCallingConv()); + auto &Call = cast(I); + SmallVector args; + args.reserve(Call.arg_size()); + for (unsigned i = 0; i < Call.arg_size(); ++i) + args.push_back(gutils->getNewFromOriginal(Call.getArgOperand(i))); + auto *Fn = getIntrinsicDeclaration(M, ID); + auto cal = cast(Builder2.CreateCall(Fn, args)); + cal->setCallingConv(Fn->getCallingConv()); cal->setDebugLoc(gutils->getNewFromOriginal(I.getDebugLoc())); return false; } diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp index ae7839d6dcd..20220ed8d50 100644 --- a/enzyme/Enzyme/EnzymeLogic.cpp +++ b/enzyme/Enzyme/EnzymeLogic.cpp @@ -4605,10 +4605,17 @@ Function *EnzymeLogic::CreatePrimalAndGradient( auto BarrierInst = Arch == Triple::amdgcn ? (llvm::Intrinsic::ID)Intrinsic::amdgcn_s_barrier : (llvm::Intrinsic::ID)Intrinsic::nvvm_barrier0; +#endif + SmallVector BarrierArgs; +#if LLVM_VERSION_MAJOR > 20 + if (Arch != Triple::amdgcn) { + BarrierArgs.push_back(ConstantInt::get( + Type::getInt32Ty(gutils->newFunc->getContext()), 0)); + } #endif instbuilder.CreateCall( getIntrinsicDeclaration(gutils->newFunc->getParent(), BarrierInst), - {}); + BarrierArgs); OldEntryInsts->moveAfter(entry); sharedBlock->moveAfter(entry); IRBuilder<> sbuilder(sharedBlock); diff --git a/enzyme/test/Enzyme/ReverseMode/cuda-barrier-sync.ll b/enzyme/test/Enzyme/ReverseMode/cuda-barrier-sync.ll new file mode 100644 index 00000000000..ae41a14daf7 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/cuda-barrier-sync.ll @@ -0,0 +1,106 @@ +; RUN: split-file %s %t +; RUN: if [ %llvmver -le 20 ]; then %opt < %t/llvm20.ll %newLoadEnzyme -enzyme-preopt=false -enzyme-detect-readthrow=0 -passes="enzyme" -S | FileCheck %t/llvm20.ll; fi +; RUN: if [ %llvmver -gt 20 ]; then %opt < %t/llvm21plus.ll %newLoadEnzyme -enzyme-preopt=false -enzyme-detect-readthrow=0 -passes="enzyme" -S | FileCheck %t/llvm21plus.ll; fi + +;--- llvm20.ll +target triple = "nvptx64-nvidia-cuda" + +declare void @llvm.nvvm.barrier0() +declare float @__enzyme_autodiff(float (float)*, ...) + +define float @f_sync(float %x) { +entry: + call void @llvm.nvvm.barrier0() + %res = fadd float %x, 1.000000e+00 + ret float %res +} + +define float @test(float %x) { +entry: + %r = call float (float (float)*, ...) @__enzyme_autodiff(float (float)* @f_sync, float %x) + ret float %r +} + +; CHECK: define internal { float } @diffef_sync(float %x, float %differeturn) +; CHECK: call void @llvm.nvvm.barrier0() +; CHECK: call void @llvm.nvvm.barrier0() +; CHECK: ret { float } + +;--- llvm21plus.ll +target triple = "nvptx64-nvidia-cuda" + +declare void @llvm.nvvm.barrier.cta.sync.aligned.all(i32) +declare void @llvm.nvvm.barrier.cta.sync.aligned.count(i32, i32) +declare float @__enzyme_autodiff(...) + +define float @f_sync_all(float %x) { +entry: + call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 7) + %res = fadd float %x, 1.000000e+00 + ret float %res +} + +define float @test_all(float %x) { +entry: + %r = call float (float (float)*, ...) @__enzyme_autodiff(float (float)* @f_sync_all, float %x) + ret float %r +} + +define float @f_sync_all_dyn(i32 %id, float %x) { +entry: + call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 %id) + %res = fadd float %x, 1.000000e+00 + ret float %res +} + +define float @test_all_dyn(i32 %id, float %x) { +entry: + %r = call float (float (i32, float)*, ...) @__enzyme_autodiff(float (i32, float)* @f_sync_all_dyn, i32 %id, float %x) + ret float %r +} + +define float @f_sync_count(float %x) { +entry: + call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 7, i32 16) + %res = fadd float %x, 1.000000e+00 + ret float %res +} + +define float @test_count(float %x) { +entry: + %r = call float (float (float)*, ...) @__enzyme_autodiff(float (float)* @f_sync_count, float %x) + ret float %r +} + +define float @f_sync_count_dyn(i32 %id, i32 %n, float %x) { +entry: + call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 %id, i32 %n) + %res = fadd float %x, 1.000000e+00 + ret float %res +} + +define float @test_count_dyn(i32 %id, i32 %n, float %x) { +entry: + %r = call float (float (i32, i32, float)*, ...) @__enzyme_autodiff(float (i32, i32, float)* @f_sync_count_dyn, i32 %id, i32 %n, float %x) + ret float %r +} + +; CHECK: define internal { float } @diffef_sync_all(float %x, float %differeturn) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 7) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 7) +; CHECK: ret { float } + +; CHECK: define internal { float } @diffef_sync_all_dyn(i32 %id, float %x, float %differeturn) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 %id) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 %id) +; CHECK: ret { float } + +; CHECK: define internal { float } @diffef_sync_count(float %x, float %differeturn) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 7, i32 16) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 7, i32 16) +; CHECK: ret { float } + +; CHECK: define internal { float } @diffef_sync_count_dyn(i32 %id, i32 %n, float %x, float %differeturn) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 %id, i32 %n) +; CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 %id, i32 %n) +; CHECK: ret { float }