[clang] [OpenMP] [Offload] Added support for Xteam min/max reduction. (llvm#1412)

Author: ronlieb

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

@@ -169,15 +169,17 @@ class CGOpenMPRuntimeGPU : public CGOpenMPRuntime {
   llvm::Value *initSpecializedKernel(CodeGenFunction &CGF);
 
   std::pair<llvm::Value *, llvm::Value *>
-  getXteamRedFunctionPtrs(CodeGenFunction &CGF, llvm::Type *RedVarType);
-
-  /// Call cross-team sum
-  llvm::Value *getXteamRedSum(CodeGenFunction &CGF, llvm::Value *Val,
-                              llvm::Value *SumPtr, llvm::Value *DTeamVals,
-                              llvm::Value *DTeamsDonePtr,
-                              llvm::Value *ThreadStartIndex,
-                              llvm::Value *NumTeams, int BlockSize,
-                              bool IsFast);
+  getXteamRedFunctionPtrs(CodeGenFunction &CGF, llvm::Type *RedVarType,
+                          CodeGenModule::XteamRedOpKind Opcode);
+
+  /// Generate a call to cross-team operation.
+  llvm::Value *getXteamRedOperation(CodeGenFunction &CGF, llvm::Value *Val,
+                                    llvm::Value *OrigVarPtr,
+                                    llvm::Value *DTeamVals,
+                                    llvm::Value *DTeamsDonePtr,
+                                    llvm::Value *ThreadStartIndex,
+                                    llvm::Value *NumTeams, int BlockSize,
+                                    CodeGenModule::XteamRedOpKind, bool IsFast);
 
   /// Emit call to Cross-team scan entry points
   llvm::Value *

clang/lib/CodeGen/CGOpenMPRuntimeGPU.h

@@ -682,7 +682,7 @@ void CodeGenFunction::EmitXteamRedCode(const OMPExecutableDirective &D,
     // EmitStmt(CapturedForStmt);
 
     // Now emit the calls to xteam_sum, one for each reduction variable
-    EmitXteamRedSum(CapturedForStmt, *Args, CGM.getXteamRedBlockSize(D));
+    EmitXteamRedOperation(CapturedForStmt, *Args, CGM.getXteamRedBlockSize(D));
   }
 
   // Xteam codegen done
@@ -706,14 +706,9 @@ void CodeGenFunction::EmitXteamLocalAggregator(const ForStmt *FStmt) {
             RedVarType->isIntegerTy()) &&
            "Unhandled type");
     llvm::AllocaInst *XteamRedInst = Builder.CreateAlloca(RedVarType);
-    llvm::Value *InitVal = nullptr;
-    if (RedVarType->isFloatTy() || RedVarType->isDoubleTy() ||
-        RedVarType->isHalfTy() || RedVarType->isBFloatTy())
-      InitVal = llvm::ConstantFP::getZero(RedVarType);
-    else if (RedVarType->isIntegerTy())
-      InitVal = llvm::ConstantInt::get(RedVarType, 0);
-    else
-      llvm_unreachable("Unhandled type");
+    // The initial value (referred to as the sentinel value) of the local
+    // reduction variable depends on the opcode.
+    llvm::Value *InitVal = getXteamRedSentinel(RedVarType, Itr->second.Opcode);
     Address XteamRedVarAddr(
         XteamRedInst, RedVarType,
         getContext().getTypeAlignInChars(RedVarExpr->getType()));
@@ -726,11 +721,56 @@ void CodeGenFunction::EmitXteamLocalAggregator(const ForStmt *FStmt) {
   }
 }
 
