Support more reduction operations

src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp

@@ -314,6 +314,10 @@ OutlinedInfoStruct CGIntrinsicsOpenMP::createOutlinedFunction(
     case DSA_REDUCTION_ADD:
     case DSA_REDUCTION_SUB:
     case DSA_REDUCTION_MUL:
+    case DSA_REDUCTION_MAX:
+    case DSA_REDUCTION_UMAX:
+    case DSA_REDUCTION_MIN:
+    case DSA_REDUCTION_UMIN:
       Reductions.push_back(V);
       break;
     default:
@@ -490,6 +494,26 @@ OutlinedInfoStruct CGIntrinsicsOpenMP::createOutlinedFunction(
           OMPBuilder.Builder, AllocaIP, AI, ReductionTy, ReductionInfos,
           IsGPUTeamsReduction);
       break;
+    case DSA_REDUCTION_MAX:
+      Priv = CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_MAX>(
+          OMPBuilder.Builder, AllocaIP, AI, ReductionTy, ReductionInfos,
+          IsGPUTeamsReduction);
+      break;
+    case DSA_REDUCTION_UMAX:
+      Priv = CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_UMAX>(
+          OMPBuilder.Builder, AllocaIP, AI, ReductionTy, ReductionInfos,
+          IsGPUTeamsReduction);
+      break;
+    case DSA_REDUCTION_MIN:
+      Priv = CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_MIN>(
+          OMPBuilder.Builder, AllocaIP, AI, ReductionTy, ReductionInfos,
+          IsGPUTeamsReduction);
+      break;
+    case DSA_REDUCTION_UMIN:
+      Priv = CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_UMIN>(
+          OMPBuilder.Builder, AllocaIP, AI, ReductionTy, ReductionInfos,
+          IsGPUTeamsReduction);
+      break;
     default:
       FATAL_ERROR("Unsupported reduction");
     }
@@ -1288,6 +1312,26 @@ void CGIntrinsicsOpenMP::emitLoop(DSAValueMapTy &DSAValueMap,
             CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_MUL>(
                 OMPBuilder.Builder, OMPBuilder.Builder.saveIP(), Orig, VTy,
                 ReductionInfos, false);
+      } else if (DSA == DSA_REDUCTION_MAX) {
+        ReplacementValue =
+            CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_MAX>(
+                OMPBuilder.Builder, OMPBuilder.Builder.saveIP(), Orig, VTy,
+                ReductionInfos, false);
+      } else if (DSA == DSA_REDUCTION_UMAX) {
+        ReplacementValue =
+            CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_UMAX>(
+                OMPBuilder.Builder, OMPBuilder.Builder.saveIP(), Orig, VTy,
+                ReductionInfos, false);
+      } else if (DSA == DSA_REDUCTION_MIN) {
+        ReplacementValue =
+            CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_MIN>(
+                OMPBuilder.Builder, OMPBuilder.Builder.saveIP(), Orig, VTy,
+                ReductionInfos, false);
+      } else if (DSA == DSA_REDUCTION_UMIN) {
+        ReplacementValue =
+            CGReduction::emitInitAndAppendInfo<DSA_REDUCTION_UMIN>(
+                OMPBuilder.Builder, OMPBuilder.Builder.saveIP(), Orig, VTy,
+                ReductionInfos, false);
       } else
         FATAL_ERROR("Unsupported privatization");
 
