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

[Dominik Adamski via llvm-commits]

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

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.

>From 8756f392e94d378f3a9af744ff77f77f4895e72f Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Mon, 11 Sep 2023 05:13:13 -0400
Subject: [PATCH] [CodeExtractor] Allow to specify address space for aggregate
 argument structure

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.
---
 .../llvm/Transforms/Utils/CodeExtractor.h     | 11 +++++++++-
 llvm/lib/Transforms/Utils/CodeExtractor.cpp   | 21 ++++++++++++-------
 2 files changed, 23 insertions(+), 9 deletions(-)

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);