-// Emit __kmpc_xteam_sum(*xteam_red_local_addr, red_var_addr) for each reduction
-// in the helper map for the given For Stmt
-void CodeGenFunction::EmitXteamRedSum(const ForStmt *FStmt,
-                                      const FunctionArgList &Args,
-                                      int BlockSize) {
+llvm::Value *
+CodeGenFunction::getXteamRedSentinel(llvm::Type *RedVarType,
+                                     CodeGenModule::XteamRedOpKind Opcode) {
+  assert((RedVarType->isFloatTy() || RedVarType->isDoubleTy() ||
+          RedVarType->isHalfTy() || RedVarType->isBFloatTy() ||
+          RedVarType->isIntegerTy()) &&
+         "Unhandled type");
+  assert(Opcode != CodeGenModule::XR_OP_unknown &&
+         "Unexpected Xteam reduction opcode");
+  if (RedVarType->isFloatTy() || RedVarType->isDoubleTy() ||
+      RedVarType->isHalfTy() || RedVarType->isBFloatTy()) {
+    if (Opcode == CodeGenModule::XR_OP_add)
+      return llvm::ConstantFP::getZero(RedVarType);
+    else if (Opcode == CodeGenModule::XR_OP_min)
+      return llvm::ConstantFP::getInfinity(RedVarType);
+    else // max operator
+      return llvm::ConstantFP::getInfinity(RedVarType, /*Negative=*/true);
+  } else {
+    // Integer type
+    if (RedVarType->getPrimitiveSizeInBits() == 16)
+      return llvm::ConstantInt::get(Int16Ty,
+                                    Opcode == CodeGenModule::XR_OP_add ? 0
+                                    : Opcode == CodeGenModule::XR_OP_min
+                                        ? std::numeric_limits<int16_t>::max()
+                                        : std::numeric_limits<int16_t>::min());
+    else if (RedVarType->getPrimitiveSizeInBits() == 32)
+      return llvm::ConstantInt::get(Int32Ty,
+                                    Opcode == CodeGenModule::XR_OP_add ? 0
+                                    : Opcode == CodeGenModule::XR_OP_min
+                                        ? std::numeric_limits<int32_t>::max()
+                                        : std::numeric_limits<int32_t>::min());
+    else {
+      assert(RedVarType->getPrimitiveSizeInBits() == 64 &&
+             "Expected a 64-bit integer");
+      return llvm::ConstantInt::get(Int64Ty,
+                                    Opcode == CodeGenModule::XR_OP_add ? 0
+                                    : Opcode == CodeGenModule::XR_OP_min
+                                        ? std::numeric_limits<int64_t>::max()
+                                        : std::numeric_limits<int64_t>::min());
+    }
+  }
+  llvm_unreachable(
+      "Unexpected type or opcode in Xteam reduction sentinel generation");
+}
+
+// Emit a call to the DeviceRTL Xteam reduction function for each reduction
+// variable in the helper map for the given For Stmt.
+void CodeGenFunction::EmitXteamRedOperation(const ForStmt *FStmt,
+                                            const FunctionArgList &Args,
+                                            int BlockSize) {
   auto &RT = static_cast<CGOpenMPRuntimeGPU &>(CGM.getOpenMPRuntime());
   const CodeGenModule::XteamRedVarMap &RedVarMap = CGM.getXteamRedVarMap(FStmt);
 
@@ -760,10 +800,12 @@ void CodeGenFunction::EmitXteamRedSum(const ForStmt *FStmt,
     const Expr *OrigRedVarExpr = RVI.RedVarExpr;
     const DeclRefExpr *DRE = cast<DeclRefExpr>(OrigRedVarExpr);
     Address OrigRedVarAddr = EmitLValue(DRE).getAddress();
-    // Pass in OrigRedVarAddr.getPointer to kmpc_xteam_sum
-    RT.getXteamRedSum(*this, Builder.CreateLoad(RVI.RedVarAddr),
-                      OrigRedVarAddr.emitRawPointer(*this), DTeamVals, DTeamsDonePtr,
-                      ThreadStartIdx, NumTeams, BlockSize, IsFast);
+    // Note that fast Xteam reduction is available only for sum operator.
+    RT.getXteamRedOperation(*this, Builder.CreateLoad(RVI.RedVarAddr),
+                            OrigRedVarAddr.emitRawPointer(*this), DTeamVals,
+                            DTeamsDonePtr, ThreadStartIdx, NumTeams, BlockSize,
+                            RVI.Opcode,
+                            IsFast && RVI.Opcode == CodeGenModule::XR_OP_add);
   }
 }
 
@@ -857,8 +899,6 @@ void CodeGenFunction::EmitXteamScanPhaseTwo(const ForStmt *FStmt,
   }
 }
 
