diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h --- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h +++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h @@ -784,7 +784,7 @@ struct OutlineInfo { using PostOutlineCBTy = std::function; PostOutlineCBTy PostOutlineCB; - BasicBlock *EntryBB, *ExitBB; + BasicBlock *EntryBB, *ExitBB, *OuterAllocaBB; SmallVector ExcludeArgsFromAggregate; /// Collect all blocks in between EntryBB and ExitBB in both the given diff --git a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h --- a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h +++ b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h @@ -92,6 +92,11 @@ BranchProbabilityInfo *BPI; AssumptionCache *AC; + // A block outside of the extraction set where any intermediate + // allocations will be placed inside. If this is null, allocations + // will be placed in the entry block of the function. + BasicBlock *AllocationBlock; + // If true, varargs functions can be extracted. bool AllowVarArgs; @@ -123,8 +128,9 @@ CodeExtractor(ArrayRef BBs, DominatorTree *DT = nullptr, bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr, BranchProbabilityInfo *BPI = nullptr, - AssumptionCache *AC = nullptr, - bool AllowVarArgs = false, bool AllowAlloca = false, + AssumptionCache *AC = nullptr, bool AllowVarArgs = false, + bool AllowAlloca = false, + BasicBlock *AllocationBlock = nullptr, std::string Suffix = ""); /// Create a code extractor for a loop body. diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp @@ -300,6 +300,7 @@ /* AssumptionCache */ nullptr, /* AllowVarArgs */ true, /* AllowAlloca */ true, + /* AllocaBlock*/ OI.OuterAllocaBB, /* Suffix */ ".omp_par"); LLVM_DEBUG(dbgs() << "Before outlining: " << *OuterFn << "\n"); @@ -878,6 +879,7 @@ InsertPointTy PreFiniIP(PRegPreFiniBB, PRegPreFiniTI->getIterator()); FiniCB(PreFiniIP); + OI.OuterAllocaBB = OuterAllocaBlock; OI.EntryBB = PRegEntryBB; OI.ExitBB = PRegExitBB; @@ -901,6 +903,7 @@ /* AssumptionCache */ nullptr, /* AllowVarArgs */ true, /* AllowAlloca */ true, + /* AllocationBlock */ OuterAllocaBlock, /* Suffix */ ".omp_par"); // Find inputs to, outputs from the code region. diff --git a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp --- a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp +++ b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp @@ -352,7 +352,7 @@ // TODO: Pass BFI and BPI to update profile information. CodeExtractor CE(Region, &DT, /* AggregateArgs */ false, /* BFI */ nullptr, /* BPI */ nullptr, AC, /* AllowVarArgs */ false, - /* AllowAlloca */ false, + /* AllowAlloca */ false, /* AllocaBlock */ nullptr, /* Suffix */ "cold." + std::to_string(Count)); // Perform a simple cost/benefit analysis to decide whether or not to permit diff --git a/llvm/lib/Transforms/IPO/IROutliner.cpp b/llvm/lib/Transforms/IPO/IROutliner.cpp --- a/llvm/lib/Transforms/IPO/IROutliner.cpp +++ b/llvm/lib/Transforms/IPO/IROutliner.cpp @@ -2679,7 +2679,7 @@ OS->Candidate->getBasicBlocks(BlocksInRegion, BE); OS->CE = new (ExtractorAllocator.Allocate()) CodeExtractor(BE, nullptr, false, nullptr, nullptr, nullptr, false, - false, "outlined"); + false, nullptr, "outlined"); findAddInputsOutputs(M, *OS, NotSame); if (!OS->IgnoreRegion) OutlinedRegions.push_back(OS); @@ -2790,7 +2790,7 @@ OS->Candidate->getBasicBlocks(BlocksInRegion, BE); OS->CE = new (ExtractorAllocator.Allocate()) CodeExtractor(BE, nullptr, false, nullptr, nullptr, nullptr, false, - false, "outlined"); + false, nullptr, "outlined"); bool FunctionOutlined = extractSection(*OS); if (FunctionOutlined) { unsigned StartIdx = OS->Candidate->getStartIdx(); diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp --- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp +++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp @@ -246,9 +246,10 @@ bool AggregateArgs, BlockFrequencyInfo *BFI, BranchProbabilityInfo *BPI, AssumptionCache *AC, bool AllowVarArgs, bool AllowAlloca, - std::string Suffix) + BasicBlock *AllocationBlock, std::string Suffix) : DT(DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI), - BPI(BPI), AC(AC), AllowVarArgs(AllowVarArgs), + BPI(BPI), AC(AC), AllocationBlock(AllocationBlock), + AllowVarArgs(AllowVarArgs), Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs, AllowAlloca)), Suffix(Suffix) {} @@ -257,7 +258,7 @@ BranchProbabilityInfo *BPI, AssumptionCache *AC, std::string Suffix) : DT(&DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI), - BPI(BPI), AC(AC), AllowVarArgs(false), + BPI(BPI), AC(AC), AllocationBlock(nullptr), AllowVarArgs(false), Blocks(buildExtractionBlockSet(L.getBlocks(), &DT, /* AllowVarArgs */ false, /* AllowAlloca */ false)), @@ -1188,9 +1189,10 @@ // Allocate a struct at the beginning of this function StructArgTy = StructType::get(newFunction->getContext(), ArgTypes); - Struct = new AllocaInst(StructArgTy, DL.getAllocaAddrSpace(), nullptr, - "structArg", - &codeReplacer->getParent()->front().front()); + Struct = new AllocaInst( + StructArgTy, DL.getAllocaAddrSpace(), nullptr, "structArg", + AllocationBlock ? &AllocationBlock->front() + : &codeReplacer->getParent()->front().front()); params.push_back(Struct); // Store aggregated inputs in the struct. diff --git a/mlir/test/Target/LLVMIR/openmp-nested.mlir b/mlir/test/Target/LLVMIR/openmp-nested.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Target/LLVMIR/openmp-nested.mlir @@ -0,0 +1,41 @@ +// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s + +module { + llvm.func @printf(!llvm.ptr, ...) -> i32 + llvm.mlir.global internal constant @str0("WG size of kernel = %d X %d\0A\00") + + llvm.func @main(%arg0: i32, %arg1: !llvm.ptr>) -> i32 { + omp.parallel { + %0 = llvm.mlir.constant(1 : index) : i64 + %1 = llvm.mlir.constant(10 : index) : i64 + %2 = llvm.mlir.constant(0 : index) : i64 + %4 = llvm.mlir.constant(0 : i32) : i32 + %12 = llvm.alloca %0 x i64 : (i64) -> !llvm.ptr + omp.wsloop (%arg2) : i64 = (%2) to (%1) step (%0) { + omp.parallel { + omp.wsloop (%arg3) : i64 = (%2) to (%0) step (%0) { + llvm.store %2, %12 : !llvm.ptr + omp.yield + } + omp.terminator + } + %19 = llvm.load %12 : !llvm.ptr + %20 = llvm.trunc %19 : i64 to i32 + %5 = llvm.mlir.addressof @str0 : !llvm.ptr> + %6 = llvm.getelementptr %5[%4, %4] : (!llvm.ptr>, i32, i32) -> !llvm.ptr + %21 = llvm.call @printf(%6, %20, %20) : (!llvm.ptr, i32, i32) -> i32 + omp.yield + } + omp.terminator + } + %a4 = llvm.mlir.constant(0 : i32) : i32 + llvm.return %a4 : i32 + } + +} + +// CHECK: call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* @1, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @[[inner1:.+]] to void (i32*, i32*, ...)*)) + +// CHECK: define internal void @[[inner1]] +// CHECK: %[[structArg:.+]] = alloca { i64* } +// CHECK: call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* @3, i32 1, void (i32*, i32*, ...)* bitcast (void (i32*, i32*, { i64* }*)* @[[inner2:.+]] to void (i32*, i32*, ...)*), { i64* }* %[[structArg]])