Move llvm private vars into task context struct manager

Author: tblah

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -1428,7 +1428,7 @@ static LogicalResult
 initFirstPrivateVars(llvm::IRBuilderBase &builder,
                      LLVM::ModuleTranslation &moduleTranslation,
                      SmallVectorImpl<mlir::Value> &mlirPrivateVars,
-                     SmallVectorImpl<llvm::Value *> &llvmPrivateVars,
+                     ArrayRef<llvm::Value *> llvmPrivateVars,
                      SmallVectorImpl<omp::PrivateClauseOp> &privateDecls,
                      llvm::BasicBlock *afterAllocas) {
   llvm::IRBuilderBase::InsertPointGuard guard(builder);
@@ -1751,11 +1751,15 @@ class TaskContextStructManager {
   /// variable, adding them to llvmPrivateVars. Null values are added where
   /// private decls were skipped so that the ordering continues to match the
   /// private decls.
-  void createGEPsToPrivateVars(SmallVectorImpl<llvm::Value *> &llvmPrivateVars);
+  void createGEPsToPrivateVars();
 
   /// De-allocate the task context structure.
   void freeStructPtr();
 
+  MutableArrayRef<llvm::Value *> getLLVMPrivateVars() {
+    return llvmPrivateVars;
+  }
+
   llvm::Value *getStructPtr() { return structPtr; }
 
 private:
@@ -1766,6 +1770,10 @@ class TaskContextStructManager {
   /// The type of each member of the structure, in order.
   SmallVector<llvm::Type *> privateVarTypes;
 
+  /// LLVM values for each private variable, or null if that private variable is
+  /// not included in the task context structure
+  SmallVector<llvm::Value *> llvmPrivateVars;
+
   /// A pointer to the structure containing context for this task.
   llvm::Value *structPtr = nullptr;
   /// The type of the structure
@@ -1801,14 +1809,14 @@ void TaskContextStructManager::generateTaskContextStruct() {
                                    "omp.task.context_ptr");
 }
 
-void TaskContextStructManager::createGEPsToPrivateVars(
-    SmallVectorImpl<llvm::Value *> &llvmPrivateVars) {
+void TaskContextStructManager::createGEPsToPrivateVars() {
   if (!structPtr) {
     assert(privateVarTypes.empty());
     return;
   }
 
   // Create GEPs for each struct member and initialize llvmPrivateVars to point
+  llvmPrivateVars.clear();
   llvmPrivateVars.reserve(privateVarTypes.size());
   llvm::Value *zero = builder.getInt32(0);
   unsigned i = 0;
@@ -1904,19 +1912,18 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
   builder.SetInsertPoint(initBlock->getTerminator());
 
   // Create task variable structure
-  llvm::SmallVector<llvm::Value *> privateVarAllocations;
   taskStructMgr.generateTaskContextStruct();
   // GEPs so that we can initialize the variables. Don't use these GEPs inside
   // of the body otherwise it will be the GEP not the struct which is fowarded
   // to the outlined function. GEPs forwarded in this way are passed in a
   // stack-allocated (by OpenMPIRBuilder) structure which is not safe for tasks
   // which may not be executed until after the current stack frame goes out of
   // scope.
-  taskStructMgr.createGEPsToPrivateVars(privateVarAllocations);
+  taskStructMgr.createGEPsToPrivateVars();
 
   for (auto [privDecl, mlirPrivVar, blockArg, llvmPrivateVarAlloc] :
        llvm::zip_equal(privateDecls, mlirPrivateVars, privateBlockArgs,
-                       privateVarAllocations)) {
+                       taskStructMgr.getLLVMPrivateVars())) {
     if (!privDecl.readsFromMold())
       // to be handled inside the task
       continue;
@@ -1950,8 +1957,8 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
 
   // firstprivate copy region
   if (failed(initFirstPrivateVars(builder, moduleTranslation, mlirPrivateVars,
-                                  privateVarAllocations, privateDecls,
-                                  copyBlock)))
+                                  taskStructMgr.getLLVMPrivateVars(),
+                                  privateDecls, copyBlock)))
     return llvm::failure();
 
   // Set up for call to createTask()
@@ -1992,7 +1999,9 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
 
     // Find and map the addresses of each variable within the task context
     // structure
-    taskStructMgr.createGEPsToPrivateVars(llvmPrivateVars);
+    taskStructMgr.createGEPsToPrivateVars();
+    llvm::copy(taskStructMgr.getLLVMPrivateVars(),
+               std::back_inserter(llvmPrivateVars));
     for (auto [blockArg, llvmPrivateVar] :
          llvm::zip_equal(privateBlockArgs, llvmPrivateVars)) {
       if (!llvmPrivateVar)