-// Emit reduction into local aggregator for a statement within the reduced loop
-// where applicable
 bool CodeGenFunction::EmitXteamRedStmt(const Stmt *S) {
   if (CGM.getCurrentXteamRedStmt() == nullptr)
     return false;
@@ -882,6 +922,10 @@ bool CodeGenFunction::EmitXteamRedStmt(const Stmt *S) {
   const CodeGenModule::XteamRedVarMap &RedVarMap =
       CGM.getXteamRedVarMap(CGM.getCurrentXteamRedStmt());
 
+  // Currently, there is limited support in Xteam reduction for calls with
+  // reduction variables in arguments. Either the call has to be at the
+  // statement level or it has to be a call to a builtin function (e.g. min/max)
+  // on the rhs of an assignment statement. Handle call at the statement level.
   if (isa<CallExpr>(S)) {
     const CallExpr *CE = cast<CallExpr>(S);
     assert(CE && "Unexpected null call expression");
@@ -962,45 +1006,131 @@ bool CodeGenFunction::EmitXteamRedStmt(const Stmt *S) {
     RedRHSExpr = RedBO->getRHS()->IgnoreImpCasts();
   } else {
     const Expr *L1RhsExpr = RedBO->getRHS()->IgnoreImpCasts();
-    assert(isa<BinaryOperator>(L1RhsExpr) &&
+    assert((isa<BinaryOperator>(L1RhsExpr) || isa<CallExpr>(L1RhsExpr) ||
+            isa<PseudoObjectExpr>(L1RhsExpr)) &&
            "Expected rhs to be a binary operator");
-    const BinaryOperator *L2BO = cast<BinaryOperator>(L1RhsExpr);
-    auto OpcL2BO = L2BO->getOpcode();
-    assert(OpcL2BO == BO_Add && "Unexpected operator");
-    // If the redvar is lhs, use the rhs in the generated reduction statement
-    // and vice-versa.
-    if (CGM.isXteamRedVarExpr(L2BO->getLHS()->IgnoreImpCasts(), RedVarDecl))
-      RedRHSExpr = L2BO->getRHS();
-    else if (CGM.isXteamRedVarExpr(L2BO->getRHS()->IgnoreImpCasts(),
-                                   RedVarDecl))
-      RedRHSExpr = L2BO->getLHS();
-    else
-      llvm_unreachable("Unhandled add expression during xteam reduction");
+    if (isa<BinaryOperator>(L1RhsExpr)) {
+      const BinaryOperator *L2BO = cast<BinaryOperator>(L1RhsExpr);
+      auto OpcL2BO = L2BO->getOpcode();
+      assert(OpcL2BO == BO_Add && "Unexpected operator");
+      // If the redvar is lhs, use the rhs in the generated reduction statement
+      // and vice-versa.
+      if (CGM.isXteamRedVarExpr(L2BO->getLHS()->IgnoreImpCasts(), RedVarDecl))
+        RedRHSExpr = L2BO->getRHS();
+      else if (CGM.isXteamRedVarExpr(L2BO->getRHS()->IgnoreImpCasts(),
+                                     RedVarDecl))
+        RedRHSExpr = L2BO->getLHS();
+      else
+        llvm_unreachable("Unhandled add expression during xteam reduction");
+    } else if (isa<CallExpr>(L1RhsExpr)) {
+      const CallExpr *Call = cast<CallExpr>(L1RhsExpr);
+      assert(CGM.getStatusOptKernelBuiltin(Call) == CodeGenModule::NxSuccess &&
+             "Expected a call to an Xteam supported builtin");
+      EmitXteamRedStmtForBuiltinCall(Call, RedVarDecl, RedVarMap);
+      return true;
+    } else {
+      assert(isa<PseudoObjectExpr>(L1RhsExpr) && "Expected a PseudoObjectExpr");
+      auto [Status, ReturnExpr] = CGM.getStatusXteamSupportedPseudoObject(
+          cast<PseudoObjectExpr>(L1RhsExpr));
+      assert(Status == CodeGenModule::NxSuccess &&
+             "Expected call expression from analysis of PseudoObjectExpr");
+      const CallExpr *Call = cast<CallExpr>(ReturnExpr);
+      assert(CGM.getStatusOptKernelBuiltin(Call) == CodeGenModule::NxSuccess &&
+             "Expected a call to an Xteam supported builtin");
+      EmitXteamRedStmtForBuiltinCall(Call, RedVarDecl, RedVarMap);
+      return true;
+    }
   }
   assert(RedRHSExpr != nullptr && "Did not find a valid reduction rhs");
-  llvm::Value *RHSValue = EmitScalarExpr(RedRHSExpr);
+
+  EmitLocalReductionStmt(RedRHSExpr, RedVarDecl, RedVarMap,
+                         CodeGenModule::XR_OP_add);
+  return true;
+}
+
+void CodeGenFunction::EmitLocalReductionStmt(
+    const Expr *E, const VarDecl *RedVarDecl,
+    const CodeGenModule::XteamRedVarMap &RedVarMap,
+    CodeGenModule::XteamRedOpKind OpKind) {
+  // For add, generate *xteam_local = *xteam_local + rhs_value
+  // For min/max, generate *xteam_local = min/max(*xteam_local, other_operand)
+
+  // First, generate the other operand.
+  llvm::Value *RHSValue = EmitScalarExpr(E);
+  // Now handle the local reduction variable accesses.
   auto It = RedVarMap.find(RedVarDecl);
   assert(It != RedVarMap.end() && "Variable must be found in reduction map");
   Address XteamRedLocalAddr = It->second.RedVarAddr;
-  // Compute *xteam_red_local_addr + rhs_value
-  llvm::Value *RedRHS = nullptr;
   llvm::Type *RedVarType = ConvertTypeForMem(It->second.RedVarExpr->getType());
+  llvm::Value *Op1 = Builder.CreateLoad(XteamRedLocalAddr);
+  llvm::Value *RedRHS = nullptr;
   if (RedVarType->isFloatTy() || RedVarType->isDoubleTy() ||
       RedVarType->isHalfTy() || RedVarType->isBFloatTy()) {
-    auto RHSOp = RHSValue->getType()->isIntegerTy()
-                     ? Builder.CreateSIToFP(RHSValue, RedVarType)
-                     : Builder.CreateFPCast(RHSValue, RedVarType);
-    RedRHS = Builder.CreateFAdd(Builder.CreateLoad(XteamRedLocalAddr), RHSOp);
+    auto Op2 = RHSValue->getType()->isIntegerTy()
+                   ? Builder.CreateSIToFP(RHSValue, RedVarType)
+                   : Builder.CreateFPCast(RHSValue, RedVarType);
+    if (OpKind == CodeGenModule::XR_OP_add)
+      RedRHS = Builder.CreateFAdd(Op1, Op2);
+    else if (OpKind == CodeGenModule::XR_OP_min)
+      RedRHS =
+          Builder.CreateMinNum(Op1, Op2, /*FMFSource=*/nullptr, "xteam.min");
+    else if (OpKind == CodeGenModule::XR_OP_max)
+      RedRHS =
+          Builder.CreateMaxNum(Op1, Op2, /*FMFSource=*/nullptr, "xteam.max");
+    else
+      llvm_unreachable("Unexpected reduction kind");
   } else if (RedVarType->isIntegerTy()) {
-    auto RHSOp = RHSValue->getType()->isIntegerTy()
-                     ? Builder.CreateIntCast(RHSValue, RedVarType, false)
-                     : Builder.CreateFPToSI(RHSValue, RedVarType);
-    RedRHS = Builder.CreateAdd(Builder.CreateLoad(XteamRedLocalAddr), RHSOp);
+    auto Op2 = RHSValue->getType()->isIntegerTy()
+                   ? Builder.CreateIntCast(RHSValue, RedVarType, false)
+                   : Builder.CreateFPToSI(RHSValue, RedVarType);
+    if (OpKind == CodeGenModule::XR_OP_add)
+      RedRHS = Builder.CreateAdd(Op1, Op2);
+    else if (OpKind == CodeGenModule::XR_OP_min)
+      // TODO Fix when unsigned
+      RedRHS = Builder.CreateBinaryIntrinsic(llvm::Intrinsic::smin, Op1, Op2,
+                                             nullptr, "xteam.min");
+    else if (OpKind == CodeGenModule::XR_OP_max)
+      // TODO fix when unsigned
+      RedRHS = Builder.CreateBinaryIntrinsic(llvm::Intrinsic::smax, Op1, Op2,
+                                             nullptr, "xteam.max");
+    else
+      llvm_unreachable("Unexpected reduction kind");
   } else
     llvm_unreachable("Unhandled type");
-  // *xteam_red_local_addr = *xteam_red_local_addr + rhs_value
+  assert(RedRHS && "Right hand side of statement cannot be null");
   Builder.CreateStore(RedRHS, XteamRedLocalAddr);
-  return true;
+}
+
+std::pair<const Expr *, CodeGenModule::XteamRedOpKind>
+CodeGenFunction::ExtractXteamRedRhsExpr(const CallExpr *Call,
+                                        const VarDecl *RedVarDecl) {
+  // Traverse arguments, identifying and ignoring the reduction variable, and
+  // then extracting the other argument.
+  CodeGenModule::XteamRedOpKind Opcode;
+  std::string CallName = Call->getDirectCallee()->getNameInfo().getAsString();
+  if (CGM.isOptKernelAMDGCNMax(CallName))
+    Opcode = CodeGenModule::XR_OP_max;
+  else if (CGM.isOptKernelAMDGCNMin(CallName))
+    Opcode = CodeGenModule::XR_OP_min;
+  else
+    llvm_unreachable("Epecting either min or max");
+
+  for (unsigned ArgIndex = 0; ArgIndex < Call->getNumArgs(); ++ArgIndex) {
+    const Expr *Arg = Call->getArg(ArgIndex);
+    while (isa<ImplicitCastExpr>(Arg))
+      Arg = cast<ImplicitCastExpr>(Arg)->getSubExpr();
+    if (CGM.isXteamRedVarExpr(Arg, RedVarDecl))
+      continue;
+    return std::make_pair(Call->getArg(ArgIndex), Opcode);
+  }
+  llvm_unreachable("Could not extract expected arg of min/max");
+}
+
+void CodeGenFunction::EmitXteamRedStmtForBuiltinCall(
+    const CallExpr *Call, const VarDecl *RedVarDecl,
+    const CodeGenModule::XteamRedVarMap &RedVarMap) {
+  auto [RhsExpr, Opcode] = ExtractXteamRedRhsExpr(Call, RedVarDecl);
+  EmitLocalReductionStmt(RhsExpr, RedVarDecl, RedVarMap, Opcode);
 }
 
 void CodeGenFunction::EmitStmt(const Stmt *S, ArrayRef<const Attr *> Attrs) {

clang/lib/CodeGen/CGStmt.cpp

@@ -956,14 +956,19 @@ llvm::Function *CodeGenFunction::GenerateOpenMPCapturedStmtFunction(
     // If Xteam found, use it. Otherwise, query again. This is required to make
     // sure that the outlined routines have the correct signature.
     if (FStmt) {
-      if (!CGM.isXteamRedKernel(FStmt))
-        isXteamKernel =
-            CGM.checkAndSetXteamRedKernel(D) == CodeGenModule::NxSuccess;
-      else
+      if (!CGM.isXteamRedKernel(FStmt)) {
+        CodeGenModule::NoLoopXteamErr NxStatus =
+            CGM.checkAndSetXteamRedKernel(D);
+        DEBUG_WITH_TYPE(NO_LOOP_XTEAM_RED,
+                        CGM.emitNxResult("[Xteam-host]", D, NxStatus));
+        isXteamKernel = (NxStatus == CodeGenModule::NxSuccess);
+      } else
         isXteamKernel = true;
     } else {
-      isXteamKernel =
-          CGM.checkAndSetXteamRedKernel(D) == CodeGenModule::NxSuccess;
+      CodeGenModule::NoLoopXteamErr NxStatus = CGM.checkAndSetXteamRedKernel(D);
+      DEBUG_WITH_TYPE(NO_LOOP_XTEAM_RED,
+                      CGM.emitNxResult("[Xteam-host]", D, NxStatus));
+      isXteamKernel = (NxStatus == CodeGenModule::NxSuccess);
     }
   }
 
@@ -7677,8 +7682,11 @@ static void emitCommonOMPTargetDirective(CodeGenFunction &CGF,
   const Stmt *OptKernelKey = CGM.getOptKernelKey(S);
   if (OptKernelKey)
     FStmt = CGM.getSingleForStmt(OptKernelKey);
-  if (FStmt && CGM.getLangOpts().OpenMPOffloadMandatory)
-    CGM.checkAndSetXteamRedKernel(S);
+  if (FStmt && CGM.getLangOpts().OpenMPOffloadMandatory) {
+    CodeGenModule::NoLoopXteamErr NxStatus = CGM.checkAndSetXteamRedKernel(S);
+    DEBUG_WITH_TYPE(NO_LOOP_XTEAM_RED,
+                    CGM.emitNxResult("[Xteam-host]", S, NxStatus));
+  }
 
   if (CGM.getLangOpts().OpenMPOffloadMandatory && !IsOffloadEntry) {
     unsigned DiagID = CGM.getDiags().getCustomDiagID(