Merge pull request kokkos#7330 from ldh4/crtrott/fix-minmaxloc

Set an initial value index during join of MincLoc, MaxLoc or MinMaxLoc
cz4rs · Sep 23, 2024 · 5f115f4 · 5f115f4
2 parents e9ab977 + 35f9a83
commit 5f115f4
Show file tree

Hide file tree

Showing 3 changed files with 296 additions and 4 deletions.
diff --git a/core/src/Kokkos_Parallel_Reduce.hpp b/core/src/Kokkos_Parallel_Reduce.hpp
@@ -438,7 +438,12 @@ struct MinLoc {
   // Required
   KOKKOS_INLINE_FUNCTION
   void join(value_type& dest, const value_type& src) const {
-    if (src.val < dest.val) dest = src;
+    if (src.val < dest.val)
+      dest = src;
+    else if (src.val == dest.val &&
+             dest.loc == reduction_identity<index_type>::min()) {
+      dest.loc = src.loc;
+    }
   }
 
   KOKKOS_INLINE_FUNCTION
@@ -493,7 +498,12 @@ struct MaxLoc {
   // Required
   KOKKOS_INLINE_FUNCTION
   void join(value_type& dest, const value_type& src) const {
-    if (src.val > dest.val) dest = src;
+    if (src.val > dest.val)
+      dest = src;
+    else if (src.val == dest.val &&
+             dest.loc == reduction_identity<index_type>::min()) {
+      dest.loc = src.loc;
+    }
   }
 
   KOKKOS_INLINE_FUNCTION
@@ -620,10 +630,16 @@ struct MinMaxLoc {
     if (src.min_val < dest.min_val) {
       dest.min_val = src.min_val;
       dest.min_loc = src.min_loc;
+    } else if (dest.min_val == src.min_val &&
+               dest.min_loc == reduction_identity<index_type>::min()) {
+      dest.min_loc = src.min_loc;
     }
     if (src.max_val > dest.max_val) {
       dest.max_val = src.max_val;
       dest.max_loc = src.max_loc;
+    } else if (dest.max_val == src.max_val &&
+               dest.max_loc == reduction_identity<index_type>::min()) {
+      dest.max_loc = src.max_loc;
     }
   }
 

diff --git a/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Reducer.hpp b/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Reducer.hpp
@@ -193,7 +193,12 @@ struct OpenMPTargetReducerWrapper<MinLoc<Scalar, Index, Space>> {
   // Required
   KOKKOS_INLINE_FUNCTION
   static void join(value_type& dest, const value_type& src) {
-    if (src.val < dest.val) dest = src;
+    if (src.val < dest.val)
+      dest = src;
+    else if (src.val == dest.val &&
+             dest.loc == reduction_identity<index_type>::min()) {
+      dest.loc = src.loc;
+    }
   }
 
   KOKKOS_INLINE_FUNCTION
@@ -215,7 +220,12 @@ struct OpenMPTargetReducerWrapper<MaxLoc<Scalar, Index, Space>> {
 
   KOKKOS_INLINE_FUNCTION
   static void join(value_type& dest, const value_type& src) {
-    if (src.val > dest.val) dest = src;
+    if (src.val > dest.val)
+      dest = src;
+    else if (src.val == dest.val &&
+             dest.loc == reduction_identity<index_type>::min()) {
+      dest.loc = src.loc;
+    }
   }
 
   KOKKOS_INLINE_FUNCTION
@@ -268,10 +278,16 @@ struct OpenMPTargetReducerWrapper<MinMaxLoc<Scalar, Index, Space>> {
     if (src.min_val < dest.min_val) {
       dest.min_val = src.min_val;
       dest.min_loc = src.min_loc;
+    } else if (dest.min_val == src.min_val &&
+               dest.min_loc == reduction_identity<index_type>::min()) {
+      dest.min_loc = src.min_loc;
     }
     if (src.max_val > dest.max_val) {
       dest.max_val = src.max_val;
       dest.max_loc = src.max_loc;
+    } else if (dest.max_val == src.max_val &&
+               dest.max_loc == reduction_identity<index_type>::min()) {
+      dest.max_loc = src.max_loc;
     }
   }