[OpenMP 6.0 ]Codegen for Reduction over private variables with reduction clause

clang/lib/CodeGen/CGOpenMPRuntime.cpp

@@ -4899,6 +4899,238 @@ void CGOpenMPRuntime::emitSingleReductionCombiner(CodeGenFunction &CGF,
   }
 }
 
+void CGOpenMPRuntime::emitPrivateReduction(
+    CodeGenFunction &CGF, SourceLocation Loc, ArrayRef<const Expr *> Privates,
+    ArrayRef<const Expr *> LHSExprs, ArrayRef<const Expr *> RHSExprs,
+    ArrayRef<const Expr *> ReductionOps) {
+  if (LHSExprs.empty() || Privates.empty() || ReductionOps.empty())
+    return;
+
+  if (LHSExprs.size() != Privates.size() ||
+      LHSExprs.size() != ReductionOps.size())
+    return;
+
+  QualType PrivateType = Privates[0]->getType();
+  llvm::Type *LLVMType = CGF.ConvertTypeForMem(PrivateType);
+
+  llvm::Constant *InitVal = nullptr;
+  const OMPDeclareReductionDecl *UDR = getReductionInit(ReductionOps[0]);
+
+  if (!UDR) {
+    InitVal = llvm::Constant::getNullValue(LLVMType);
+    if (const auto *DRE = dyn_cast<DeclRefExpr>(Privates[0])) {
+      if (const auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
+        const Expr *InitExpr = VD->getInit();
+        if (InitExpr && !PrivateType->isAggregateType() &&
+            !PrivateType->isAnyComplexType()) {
+          Expr::EvalResult Result;
+          if (InitExpr->EvaluateAsRValue(Result, CGF.getContext())) {
+            APValue &InitValue = Result.Val;
+            if (InitValue.isInt()) {
+              InitVal = llvm::ConstantInt::get(LLVMType, InitValue.getInt());
+            }
+          }
+        }
+      }
+    }
+  } else {
+    InitVal = llvm::Constant::getNullValue(LLVMType);
+  }
+
+  // Create an internal shared variable
+  std::string SharedName = getName({"internal_private_var"});
+  llvm::GlobalVariable *SharedVar = new llvm::GlobalVariable(
+      CGM.getModule(), LLVMType, false, llvm::GlobalValue::CommonLinkage,
+      InitVal, ".omp.reduction." + SharedName, nullptr,
+      llvm::GlobalVariable::NotThreadLocal);
+
+  SharedVar->setAlignment(
+      llvm::MaybeAlign(CGF.getContext().getTypeAlign(PrivateType) / 8));
+
+  Address SharedResult(SharedVar, SharedVar->getValueType(),
+                       CGF.getContext().getTypeAlignInChars(PrivateType));
+
+  llvm::Value *ThreadId = getThreadID(CGF, Loc);
+  llvm::Value *BarrierLoc = emitUpdateLocation(CGF, Loc, OMP_ATOMIC_REDUCE);
+  llvm::Value *BarrierArgs[] = {BarrierLoc, ThreadId};
+
+  CGF.EmitRuntimeCall(OMPBuilder.getOrCreateRuntimeFunction(
+                          CGM.getModule(), OMPRTL___kmpc_barrier),
+                      BarrierArgs);
+
+  llvm::BasicBlock *InitBB = CGF.createBasicBlock("init");
+  llvm::BasicBlock *InitEndBB = CGF.createBasicBlock("init.end");
+
+  llvm::Value *IsWorker = CGF.Builder.CreateICmpEQ(
+      ThreadId, llvm::ConstantInt::get(ThreadId->getType(), 0));
+  CGF.Builder.CreateCondBr(IsWorker, InitBB, InitEndBB);
+
+  CGF.EmitBlock(InitBB);
+
+  if (const auto *DRE = dyn_cast<DeclRefExpr>(Privates[0])) {
+    if (const auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
+      const Expr *InitExpr = VD->getInit();
+      if (InitExpr &&
+          (PrivateType->isAggregateType() || PrivateType->isAnyComplexType())) {
+        CGF.EmitAnyExprToMem(InitExpr, SharedResult,
+                             PrivateType.getQualifiers(), true);
+      } else if (!InitVal->isNullValue()) {
+        CGF.EmitStoreOfScalar(InitVal,
+                              CGF.MakeAddrLValue(SharedResult, PrivateType));
+      } else {
+        CGF.EmitNullInitialization(SharedResult, PrivateType);
+      }
+    } else {
+      CGF.EmitNullInitialization(SharedResult, PrivateType);
+    }
+  } else {
+    CGF.EmitNullInitialization(SharedResult, PrivateType);
+  }
+
+  CGF.Builder.CreateBr(InitEndBB);
+  CGF.EmitBlock(InitEndBB);
+  CGF.EmitRuntimeCall(OMPBuilder.getOrCreateRuntimeFunction(
+                          CGM.getModule(), OMPRTL___kmpc_barrier),
+                      BarrierArgs);
+  for (unsigned I :
+       llvm::seq<unsigned>(std::min(ReductionOps.size(), LHSExprs.size()))) {
+
+    const Expr *ReductionOp = ReductionOps[I];
+    LValue SharedLV = CGF.MakeAddrLValue(SharedResult, PrivateType);
+    LValue LHSLV = CGF.EmitLValue(LHSExprs[I]);
+    // If UDR
+    const OMPDeclareReductionDecl *CurrentUDR =
+        getReductionInit(ReductionOps[I]);
+    if (CurrentUDR) {
+      auto ReductionGen = [&](CodeGenFunction &CGF, PrePostActionTy &Action) {
+        Action.Enter(CGF);
+        std::pair<llvm::Function *, llvm::Function *> ReductionFnPair =
+            getUserDefinedReduction(CurrentUDR);
+        llvm::Function *CombinerFn = ReductionFnPair.first;
+        if (const auto *CE = dyn_cast<CallExpr>(ReductionOp)) {
+          if (CE && CombinerFn) {
+            const auto *CE = cast<CallExpr>(ReductionOps[I]);
+            const auto *OutDRE = cast<DeclRefExpr>(
+                cast<UnaryOperator>(CE->getArg(0)->IgnoreParenImpCasts())
+                    ->getSubExpr());
+            const auto *InDRE = cast<DeclRefExpr>(
+                cast<UnaryOperator>(CE->getArg(1)->IgnoreParenImpCasts())
+                    ->getSubExpr());
+            CodeGenFunction::OMPPrivateScope LocalScope(CGF);
+            LocalScope.addPrivate(cast<VarDecl>(OutDRE->getDecl()),
+                                  SharedLV.getAddress());
+            LocalScope.addPrivate(cast<VarDecl>(InDRE->getDecl()),
+                                  LHSLV.getAddress());
+            (void)LocalScope.Privatize();
+            emitReductionCombiner(CGF, ReductionOp);
+          }
+        }
+      };
+      std::string CriticalName = getName({"reduction_critical"});
+      emitCriticalRegion(CGF, CriticalName, ReductionGen, Loc);
+    } else {
+      // Built-in Operator Combination
+      const Expr *ReductionClauseExpr = ReductionOp->IgnoreParenCasts();
+      if (const auto *Cleanup = dyn_cast<ExprWithCleanups>(ReductionClauseExpr))
+        ReductionClauseExpr = Cleanup->getSubExpr()->IgnoreParenCasts();
+      const Expr *AssignRHS = nullptr;
+      if (const auto *BinOp = dyn_cast<BinaryOperator>(ReductionClauseExpr)) {
+        if (BinOp->getOpcode() == BO_Assign)
+          AssignRHS = BinOp->getRHS();
+      } else if (const auto *OpCall =
+                     dyn_cast<CXXOperatorCallExpr>(ReductionClauseExpr)) {
+        if (OpCall->getOperator() == OO_Equal)
+          AssignRHS = OpCall->getArg(1);
+      }
+      if (!AssignRHS)
+        continue;
+      const Expr *ReductionCombinerExpr = AssignRHS->IgnoreParenImpCasts();
+      if (const auto *MTE =
+              dyn_cast<MaterializeTemporaryExpr>(ReductionCombinerExpr))
+        ReductionCombinerExpr = MTE->getSubExpr()->IgnoreParenImpCasts();
+
+      BinaryOperatorKind BO = BO_Assign;
+      RValue PrivateRV = CGF.EmitLoadOfLValue(LHSLV, Loc);
+      if (const auto *BinOp = dyn_cast<BinaryOperator>(ReductionCombinerExpr)) {
+        BO = BinOp->getOpcode();
+        auto UpdateOp = [&](RValue OldVal) {
+          if (BO == BO_Mul) {
+            llvm::Value *OldScalar = OldVal.getScalarVal();
+            llvm::Value *PrivateScalar = PrivateRV.getScalarVal();
+            llvm::Value *Result =
+                CGF.Builder.CreateMul(OldScalar, PrivateScalar);
+            return RValue::get(Result);
+          } else {
+            OpaqueValueExpr OVE(BinOp->getLHS()->getExprLoc(),
+                                BinOp->getLHS()->getType(),
+                                ExprValueKind::VK_PRValue);
+            CodeGenFunction::OpaqueValueMapping OldValMapping(CGF, &OVE,
+                                                              OldVal);
+            return CGF.EmitAnyExpr(BinOp->getRHS());
+          }
+        };
+
+        (void)CGF.EmitOMPAtomicSimpleUpdateExpr(
+            SharedLV, PrivateRV, BO, true,
+            llvm::AtomicOrdering::SequentiallyConsistent, Loc, UpdateOp);
+      } else if (dyn_cast<CXXOperatorCallExpr>(ReductionClauseExpr)) {
+        // Implicit Reduction Identifiers ( openmp 6.0 sec 7.6.5 )
+        auto ReductionGen = [&](CodeGenFunction &CGF, PrePostActionTy &Action) {
+          Action.Enter(CGF);
+          const auto *OmpOutDRE =
+              dyn_cast<DeclRefExpr>(LHSExprs[I]->IgnoreParenImpCasts());
+          const auto *OmpInDRE =
+              dyn_cast<DeclRefExpr>(RHSExprs[I]->IgnoreParenImpCasts());
+
+          if (!OmpOutDRE || !OmpInDRE) {
+            return;
+          }
+          const VarDecl *OmpOutVD = cast<VarDecl>(OmpOutDRE->getDecl());
+          const VarDecl *OmpInVD = cast<VarDecl>(OmpInDRE->getDecl());
+          CodeGenFunction::OMPPrivateScope LocalScope(CGF);
+          LocalScope.addPrivate(OmpOutVD, SharedLV.getAddress());
+          LocalScope.addPrivate(OmpInVD, LHSLV.getAddress());
+          (void)LocalScope.Privatize();
+          CGF.EmitIgnoredExpr(ReductionOp);
+        };
+        std::string CriticalName = getName({"reduction_critical"});
+        emitCriticalRegion(CGF, CriticalName, ReductionGen, Loc);
+      }
+    }
+  }
+  // Final barrier
+  CGF.EmitRuntimeCall(OMPBuilder.getOrCreateRuntimeFunction(
+                          CGM.getModule(), OMPRTL___kmpc_barrier),
+                      BarrierArgs);
+
+  // Broadcast final result
+  bool IsAggregate = PrivateType->isAggregateType();
+  llvm::Value *FinalResultVal = nullptr;
+  LValue SharedLV = CGF.MakeAddrLValue(SharedResult, PrivateType);
+  Address FinalResultAddr = Address::invalid();
+  if (IsAggregate) {
+    FinalResultAddr = SharedResult;
+  } else {
+    FinalResultVal = CGF.EmitLoadOfScalar(SharedLV, Loc);
+  }
+
+  for (unsigned I : llvm::seq<unsigned>(Privates.size())) {
+    LValue TargetLHSLV = CGF.EmitLValue(LHSExprs[I]);
+    if (IsAggregate) {
+      CGF.EmitAggregateCopy(TargetLHSLV,
+                            CGF.MakeAddrLValue(FinalResultAddr, PrivateType),
+                            PrivateType, AggValueSlot::DoesNotOverlap, false);
+    } else {
+      CGF.EmitStoreOfScalar(FinalResultVal, TargetLHSLV);
+    }
+  }
+
+  // Final synchronization
+  CGF.EmitRuntimeCall(OMPBuilder.getOrCreateRuntimeFunction(
+                          CGM.getModule(), OMPRTL___kmpc_barrier),
+                      BarrierArgs);
+}
+
 void CGOpenMPRuntime::emitReduction(CodeGenFunction &CGF, SourceLocation Loc,
                                     ArrayRef<const Expr *> Privates,
                                     ArrayRef<const Expr *> LHSExprs,
@@ -5201,6 +5433,8 @@ void CGOpenMPRuntime::emitReduction(CodeGenFunction &CGF, SourceLocation Loc,
 
   CGF.EmitBranch(DefaultBB);
   CGF.EmitBlock(DefaultBB, /*IsFinished=*/true);
+  if (Options.IsPrivateVarReduction)
+    emitPrivateReduction(CGF, Loc, Privates, LHSExprs, RHSExprs, ReductionOps);
 }
 
 /// Generates unique name for artificial threadprivate variables.

clang/lib/CodeGen/CGOpenMPRuntime.h

@@ -1201,8 +1201,22 @@ class CGOpenMPRuntime {
   struct ReductionOptionsTy {
     bool WithNowait;
     bool SimpleReduction;
+    bool IsPrivateVarReduction;
     OpenMPDirectiveKind ReductionKind;
   };
+
+  /// Emits code for private variable reduction
+  /// \param Privates List of private copies for original reduction arguments.
+  /// \param LHSExprs List of LHS in \a ReductionOps reduction operations.
+  /// \param RHSExprs List of RHS in \a ReductionOps reduction operations.
+  /// \param ReductionOps List of reduction operations in form 'LHS binop RHS'
+  /// or 'operator binop(LHS, RHS)'.
+  void emitPrivateReduction(CodeGenFunction &CGF, SourceLocation Loc,
+                            ArrayRef<const Expr *> Privates,
+                            ArrayRef<const Expr *> LHSExprs,
+                            ArrayRef<const Expr *> RHSExprs,
+                            ArrayRef<const Expr *> ReductionOps);
+
   /// Emit a code for reduction clause. Next code should be emitted for
   /// reduction:
   /// \code

clang/lib/CodeGen/CGStmtOpenMP.cpp

@@ -1470,6 +1470,7 @@ void CodeGenFunction::EmitOMPReductionClauseFinal(
   llvm::SmallVector<const Expr *, 8> LHSExprs;
   llvm::SmallVector<const Expr *, 8> RHSExprs;
   llvm::SmallVector<const Expr *, 8> ReductionOps;
+  llvm::SmallVector<bool, 8> IsPrivate;
   bool HasAtLeastOneReduction = false;
   bool IsReductionWithTaskMod = false;
   for (const auto *C : D.getClausesOfKind<OMPReductionClause>()) {
@@ -1480,6 +1481,8 @@ void CodeGenFunction::EmitOMPReductionClauseFinal(
     Privates.append(C->privates().begin(), C->privates().end());
     LHSExprs.append(C->lhs_exprs().begin(), C->lhs_exprs().end());
     RHSExprs.append(C->rhs_exprs().begin(), C->rhs_exprs().end());
+    IsPrivate.append(C->private_var_reduction_flags().begin(),
+                     C->private_var_reduction_flags().end());
     ReductionOps.append(C->reduction_ops().begin(), C->reduction_ops().end());
     IsReductionWithTaskMod =
         IsReductionWithTaskMod || C->getModifier() == OMPC_REDUCTION_task;
@@ -1499,9 +1502,11 @@ void CodeGenFunction::EmitOMPReductionClauseFinal(
     bool SimpleReduction = ReductionKind == OMPD_simd;
     // Emit nowait reduction if nowait clause is present or directive is a
     // parallel directive (it always has implicit barrier).
+    bool IsPrivateVarReduction =
+        llvm::any_of(IsPrivate, [](bool IsPriv) { return IsPriv; });
     CGM.getOpenMPRuntime().emitReduction(
         *this, D.getEndLoc(), Privates, LHSExprs, RHSExprs, ReductionOps,
-        {WithNowait, SimpleReduction, ReductionKind});
+        {WithNowait, SimpleReduction, IsPrivateVarReduction, ReductionKind});
   }
 }
 
@@ -3943,7 +3948,8 @@ static void emitScanBasedDirective(
       PrivScope.Privatize();
       CGF.CGM.getOpenMPRuntime().emitReduction(
           CGF, S.getEndLoc(), Privates, LHSs, RHSs, ReductionOps,
-          {/*WithNowait=*/true, /*SimpleReduction=*/true, OMPD_unknown});
+          {/*WithNowait=*/true, /*SimpleReduction=*/true,
+           /*IsPrivateVarReduction */ false, OMPD_unknown});
     }
     llvm::Value *NextIVal =
         CGF.Builder.CreateNUWSub(IVal, llvm::ConstantInt::get(CGF.SizeTy, 1));
@@ -5747,7 +5753,7 @@ void CodeGenFunction::EmitOMPScanDirective(const OMPScanDirective &S) {
       }
       CGM.getOpenMPRuntime().emitReduction(
           *this, ParentDir.getEndLoc(), Privates, LHSs, RHSs, ReductionOps,
-          {/*WithNowait=*/true, /*SimpleReduction=*/true, OMPD_simd});
+          {/*WithNowait=*/true, /*SimpleReduction=*/true, false, OMPD_simd});
       for (unsigned I = 0, E = CopyArrayElems.size(); I < E; ++I) {
         const Expr *PrivateExpr = Privates[I];
         LValue DestLVal;