[OpenACC] Implement 'firstprivate' clause copy lowering (#154150)

This patch is the last of the 'firstprivate' clause lowering patches. It
takes the already generated 'copy' init from Sema and uses it to
generate the IR for the copy section of the recipe.

However, one thing that this patch had to do, was come up with a way to
hijack the decl registration in CIRGenFunction. Because these decls are
being created in a 'different' place, we need to remove the things we've
added. We could alternatively generate these 'differently', but it seems
worth a little extra effort here to avoid having to re-implement
variable initialization.

Author: erichkeane

clang/lib/CIR/CodeGen/CIRGenDecl.cpp

@@ -154,15 +154,19 @@ void CIRGenFunction::emitAutoVarInit(
     initializeWhatIsTechnicallyUninitialized(addr);
     LValue lv = makeAddrLValue(addr, type, AlignmentSource::Decl);
     emitExprAsInit(init, &d, lv);
-    // In case lv has uses it means we indeed initialized something
-    // out of it while trying to build the expression, mark it as such.
-    mlir::Value val = lv.getAddress().getPointer();
-    assert(val && "Should have an address");
-    auto allocaOp = val.getDefiningOp<cir::AllocaOp>();
-    assert(allocaOp && "Address should come straight out of the alloca");
-
-    if (!allocaOp.use_empty())
-      allocaOp.setInitAttr(mlir::UnitAttr::get(&getMLIRContext()));
+
+    if (!emission.wasEmittedAsOffloadClause()) {
+      // In case lv has uses it means we indeed initialized something
+      // out of it while trying to build the expression, mark it as such.
+      mlir::Value val = lv.getAddress().getPointer();
+      assert(val && "Should have an address");
+      auto allocaOp = val.getDefiningOp<cir::AllocaOp>();
+      assert(allocaOp && "Address should come straight out of the alloca");
+
+      if (!allocaOp.use_empty())
+        allocaOp.setInitAttr(mlir::UnitAttr::get(&getMLIRContext()));
+    }
+
     return;
   }
 

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -471,6 +471,10 @@ class CIRGenFunction : public CIRGenTypeCache {
     /// escaping block.
     bool IsEscapingByRef = false;
 
+    /// True if the variable was emitted as an offload recipe, and thus doesn't
+    /// have the same sort of alloca initialization.
+    bool EmittedAsOffload = false;
+
     mlir::Value NRVOFlag{};
 
     struct Invalid {};
@@ -483,11 +487,18 @@ class CIRGenFunction : public CIRGenTypeCache {
 
     bool wasEmittedAsGlobal() const { return !Addr.isValid(); }
 
+    bool wasEmittedAsOffloadClause() const { return EmittedAsOffload; }
+
     /// Returns the raw, allocated address, which is not necessarily
     /// the address of the object itself. It is casted to default
     /// address space for address space agnostic languages.
     Address getAllocatedAddress() const { return Addr; }
 
+    // Changes the stored address for the emission.  This function should only
+    // be used in extreme cases, and isn't required to model normal AST
+    // initialization/variables.
+    void setAllocatedAddress(Address A) { Addr = A; }
+
     /// Returns the address of the object within this declaration.
     /// Note that this does not chase the forwarding pointer for
     /// __block decls.
@@ -519,12 +530,34 @@ class CIRGenFunction : public CIRGenTypeCache {
     symbolTable.insert(vd, addr.getPointer());
   }
 
-  /// Removes a declaration from the address-relationship.  This is a function
-  /// that shouldn't need to be used except in cases where we're adding/removing
-  /// things that aren't part of the language-semantics AST.
-  void removeAddrOfLocalVar(const clang::VarDecl *vd) {
-    localDeclMap.erase(vd);
-  }
+  // A class to allow reverting changes to a var-decl's registration to the
+  // localDeclMap. This is used in cases where things are being inserted into
+  // the variable list but don't follow normal lookup/search rules, like in
+  // OpenACC recipe generation.
+  class DeclMapRevertingRAII {
+    CIRGenFunction &cgf;
+    const VarDecl *vd;
+    bool shouldDelete = false;
+    Address oldAddr = Address::invalid();
+
+  public:
+    DeclMapRevertingRAII(CIRGenFunction &cgf, const VarDecl *vd)
+        : cgf(cgf), vd(vd) {
+      auto mapItr = cgf.localDeclMap.find(vd);
+
+      if (mapItr != cgf.localDeclMap.end())
+        oldAddr = mapItr->second;
+      else
+        shouldDelete = true;
+    }
+
+    ~DeclMapRevertingRAII() {
+      if (shouldDelete)
+        cgf.localDeclMap.erase(vd);
+      else
+        cgf.localDeclMap.insert_or_assign(vd, oldAddr);
+    }
+  };
 
   bool shouldNullCheckClassCastValue(const CastExpr *ce);
 
@@ -970,6 +1003,11 @@ class CIRGenFunction : public CIRGenTypeCache {
   void emitAutoVarDecl(const clang::VarDecl &d);
 
   void emitAutoVarCleanups(const AutoVarEmission &emission);
+  /// Emit the initializer for an allocated variable.  If this call is not
+  /// associated with the call to emitAutoVarAlloca (as the address of the
+  /// emission is not directly an alloca), the allocatedSeparately parameter can
+  /// be used to suppress the assertions.  However, this should only be used in
+  /// extreme cases, as it doesn't properly reflect the language/AST.
   void emitAutoVarInit(const AutoVarEmission &emission);
   void emitAutoVarTypeCleanup(const AutoVarEmission &emission,
                               clang::QualType::DestructionKind dtorKind);

clang/lib/CIR/CodeGen/CIRGenOpenACCClause.cpp

@@ -392,12 +392,29 @@ class OpenACCClauseCIREmitter final
       CIRGenFunction::AutoVarEmission tempDeclEmission,
       mlir::acc::FirstprivateRecipeOp recipe, const VarDecl *varRecipe,
       const VarDecl *temporary) {
-    builder.createBlock(&recipe.getCopyRegion(), recipe.getCopyRegion().end(),
-                        {mainOp.getType(), mainOp.getType()}, {loc, loc});
+    mlir::Block *block = builder.createBlock(
+        &recipe.getCopyRegion(), recipe.getCopyRegion().end(),
+        {mainOp.getType(), mainOp.getType()}, {loc, loc});
     builder.setInsertionPointToEnd(&recipe.getCopyRegion().back());
 
-    // TODO: OpenACC: Implement this copy to actually do something.
+    mlir::BlockArgument fromArg = block->getArgument(0);
+    mlir::BlockArgument toArg = block->getArgument(1);
+
+    mlir::Type elementTy =
+        mlir::cast<cir::PointerType>(mainOp.getType()).getPointee();
+
+    // Set the address of the emission to be the argument, so that we initialize
+    // that instead of the variable in the other block.
+    tempDeclEmission.setAllocatedAddress(
+        Address{toArg, elementTy, cgf.getContext().getDeclAlign(varRecipe)});
+    tempDeclEmission.EmittedAsOffload = true;
+
+    CIRGenFunction::DeclMapRevertingRAII declMapRAII{cgf, temporary};
+    cgf.setAddrOfLocalVar(
+        temporary,
+        Address{fromArg, elementTy, cgf.getContext().getDeclAlign(varRecipe)});
 
+    cgf.emitAutoVarInit(tempDeclEmission);
     mlir::acc::YieldOp::create(builder, locEnd);
   }
 
@@ -417,6 +434,7 @@ class OpenACCClauseCIREmitter final
 
     CIRGenFunction::AutoVarEmission tempDeclEmission{
         CIRGenFunction::AutoVarEmission::invalid()};
+    CIRGenFunction::DeclMapRevertingRAII declMapRAII{cgf, varRecipe};
 
     // Do the 'init' section of the recipe IR, which does an alloca, then the
     // initialization (except for firstprivate).
@@ -425,6 +443,7 @@ class OpenACCClauseCIREmitter final
     builder.setInsertionPointToEnd(&recipe.getInitRegion().back());
     tempDeclEmission =
         cgf.emitAutoVarAlloca(*varRecipe, builder.saveInsertionPoint());
+
     // 'firstprivate' doesn't do its initialization in the 'init' section,
     // instead does it in the 'copy' section.  SO only do init here.
     // 'reduction' appears to use it too (rather than a 'copy' section), so
@@ -450,16 +469,19 @@ class OpenACCClauseCIREmitter final
     mlir::acc::YieldOp::create(builder, locEnd);
 
     if constexpr (std::is_same_v<RecipeTy, mlir::acc::FirstprivateRecipeOp>) {
-      // TODO: OpenACC: we should have a errorNYI call here if
-      // !varRecipe->getInit(), but as that generation isn't currently
-      // implemented, it ends up being too noisy. So when we implement copy-init
-      // generation both in Sema and here, we should have a diagnostic here.
+      if (!varRecipe->getInit()) {
+        // If we don't have any initialization recipe, we failed during Sema to
+        // initialize this correctly. If we disable the
+        // Sema::TentativeAnalysisScopes in SemaOpenACC::CreateInitRecipe, it'll
+        // emit an error to tell us.  However, emitting those errors during
+        // production is a violation of the standard, so we cannot do them.
+        cgf.cgm.errorNYI(
+            exprRange, "firstprivate copy-init recipe not properly generated");
+      }
+
       createFirstprivateRecipeCopy(loc, locEnd, mainOp, tempDeclEmission,
                                    recipe, varRecipe, temporary);
     }
-
-    // Make sure we cleanup after ourselves here.
-    cgf.removeAddrOfLocalVar(varRecipe);
   }
 
   void createRecipeDestroySection(mlir::Location loc, mlir::Location locEnd,