@@ -1773,6 +1817,10 @@ void CGIntrinsicsOpenMP::emitOMPOffloadingMappings(
     case DSA_REDUCTION_ADD:
     case DSA_REDUCTION_SUB:
     case DSA_REDUCTION_MUL:
+    case DSA_REDUCTION_MAX:
+    case DSA_REDUCTION_UMAX:
+    case DSA_REDUCTION_MIN:
+    case DSA_REDUCTION_UMIN:
     case DSA_MAP_TOFROM:
       MapType = OMP_TGT_MAPTYPE_TO | OMP_TGT_MAPTYPE_FROM;
       if (IsTargetRegion)
@@ -1822,6 +1870,10 @@ void CGIntrinsicsOpenMP::emitOMPOffloadingMappings(
     case DSA_REDUCTION_ADD:
     case DSA_REDUCTION_SUB:
     case DSA_REDUCTION_MUL:
+    case DSA_REDUCTION_MAX:
+    case DSA_REDUCTION_UMAX:
+    case DSA_REDUCTION_MIN:
+    case DSA_REDUCTION_UMIN:
       Size = ConstantInt::get(OMPBuilder.SizeTy,
                               M.getDataLayout().getTypeAllocSize(V->getType()));
       EmitMappingEntry(Size, GetMapType(DSA), V, V);
@@ -2917,6 +2969,58 @@ Value *CGReduction::emitOperation<DSA_REDUCTION_ADD>(IRBuilderBase &IRB,
     FATAL_ERROR("Unsupported type for reduction operation");
 }
 
+template <>
+Value *CGReduction::emitOperation<DSA_REDUCTION_MAX>(IRBuilderBase &IRB,
+                                                     Value *LHS, Value *RHS) {
+  Type *VTy = RHS->getType();
+  if (VTy->isIntegerTy()) {
+    Value *Cmp = IRB.CreateICmpSGT(LHS, RHS, "red.max.cmp");
+    return IRB.CreateSelect(Cmp, LHS, RHS, "red.max");
+  }
+  if (VTy->isFloatTy() || VTy->isDoubleTy()) {
+    Value *Cmp = IRB.CreateFCmpOGT(LHS, RHS, "red.max.cmp");
+    return IRB.CreateSelect(Cmp, LHS, RHS, "red.max");
+  }
+  FATAL_ERROR("Unsupported type for maximum reduction operation");
+}
+
+template <>
+Value *CGReduction::emitOperation<DSA_REDUCTION_UMAX>(IRBuilderBase &IRB,
+                                                      Value *LHS, Value *RHS) {
+  Type *VTy = RHS->getType();
+  if (VTy->isIntegerTy()) {
+    Value *Cmp = IRB.CreateICmpUGT(LHS, RHS, "red.umax.cmp");
+    return IRB.CreateSelect(Cmp, LHS, RHS, "red.umax");
+  }
+  FATAL_ERROR("Unsupported type for unsigned maximum reduction operation");
+}
+
+template <>
+Value *CGReduction::emitOperation<DSA_REDUCTION_MIN>(IRBuilderBase &IRB,
+                                                     Value *LHS, Value *RHS) {
+  Type *VTy = RHS->getType();
+  if (VTy->isIntegerTy()) {
+    Value *Cmp = IRB.CreateICmpSLT(LHS, RHS, "red.min.cmp");
+    return IRB.CreateSelect(Cmp, LHS, RHS, "red.min");
+  }
+  if (VTy->isFloatTy() || VTy->isDoubleTy()) {
+    Value *Cmp = IRB.CreateFCmpOLT(LHS, RHS, "red.min.cmp");
+    return IRB.CreateSelect(Cmp, LHS, RHS, "red.min");
+  }
+  FATAL_ERROR("Unsupported type for minimum reduction operation");
+}
+
+template <>
+Value *CGReduction::emitOperation<DSA_REDUCTION_UMIN>(IRBuilderBase &IRB,
+                                                      Value *LHS, Value *RHS) {
+  Type *VTy = RHS->getType();
+  if (VTy->isIntegerTy()) {
+    Value *Cmp = IRB.CreateICmpULT(LHS, RHS, "red.umin.cmp");
+    return IRB.CreateSelect(Cmp, LHS, RHS, "red.umin");
+  }
+  FATAL_ERROR("Unsupported type for unsigned minimum reduction operation");
+}
+
 // OpenMP 5.1, 2.21.5, sub is the same as add.
 template <>
 Value *CGReduction::emitOperation<DSA_REDUCTION_SUB>(IRBuilderBase &IRB,
@@ -2950,3 +3054,35 @@ InsertPointTy CGReduction::emitAtomicOperationRMW<DSA_REDUCTION_SUB>(
     IRBuilderBase &IRB, Value *LHS, Value *Partial) {
   return emitAtomicOperationRMW<DSA_REDUCTION_ADD>(IRB, LHS, Partial);
 }
+
+template <>
+InsertPointTy CGReduction::emitAtomicOperationRMW<DSA_REDUCTION_MAX>(
+    IRBuilderBase &IRB, Value *LHS, Value *Partial) {
+  IRB.CreateAtomicRMW(AtomicRMWInst::Max, LHS, Partial, MaybeAlign(),
+                      AtomicOrdering::Monotonic);
+  return IRB.saveIP();
+}
+
+template <>
+InsertPointTy CGReduction::emitAtomicOperationRMW<DSA_REDUCTION_UMAX>(
+    IRBuilderBase &IRB, Value *LHS, Value *Partial) {
+  IRB.CreateAtomicRMW(AtomicRMWInst::UMax, LHS, Partial, MaybeAlign(),
+                      AtomicOrdering::Monotonic);
+  return IRB.saveIP();
+}
+
+template <>
+InsertPointTy CGReduction::emitAtomicOperationRMW<DSA_REDUCTION_MIN>(
+    IRBuilderBase &IRB, Value *LHS, Value *Partial) {
+  IRB.CreateAtomicRMW(AtomicRMWInst::Min, LHS, Partial, MaybeAlign(),
+                      AtomicOrdering::Monotonic);
+  return IRB.saveIP();
+}
+
+template <>
+InsertPointTy CGReduction::emitAtomicOperationRMW<DSA_REDUCTION_UMIN>(
+    IRBuilderBase &IRB, Value *LHS, Value *Partial) {
+  IRB.CreateAtomicRMW(AtomicRMWInst::UMin, LHS, Partial, MaybeAlign(),
+                      AtomicOrdering::Monotonic);
+  return IRB.saveIP();
+}

src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h

@@ -35,6 +35,10 @@ enum DSAType {
   DSA_REDUCTION_ADD,
   DSA_REDUCTION_SUB,
   DSA_REDUCTION_MUL,
+  DSA_REDUCTION_MAX,
+  DSA_REDUCTION_UMAX,
+  DSA_REDUCTION_MIN,
+  DSA_REDUCTION_UMIN,
   DSA_MAP_ALLOC,
   DSA_MAP_TO,
   DSA_MAP_FROM,
@@ -98,6 +102,10 @@ static const DenseMap<StringRef, DSAType> StringToDSA = {
     {"QUAL.OMP.REDUCTION.ADD", DSA_REDUCTION_ADD},
     {"QUAL.OMP.REDUCTION.SUB", DSA_REDUCTION_SUB},
     {"QUAL.OMP.REDUCTION.MUL", DSA_REDUCTION_MUL},
+    {"QUAL.OMP.REDUCTION.MAX", DSA_REDUCTION_MAX},
+    {"QUAL.OMP.REDUCTION.UMAX", DSA_REDUCTION_UMAX},
+    {"QUAL.OMP.REDUCTION.MIN", DSA_REDUCTION_MIN},
+    {"QUAL.OMP.REDUCTION.UMIN", DSA_REDUCTION_UMIN},
     {"QUAL.OMP.MAP.ALLOC", DSA_MAP_ALLOC},
     {"QUAL.OMP.MAP.TO", DSA_MAP_TO},
     {"QUAL.OMP.MAP.FROM", DSA_MAP_FROM},
@@ -129,6 +137,14 @@ inline std::string toString(const DSAType &DSA) {
     return "DSA_REDUCTION_SUB";
   case DSA_REDUCTION_MUL:
     return "DSA_REDUCTION_MUL";
+  case DSA_REDUCTION_MAX:
+    return "DSA_REDUCTION_MAX";
+  case DSA_REDUCTION_UMAX:
+    return "DSA_REDUCTION_UMAX";
+  case DSA_REDUCTION_MIN:
+    return "DSA_REDUCTION_MIN";
+  case DSA_REDUCTION_UMIN:
+    return "DSA_REDUCTION_UMIN";
   case DSA_MAP_ALLOC:
     return "DSA_MAP_ALLOC";
   case DSA_MAP_TO:
@@ -346,6 +362,10 @@ struct CGReduction {
       switch (ReductionOperator) {
       case DSA_REDUCTION_ADD:
       case DSA_REDUCTION_SUB:
+      case DSA_REDUCTION_MAX:
+      case DSA_REDUCTION_UMAX:
+      case DSA_REDUCTION_MIN:
+      case DSA_REDUCTION_UMIN:
         return emitAtomicOperationRMW<ReductionOperator>(Builder, LHS, Partial);
         break;
       case DSA_REDUCTION_MUL:
@@ -373,13 +393,43 @@ struct CGReduction {
       Type *ReductionTy,
       SmallVectorImpl<OpenMPIRBuilder::ReductionInfo> &ReductionInfos,
       bool IsGPUTeamsReduction) {
-    auto GetIdentityValue = []() {
+    auto GetIdentityValue = [ReductionTy]() -> Constant * {
       switch (ReductionOperator) {
       case DSA_REDUCTION_ADD:
       case DSA_REDUCTION_SUB:
-        return 0;
+        return Constant::getNullValue(ReductionTy);
       case DSA_REDUCTION_MUL:
-        return 1;
+        if (ReductionTy->isIntegerTy())
+          return ConstantInt::get(ReductionTy, 1);
+        if (ReductionTy->isFloatingPointTy())
+          return ConstantFP::get(ReductionTy, 1.0);
+        FATAL_ERROR("Invalid value type");
+      case DSA_REDUCTION_MAX:
+        if (auto *IntegerTy = dyn_cast<IntegerType>(ReductionTy)) {
+          APInt Lowest = APInt::getSignedMinValue(IntegerTy->getBitWidth());
+          return ConstantInt::get(IntegerTy, Lowest);
+        }
+        if (ReductionTy->isFloatingPointTy())
+          return ConstantFP::getInfinity(ReductionTy, true);
+        FATAL_ERROR("Invalid value type");
+      case DSA_REDUCTION_UMAX:
+        if (ReductionTy->isIntegerTy())
+          return Constant::getNullValue(ReductionTy);
+        FATAL_ERROR("Invalid value type");
+      case DSA_REDUCTION_MIN:
+        if (auto *IntegerTy = dyn_cast<IntegerType>(ReductionTy)) {
+          APInt Highest = APInt::getSignedMaxValue(IntegerTy->getBitWidth());
+          return ConstantInt::get(IntegerTy, Highest);
+        }
+        if (ReductionTy->isFloatingPointTy())
+          return ConstantFP::getInfinity(ReductionTy);
+        FATAL_ERROR("Invalid value type");
+      case DSA_REDUCTION_UMIN:
+        if (auto *IntegerTy = dyn_cast<IntegerType>(ReductionTy)) {
+          APInt Highest = APInt::getMaxValue(IntegerTy->getBitWidth());
+          return ConstantInt::get(IntegerTy, Highest);
+        }
+        FATAL_ERROR("Invalid value type");
       default:
         FATAL_ERROR("Unknown reduction type");
       }
@@ -404,12 +454,7 @@ struct CGReduction {
     IRB.restoreIP(SaveIP);
 
     // Store identity value based on operation and type.
-    if (ReductionTy->isIntegerTy()) {
-      IRB.CreateStore(ConstantInt::get(ReductionTy, GetIdentityValue()), Priv);
-    } else if (ReductionTy->isFloatTy() || ReductionTy->isDoubleTy()) {
-      IRB.CreateStore(ConstantFP::get(ReductionTy, GetIdentityValue()), Priv);
-    } else
-      FATAL_ERROR("Unsupported type to init with identity reduction value");
+    IRB.CreateStore(GetIdentityValue(), Priv);
 
 #if LLVM_VERSION_MAJOR <= 16
     ReductionInfos.push_back(

src/numba/openmp/omp_grammar.py

@@ -698,7 +698,9 @@
     PLUS: "+"
     MINUS: "-"
     STAR: "*"
-    reduction_operator: PLUS | "\\" | STAR | MINUS | "&" | "^" | "|" | "&&" | "||"
+    MAX: "max"
+    MIN: "min"
+    reduction_operator: PLUS | "\\" | STAR | MINUS | "&" | "^" | "|" | "&&" | "||" | MAX | MIN
     threadprivate_directive: "threadprivate" "(" var_list ")"
     cancellation_point_directive: "cancellation point" construct_type_clause
     construct_type_clause: PARALLEL

src/numba/openmp/omp_lower.py

@@ -2925,6 +2925,10 @@ def reduction_operator(self, args):
             return "SUB"
         elif arg == "*":
             return "MUL"
+        elif arg == "max":
+            return "MAX"
+        elif arg == "min":
+            return "MIN"
         assert 0
 
     def threadprivate_directive(self, args):

src/numba/openmp/tags.py

@@ -263,6 +263,16 @@ def lower(self, lowerer, debug):
 
         name_to_use = self.name
 
+        if name_to_use in [
+            "QUAL.OMP.REDUCTION.MAX",
+            "QUAL.OMP.REDUCTION.MIN",
+        ]:
+            reduction_type = typemap_lookup(typemap, self.arg)
+            if isinstance(reduction_type, types.Integer) and not reduction_type.signed:
+                name_to_use = name_to_use.replace(
+                    "QUAL.OMP.REDUCTION.", "QUAL.OMP.REDUCTION.U"
+                )
+
         is_array = self.arg in typemap and isinstance(
             typemap[self.arg], types.npytypes.Array
         )