https://github.com/luporl created https://github.com/llvm/llvm-project/pull/80488
Use the new copyprivate list from omp.single to emit calls to __kmpc_copyprivate, during the creation of the single operation in OMPIRBuilder. This is patch 4 of 4, to add support for COPYPRIVATE in Flang. Original PR: https://github.com/llvm/llvm-project/pull/73128 >From 52462e6790194c19465f017b81e51e4a99136d3a Mon Sep 17 00:00:00 2001 From: Leandro Lupori <leandro.lup...@linaro.org> Date: Fri, 2 Feb 2024 17:16:34 -0300 Subject: [PATCH] [llvm][mlir][OMPIRBuilder] Translate omp.single's copyprivate Use the new copyprivate list from omp.single to emit calls to __kmpc_copyprivate, during the creation of the single operation in OMPIRBuilder. This is patch 4 of 4, to add support for COPYPRIVATE in Flang. Original PR: https://github.com/llvm/llvm-project/pull/73128 --- .../llvm/Frontend/OpenMP/OMPIRBuilder.h | 6 +- llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 23 +++- .../Frontend/OpenMPIRBuilderTest.cpp | 111 ++++++++++++++++++ .../OpenMP/OpenMPToLLVMIRTranslation.cpp | 20 +++- mlir/test/Target/LLVMIR/openmp-llvm.mlir | 32 +++++ 5 files changed, 187 insertions(+), 5 deletions(-) diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h index 669104307fa0e..ab92c172c75ae 100644 --- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h +++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h @@ -1819,12 +1819,16 @@ class OpenMPIRBuilder { /// \param FiniCB Callback to finalize variable copies. /// \param IsNowait If false, a barrier is emitted. /// \param DidIt Local variable used as a flag to indicate 'single' thread + /// \param CPVars copyprivate variables. + /// \param CPFuncs copy functions to use for each copyprivate variable. /// /// \returns The insertion position *after* the single call. InsertPointTy createSingle(const LocationDescription &Loc, BodyGenCallbackTy BodyGenCB, FinalizeCallbackTy FiniCB, bool IsNowait, - llvm::Value *DidIt); + llvm::Value *DidIt, + ArrayRef<llvm::Value *> CPVars = {}, + ArrayRef<llvm::Function *> CPFuncs = {}); /// Generator for '#omp master' /// diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp index f6cf358119fb7..7abac0f660ef8 100644 --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp @@ -3992,7 +3992,8 @@ OpenMPIRBuilder::createCopyPrivate(const LocationDescription &Loc, OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSingle( const LocationDescription &Loc, BodyGenCallbackTy BodyGenCB, - FinalizeCallbackTy FiniCB, bool IsNowait, llvm::Value *DidIt) { + FinalizeCallbackTy FiniCB, bool IsNowait, llvm::Value *DidIt, + ArrayRef<llvm::Value *> CPVars, ArrayRef<llvm::Function *> CPFuncs) { if (!updateToLocation(Loc)) return Loc.IP; @@ -4015,17 +4016,33 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSingle( Function *ExitRTLFn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_end_single); Instruction *ExitCall = Builder.CreateCall(ExitRTLFn, Args); + auto FiniCBWrapper = [&](InsertPointTy IP) { + FiniCB(IP); + + if (DidIt) + Builder.CreateStore(Builder.getInt32(1), DidIt); + }; + // generates the following: // if (__kmpc_single()) { // .... single region ... // __kmpc_end_single // } + // __kmpc_copyprivate // __kmpc_barrier - EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB, + EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCBWrapper, /*Conditional*/ true, /*hasFinalize*/ true); - if (!IsNowait) + + if (DidIt) { + for (size_t I = 0, E = CPVars.size(); I < E; ++I) + // NOTE BufSize is currently unused, so just pass 0. + createCopyPrivate(LocationDescription(Builder.saveIP(), Loc.DL), + /*BufSize=*/ConstantInt::get(Int64, 0), CPVars[I], + CPFuncs[I], DidIt); + // NOTE __kmpc_copyprivate already inserts a barrier + } else if (!IsNowait) createBarrier(LocationDescription(Builder.saveIP(), Loc.DL), omp::Directive::OMPD_unknown, /* ForceSimpleCall */ false, /* CheckCancelFlag */ false); diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp index e79d0bb2f65ae..0eb1039aa442c 100644 --- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp +++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp @@ -3464,6 +3464,117 @@ TEST_F(OpenMPIRBuilderTest, SingleDirectiveNowait) { EXPECT_EQ(ExitBarrier, nullptr); } +TEST_F(OpenMPIRBuilderTest, SingleDirectiveCopyPrivate) { + using InsertPointTy = OpenMPIRBuilder::InsertPointTy; + OpenMPIRBuilder OMPBuilder(*M); + OMPBuilder.initialize(); + F->setName("func"); + IRBuilder<> Builder(BB); + + OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL}); + + AllocaInst *PrivAI = nullptr; + + BasicBlock *EntryBB = nullptr; + BasicBlock *ThenBB = nullptr; + + Value *CPVar = Builder.CreateAlloca(F->arg_begin()->getType()); + Builder.CreateStore(F->arg_begin(), CPVar); + + FunctionType *CopyFuncTy = FunctionType::get( + Builder.getVoidTy(), {Builder.getPtrTy(), Builder.getPtrTy()}, false); + Function *CopyFunc = + Function::Create(CopyFuncTy, Function::PrivateLinkage, "copy_var", *M); + + Value *DidIt = Builder.CreateAlloca(Type::getInt32Ty(Builder.getContext())); + + auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) { + if (AllocaIP.isSet()) + Builder.restoreIP(AllocaIP); + else + Builder.SetInsertPoint(&*(F->getEntryBlock().getFirstInsertionPt())); + PrivAI = Builder.CreateAlloca(F->arg_begin()->getType()); + Builder.CreateStore(F->arg_begin(), PrivAI); + + llvm::BasicBlock *CodeGenIPBB = CodeGenIP.getBlock(); + llvm::Instruction *CodeGenIPInst = &*CodeGenIP.getPoint(); + EXPECT_EQ(CodeGenIPBB->getTerminator(), CodeGenIPInst); + + Builder.restoreIP(CodeGenIP); + + // collect some info for checks later + ThenBB = Builder.GetInsertBlock(); + EntryBB = ThenBB->getUniquePredecessor(); + + // simple instructions for body + Value *PrivLoad = + Builder.CreateLoad(PrivAI->getAllocatedType(), PrivAI, "local.use"); + Builder.CreateICmpNE(F->arg_begin(), PrivLoad); + }; + + auto FiniCB = [&](InsertPointTy IP) { + BasicBlock *IPBB = IP.getBlock(); + EXPECT_NE(IPBB->end(), IP.getPoint()); + }; + + Builder.restoreIP(OMPBuilder.createSingle(Builder, BodyGenCB, FiniCB, + /*IsNowait*/ false, DidIt, {CPVar}, + {CopyFunc})); + Value *EntryBBTI = EntryBB->getTerminator(); + EXPECT_NE(EntryBBTI, nullptr); + EXPECT_TRUE(isa<BranchInst>(EntryBBTI)); + BranchInst *EntryBr = cast<BranchInst>(EntryBB->getTerminator()); + EXPECT_TRUE(EntryBr->isConditional()); + EXPECT_EQ(EntryBr->getSuccessor(0), ThenBB); + BasicBlock *ExitBB = ThenBB->getUniqueSuccessor(); + EXPECT_EQ(EntryBr->getSuccessor(1), ExitBB); + + CmpInst *CondInst = cast<CmpInst>(EntryBr->getCondition()); + EXPECT_TRUE(isa<CallInst>(CondInst->getOperand(0))); + + CallInst *SingleEntryCI = cast<CallInst>(CondInst->getOperand(0)); + EXPECT_EQ(SingleEntryCI->arg_size(), 2U); + EXPECT_EQ(SingleEntryCI->getCalledFunction()->getName(), "__kmpc_single"); + EXPECT_TRUE(isa<GlobalVariable>(SingleEntryCI->getArgOperand(0))); + + CallInst *SingleEndCI = nullptr; + for (auto &FI : *ThenBB) { + Instruction *Cur = &FI; + if (isa<CallInst>(Cur)) { + SingleEndCI = cast<CallInst>(Cur); + if (SingleEndCI->getCalledFunction()->getName() == "__kmpc_end_single") + break; + SingleEndCI = nullptr; + } + } + EXPECT_NE(SingleEndCI, nullptr); + EXPECT_EQ(SingleEndCI->arg_size(), 2U); + EXPECT_TRUE(isa<GlobalVariable>(SingleEndCI->getArgOperand(0))); + EXPECT_EQ(SingleEndCI->getArgOperand(1), SingleEntryCI->getArgOperand(1)); + + CallInst *CopyPrivateCI = nullptr; + bool FoundBarrier = false; + for (auto &FI : *ExitBB) { + Instruction *Cur = &FI; + if (auto *CI = dyn_cast<CallInst>(Cur)) { + if (CI->getCalledFunction()->getName() == "__kmpc_barrier") + FoundBarrier = true; + else if (CI->getCalledFunction()->getName() == "__kmpc_copyprivate") + CopyPrivateCI = CI; + } + } + EXPECT_FALSE(FoundBarrier); + EXPECT_NE(CopyPrivateCI, nullptr); + EXPECT_EQ(CopyPrivateCI->arg_size(), 6U); + EXPECT_TRUE(isa<AllocaInst>(CopyPrivateCI->getArgOperand(3))); + EXPECT_EQ(CopyPrivateCI->getArgOperand(3), CPVar); + EXPECT_TRUE(isa<Function>(CopyPrivateCI->getArgOperand(4))); + EXPECT_EQ(CopyPrivateCI->getArgOperand(4), CopyFunc); + EXPECT_TRUE(isa<LoadInst>(CopyPrivateCI->getArgOperand(5))); + LoadInst *DidItLI = cast<LoadInst>(CopyPrivateCI->getArgOperand(5)); + EXPECT_EQ(DidItLI->getOperand(0), DidIt); +} + TEST_F(OpenMPIRBuilderTest, OMPAtomicReadFlt) { OpenMPIRBuilder OMPBuilder(*M); OMPBuilder.initialize(); diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp index 23e101f1e4527..964a1aeb5a00d 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp @@ -656,8 +656,26 @@ convertOmpSingle(omp::SingleOp &singleOp, llvm::IRBuilderBase &builder, moduleTranslation, bodyGenStatus); }; auto finiCB = [&](InsertPointTy codeGenIP) {}; + + // Handle copyprivate + Operation::operand_range cpVars = singleOp.getCopyprivateVars(); + std::optional<ArrayAttr> cpFuncs = singleOp.getCopyprivateFuncs(); + llvm::SmallVector<llvm::Value *> llvmCPVars; + llvm::SmallVector<llvm::Function *> llvmCPFuncs; + for (size_t i = 0, e = cpVars.size(); i < e; ++i) { + llvmCPVars.push_back(moduleTranslation.lookupValue(cpVars[i])); + auto llvmFuncOp = SymbolTable::lookupNearestSymbolFrom<LLVM::LLVMFuncOp>( + singleOp, cast<SymbolRefAttr>((*cpFuncs)[i])); + llvmCPFuncs.push_back( + moduleTranslation.lookupFunction(llvmFuncOp.getName())); + } + llvm::Value *didIt = nullptr; + if (!llvmCPVars.empty()) + didIt = builder.CreateAlloca(llvm::Type::getInt32Ty(builder.getContext())); + builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createSingle( - ompLoc, bodyCB, finiCB, singleOp.getNowait(), /*DidIt=*/nullptr)); + ompLoc, bodyCB, finiCB, singleOp.getNowait(), didIt, llvmCPVars, + llvmCPFuncs)); return bodyGenStatus; } diff --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir index 29baa84e7e19d..8a3d5d6407659 100644 --- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir +++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir @@ -2165,6 +2165,38 @@ llvm.func @single_nowait(%x: i32, %y: i32, %zaddr: !llvm.ptr) { // ----- +llvm.func @copy_i32(!llvm.ptr, !llvm.ptr) +llvm.func @copy_f32(!llvm.ptr, !llvm.ptr) + +// CHECK-LABEL: @single_copyprivate +// CHECK-SAME: (ptr %[[ip:.*]], ptr %[[fp:.*]]) +llvm.func @single_copyprivate(%ip: !llvm.ptr, %fp: !llvm.ptr) { + // CHECK: call i32 @__kmpc_single + omp.single copyprivate(%ip -> @copy_i32 : !llvm.ptr, %fp -> @copy_f32 : !llvm.ptr) { + // CHECK: %[[i:.*]] = load i32, ptr %[[ip]] + %i = llvm.load %ip : !llvm.ptr -> i32 + // CHECK: %[[i2:.*]] = add i32 %[[i]], %[[i]] + %i2 = llvm.add %i, %i : i32 + // CHECK: store i32 %[[i2]], ptr %[[ip]] + llvm.store %i2, %ip : i32, !llvm.ptr + // CHECK: %[[f:.*]] = load float, ptr %[[fp]] + %f = llvm.load %fp : !llvm.ptr -> f32 + // CHECK: %[[f2:.*]] = fadd float %[[f]], %[[f]] + %f2 = llvm.fadd %f, %f : f32 + // CHECK: store float %[[f2]], ptr %[[fp]] + llvm.store %f2, %fp : f32, !llvm.ptr + // CHECK: call void @__kmpc_end_single + // CHECK: call void @__kmpc_copyprivate({{.*}}, ptr %[[ip]], ptr @copy_i32, {{.*}}) + // CHECK: call void @__kmpc_copyprivate({{.*}}, ptr %[[fp]], ptr @copy_f32, {{.*}}) + // CHECK-NOT: call void @__kmpc_barrier + omp.terminator + } + // CHECK: ret void + llvm.return +} + +// ----- + // CHECK: @_QFsubEx = internal global i32 undef // CHECK: @_QFsubEx.cache = common global ptr null _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits