diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp index ce2fc390db2b0..da11ee9960e1f 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp @@ -1254,13 +1254,13 @@ static LogicalResult allocAndInitializeReductionVars( /// Allocate delayed private variables. Returns the basic block which comes /// after all of these allocations. llvm::Value * for each of these private /// variables are populated in llvmPrivateVars. -template static llvm::Expected -allocatePrivateVars(OP opInst, llvm::IRBuilderBase &builder, +allocatePrivateVars(llvm::IRBuilderBase &builder, LLVM::ModuleTranslation &moduleTranslation, MutableArrayRef privateBlockArgs, MutableArrayRef privateDecls, - llvm::SmallVector &llvmPrivateVars, + MutableArrayRef mlirPrivateVars, + llvm::SmallVectorImpl &llvmPrivateVars, const llvm::OpenMPIRBuilder::InsertPointTy &allocaIP) { // Allocate private vars llvm::BranchInst *allocaTerminator = @@ -1285,19 +1285,18 @@ allocatePrivateVars(OP opInst, llvm::IRBuilderBase &builder, llvm::BasicBlock *privAllocBlock = nullptr; if (!privateBlockArgs.empty()) privAllocBlock = splitBB(builder, true, "omp.private.latealloc"); - for (unsigned i = 0; i < privateBlockArgs.size(); ++i) { - Region &allocRegion = privateDecls[i].getAllocRegion(); + for (auto [privDecl, mlirPrivVar, blockArg] : + llvm::zip_equal(privateDecls, mlirPrivateVars, privateBlockArgs)) { + Region &allocRegion = privDecl.getAllocRegion(); // map allocation region block argument - llvm::Value *nonPrivateVar = - moduleTranslation.lookupValue(opInst.getPrivateVars()[i]); + llvm::Value *nonPrivateVar = moduleTranslation.lookupValue(mlirPrivVar); assert(nonPrivateVar); - moduleTranslation.mapValue(privateDecls[i].getAllocMoldArg(), - nonPrivateVar); + moduleTranslation.mapValue(privDecl.getAllocMoldArg(), nonPrivateVar); // in-place convert the private allocation region SmallVector phis; - if (privateDecls[i].getAllocMoldArg().getUses().empty()) { + if (privDecl.getAllocMoldArg().getUses().empty()) { // TODO this should use // allocaIP.getBlock()->getFirstNonPHIOrDbgOrAlloca() so it goes before // the code for fetching the thread id. Not doing this for now to avoid @@ -1313,7 +1312,7 @@ allocatePrivateVars(OP opInst, llvm::IRBuilderBase &builder, assert(phis.size() == 1 && "expected one allocation to be yielded"); - moduleTranslation.mapValue(privateBlockArgs[i], phis[0]); + moduleTranslation.mapValue(blockArg, phis[0]); llvmPrivateVars.push_back(phis[0]); // clear alloc region block argument mapping in case it needs to be @@ -1561,11 +1560,14 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder, // Collect delayed privatisation declarations MutableArrayRef privateBlockArgs = cast(*taskOp).getPrivateBlockArgs(); + SmallVector mlirPrivateVars; SmallVector llvmPrivateVars; SmallVector privateDecls; + mlirPrivateVars.reserve(privateBlockArgs.size()); llvmPrivateVars.reserve(privateBlockArgs.size()); - privateDecls.reserve(privateBlockArgs.size()); collectPrivatizationDecls(taskOp, privateDecls); + for (mlir::Value privateVar : taskOp.getPrivateVars()) + mlirPrivateVars.push_back(privateVar); auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP) -> llvm::Error { @@ -1575,8 +1577,8 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder, moduleTranslation, allocaIP); llvm::Expected afterAllocas = allocatePrivateVars( - taskOp, builder, moduleTranslation, privateBlockArgs, privateDecls, - llvmPrivateVars, allocaIP); + builder, moduleTranslation, privateBlockArgs, privateDecls, + mlirPrivateVars, llvmPrivateVars, allocaIP); if (handleError(afterAllocas, *taskOp).failed()) return llvm::make_error(); @@ -1879,11 +1881,14 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder, // Collect delayed privatization declarations MutableArrayRef privateBlockArgs = cast(*opInst).getPrivateBlockArgs(); + SmallVector mlirPrivateVars; SmallVector llvmPrivateVars; SmallVector privateDecls; + mlirPrivateVars.reserve(privateBlockArgs.size()); llvmPrivateVars.reserve(privateBlockArgs.size()); - privateDecls.reserve(privateBlockArgs.size()); collectPrivatizationDecls(opInst, privateDecls); + for (mlir::Value privateVar : opInst.getPrivateVars()) + mlirPrivateVars.push_back(privateVar); // Collect reduction declarations SmallVector reductionDecls; @@ -1895,8 +1900,8 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder, auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP) -> llvm::Error { llvm::Expected afterAllocas = allocatePrivateVars( - opInst, builder, moduleTranslation, privateBlockArgs, privateDecls, - llvmPrivateVars, allocaIP); + builder, moduleTranslation, privateBlockArgs, privateDecls, + mlirPrivateVars, llvmPrivateVars, allocaIP); if (handleError(afterAllocas, *opInst).failed()) return llvm::make_error();