[llvm] [CodeExtractor][OpenMP] Allow to specify address space for aggregate arguments structure (PR #66998)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-openmp

<details>
<summary>Changes</summary>

The user of CodeExtractor should be able to specify the alloca address space for aggregate arguments object.

The information about alloca address space for aggregate arguments object is used to define type of outlined function arguments.

CodeExtractor is used to generate outlined functions required by OpenMP runtime. The arguments of the outlined functions for OpenMP GPU code are in 0 address space. 0 address space does not need to be the default address space for GPU device. That's why there is a need to allow user of CodeExtractor to specify the address space for aggregate args structure.

---
Full diff: https://github.com/llvm/llvm-project/pull/66998.diff


2 Files Affected:

- (modified) llvm/include/llvm/Transforms/Utils/CodeExtractor.h (+10-1) 
- (modified) llvm/lib/Transforms/Utils/CodeExtractor.cpp (+13-8) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
index bb23cf4a9a3cbbb..c610a4c0bc93749 100644
--- a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
+++ b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
@@ -114,6 +114,11 @@ class CodeExtractorAnalysisCache {
     // label, if non-empty, otherwise "extracted".
     std::string Suffix;
 
+    // Optional value of alloca address space which should be used for
+    // allocation of structure which aggregates arguments. If empty, default
+    // alloca address space for given data layout is used.
+    std::optional<unsigned> AggregateArgsAllocaAddrSpace;
+
   public:
     /// Create a code extractor for a sequence of blocks.
     ///
@@ -128,13 +133,17 @@ class CodeExtractorAnalysisCache {
     /// Any new allocations will be placed in the AllocationBlock, unless
     /// it is null, in which case it will be placed in the entry block of
     /// the function from which the code is being extracted.
+    /// AggregateArgsAllocaAddrSpace param specifies the alloca address space
+    /// where aggregate argument structure is allocated. If it's not specified,
+    /// then the default alloca address space is used.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr,
                   AssumptionCache *AC = nullptr, bool AllowVarArgs = false,
                   bool AllowAlloca = false,
                   BasicBlock *AllocationBlock = nullptr,
-                  std::string Suffix = "");
+                  std::string Suffix = "",
+                  std::optional<unsigned> AggregateArgsAllocaAddrSpace = {});
 
     /// Create a code extractor for a loop body.
     ///
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index cafa99491f5b5f6..e03f31d16526910 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -241,16 +241,17 @@ buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
   return Result;
 }
 
-CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
-                             bool AggregateArgs, BlockFrequencyInfo *BFI,
-                             BranchProbabilityInfo *BPI, AssumptionCache *AC,
-                             bool AllowVarArgs, bool AllowAlloca,
-                             BasicBlock *AllocationBlock, std::string Suffix)
+CodeExtractor::CodeExtractor(
+    ArrayRef<BasicBlock *> BBs, DominatorTree *DT, bool AggregateArgs,
+    BlockFrequencyInfo *BFI, BranchProbabilityInfo *BPI, AssumptionCache *AC,
+    bool AllowVarArgs, bool AllowAlloca, BasicBlock *AllocationBlock,
+    std::string Suffix, std::optional<unsigned> AggregateArgsAllocaAddrSpace)
     : DT(DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
       BPI(BPI), AC(AC), AllocationBlock(AllocationBlock),
       AllowVarArgs(AllowVarArgs),
       Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs, AllowAlloca)),
-      Suffix(Suffix) {}
+      Suffix(Suffix),
+      AggregateArgsAllocaAddrSpace(AggregateArgsAllocaAddrSpace) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs,
                              BlockFrequencyInfo *BFI,
@@ -866,7 +867,9 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   StructType *StructTy = nullptr;
   if (AggregateArgs && !AggParamTy.empty()) {
     StructTy = StructType::get(M->getContext(), AggParamTy);
-    ParamTy.push_back(PointerType::get(StructTy, DL.getAllocaAddrSpace()));
+    ParamTy.push_back(PointerType::get(
+        StructTy,
+        AggregateArgsAllocaAddrSpace.value_or(DL.getAllocaAddrSpace())));
   }
 
   LLVM_DEBUG({
@@ -1183,7 +1186,9 @@ CallInst *CodeExtractor::emitCallAndSwitchStatement(Function *newFunction,
     // Allocate a struct at the beginning of this function
     StructArgTy = StructType::get(newFunction->getContext(), ArgTypes);
     Struct = new AllocaInst(
-        StructArgTy, DL.getAllocaAddrSpace(), nullptr, "structArg",
+        StructArgTy,
+        AggregateArgsAllocaAddrSpace.value_or(DL.getAllocaAddrSpace()), nullptr,
+        "structArg",
         AllocationBlock ? &*AllocationBlock->getFirstInsertionPt()
                         : &codeReplacer->getParent()->front().front());
     params.push_back(Struct);

``````````

</details>


https://github.com/llvm/llvm-project/pull/66998