[flang] Fix MAXLOC/MINLOC when MASK is scalar .FALSE.

[PeteSteinfeld]

Summary

When passing a scalar .FALSE. as the MASK argument to MAXLOC, we were getting
bad memory references. We were falling into the code intended when the MASK
argument was missing.

I fixed this by passing extra information to determine if the MASK argument is
missing versus a scalar value. I also added tests for MAXLOC and MINLOC with
scalar values of .TRUE. and .FALSE. for the MASK argument.

Along the way, I eliminated the unused "chars" argument from the constructor
for ExtremumLocAccumulator.

[jeanPerier]

flang/runtime/extrema.cpp

84

It is really not straightforward to me how you can assume that mask_ is a scalar false here when the if is only testing for the presence of a MASK.

I assume this may come from the way GetResult() is called in flang/runtime/reduction-templates.h, but that is a big assumption.

Have you tried fixing the issue on the reduction-templates side that is calling GetResult ?

[PeteSteinfeld]

Responding to Jean's comments. Rather than using the presence of a mask in
the ExtremumLocAccumulator object to determine if the MASK argument is a
scalar .FALSE. value, I added an argument to the GetResult function.

[PeteSteinfeld]

flang/runtime/extrema.cpp

84

Thanks, Jean. Your point is well taken.

I changed things to add an argument to GetResult to explicitly cover the situation where the MASK argument is a scalar .FALSE. value. Let me know if there's a better way to handle this.

[klausler]

When the mask is scalar .false. the result will be zero(es). Seems easier and clearer to check for that case up front than to run the reduction and then override the result.

[PeteSteinfeld]

Responding to Peter's comments. I restricted all changes to extrema.cpp where
I modified TypedPartialMaxOrMinLoc to test for the scalar .FALSE. value of MASK
and, in the case, create the result and set it to all zeroes.

[PeteSteinfeld]

Thanks @klausler. I'm tried to follow your suggestion. Please take another look.

[klausler]

flang/runtime/extrema.cpp

262

In the case where the mask is scalar and .TRUE., would it be best here to set mask = nullptr;?

[PeteSteinfeld]

Responding to Peter's latest suggestion.

flang/runtime/extrema.cpp

 
 private:
   std::size_t chars_;
 };
 
 template <typename COMPARE> class ExtremumLocAccumulator {
 public:
   using Type = typename COMPARE::Type;
-  ExtremumLocAccumulator(const Descriptor &array, std::size_t chars = 0)
-      : array_{array}, argRank_{array.rank()}, compare_{array.ElementBytes()} {
+  ExtremumLocAccumulator(
+      const Descriptor &array, const Descriptor *mask = nullptr)
+      : array_{array}, argRank_{array.rank()}, compare_{array.ElementBytes()},
+        mask_(mask) {
     Reinitialize();
   }
   void Reinitialize() {
     // per standard: result indices are all zero if no data
     for (int j{0}; j < argRank_; ++j) {
       extremumLoc_[j] = 0;
     }
     previous_ = nullptr;
   }
   int argRank() const { return argRank_; }
   template <typename A> void GetResult(A *p, int zeroBasedDim = -1) {
     if (zeroBasedDim >= 0) {
       *p = extremumLoc_[zeroBasedDim] -
           array_.GetDimension(zeroBasedDim).LowerBound() + 1;
     } else {
+      // Either mask is missing or it's scalar .FALSE.
+      if (mask_) {
+        // Scalar .FALSE. -- all values are zero
+        *p = 0;
+      } else {
+        // No mask
         for (int j{0}; j < argRank_; ++j) {
           p[j] = extremumLoc_[j] - array_.GetDimension(j).LowerBound() + 1;
         }
       }
     }
+  }
   template <typename IGNORED> bool AccumulateAt(const SubscriptValue at[]) {
     const auto &value{*array_.Element<Type>(at)};
     if (!previous_ || compare_(value, *previous_)) {
       previous_ = &value;
       for (int j{0}; j < argRank_; ++j) {
         extremumLoc_[j] = at[j];
       }
     }
     return true;
   }
 
 private:
   const Descriptor &array_;
   int argRank_;
   SubscriptValue extremumLoc_[maxRank];
   const Type *previous_{nullptr};
   COMPARE compare_;
+  const Descriptor *mask_;
 };
 
 template <typename ACCUMULATOR, typename CPPTYPE>
 static void LocationHelper(const char *intrinsic, Descriptor &result,
     const Descriptor &x, int kind, const Descriptor *mask,
     Terminator &terminator) {
   ACCUMULATOR accumulator{x};
   DoTotalReduction<CPPTYPE>(x, 0, mask, accumulator, intrinsic, terminator);

 
 template <TypeCategory CAT, int KIND, bool IS_MAX,
     template <typename, bool, bool> class COMPARE, bool BACK>
 static void DoPartialMaxOrMinLocDirection(const char *intrinsic,
     Descriptor &result, const Descriptor &x, int kind, int dim,
     const Descriptor *mask, Terminator &terminator) {
   using CppType = CppTypeFor<CAT, KIND>;
   using Accumulator = ExtremumLocAccumulator<COMPARE<CppType, IS_MAX, BACK>>;
-  Accumulator accumulator{x};
+  Accumulator accumulator{x, mask};
   ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor, void>(
       kind, terminator, result, x, dim, mask, terminator, intrinsic,
       accumulator);
 }
 
 template <TypeCategory CAT, int KIND, bool IS_MAX,
     template <typename, bool, bool> class COMPARE>
 inline void DoPartialMaxOrMinLoc(const char *intrinsic, Descriptor &result,
inline void TypedPartialMaxOrMinLoc(const char *intrinsic, Descriptor &result,
   switch (catKind->first) {
   case TypeCategory::Integer:
     ApplyIntegerKind<DoPartialMaxOrMinLocHelper<TypeCategory::Integer, IS_MAX,
                          NumericCompare>::template Functor,
         void>(catKind->second, terminator, intrinsic, result, x, kind, dim,
         mask, back, terminator);
     break;
   case TypeCategory::Real:
     ApplyFloatingPointKind<DoPartialMaxOrMinLocHelper<TypeCategory::Real,
                                IS_MAX, NumericCompare>::template Functor,
         void>(catKind->second, terminator, intrinsic, result, x, kind, dim,
         mask, back, terminator);
     break;
   case TypeCategory::Character:
     ApplyCharacterKind<DoPartialMaxOrMinLocHelper<TypeCategory::Character,
                            IS_MAX, CharacterCompare>::template Functor,
         void>(catKind->second, terminator, intrinsic, result, x, kind, dim,

flang/unittests/Runtime/Reduction.cpp

TEST(Reductions, DoubleMaxMinNorm2) {
   StaticDescriptor<1, true> statDesc0[1];
   Descriptor &scalarResult{statDesc0[0].descriptor()};
   RTNAME(MaxlocDim)
   (scalarResult, *array1, /*KIND=*/2, /*DIM=*/1, __FILE__, __LINE__,
       /*MASK=*/nullptr, /*BACK=*/false);
   EXPECT_EQ(scalarResult.rank(), 0);
   EXPECT_EQ(*scalarResult.ZeroBasedIndexedElement<std::int16_t>(0), 23);
   scalarResult.Destroy();
+
+  // Test .FALSE. scalar MASK argument
+  auto falseMask{MakeArray<TypeCategory::Logical, 4>(
+      std::vector<int>{}, std::vector<std::int32_t>{0})};
+  RTNAME(MaxlocDim)
+  (loc, *array, /*KIND=*/4, /*DIM=*/2, __FILE__, __LINE__,
+      /*MASK=*/&*falseMask, /*BACK=*/false);
+  EXPECT_EQ(loc.rank(), 2);
+  EXPECT_EQ(loc.type().raw(), (TypeCode{TypeCategory::Integer, 4}.raw()));
+  EXPECT_EQ(loc.GetDimension(0).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(0).Extent(), 3);
+  EXPECT_EQ(loc.GetDimension(1).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(1).Extent(), 2);
+  for (int i{0}; i < 6; ++i) {
+    EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(0), 0);
+  }
+  loc.Destroy();
+
+  // Test .TRUE. scalar MASK argument
+  auto trueMask{MakeArray<TypeCategory::Logical, 4>(
+      std::vector<int>{}, std::vector<std::int32_t>{1})};
+  RTNAME(MaxlocDim)
+  (loc, *array, /*KIND=*/4, /*DIM=*/2, __FILE__, __LINE__,
+      /*MASK=*/&*trueMask, /*BACK=*/false);
+  EXPECT_EQ(loc.rank(), 2);
+  EXPECT_EQ(loc.type().raw(), (TypeCode{TypeCategory::Integer, 4}.raw()));
+  EXPECT_EQ(loc.GetDimension(0).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(0).Extent(), 3);
+  EXPECT_EQ(loc.GetDimension(1).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(1).Extent(), 2);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(0), 3);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(1), 4);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(2), 3);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(3), 3);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(4), 4);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(5), 4);
+  loc.Destroy();
+
   RTNAME(MinlocDim)
   (scalarResult, *array1, /*KIND=*/2, /*DIM=*/1, __FILE__, __LINE__,
       /*MASK=*/nullptr, /*BACK=*/true);
   EXPECT_EQ(scalarResult.rank(), 0);
   EXPECT_EQ(*scalarResult.ZeroBasedIndexedElement<std::int16_t>(0), 22);
   scalarResult.Destroy();
+
+  // Test .FALSE. scalar MASK argument
+  RTNAME(MinlocDim)
+  (loc, *array, /*KIND=*/4, /*DIM=*/2, __FILE__, __LINE__,
+      /*MASK=*/&*falseMask, /*BACK=*/false);
+  EXPECT_EQ(loc.rank(), 2);
+  EXPECT_EQ(loc.type().raw(), (TypeCode{TypeCategory::Integer, 4}.raw()));
+  EXPECT_EQ(loc.GetDimension(0).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(0).Extent(), 3);
+  EXPECT_EQ(loc.GetDimension(1).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(1).Extent(), 2);
+  for (int i{0}; i < 6; ++i) {
+    EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(0), 0);
+  }
+  loc.Destroy();
+
+  // Test .TRUE. scalar MASK argument
+  RTNAME(MinlocDim)
+  (loc, *array, /*KIND=*/4, /*DIM=*/2, __FILE__, __LINE__,
+      /*MASK=*/&*trueMask, /*BACK=*/false);
+  EXPECT_EQ(loc.rank(), 2);
+  EXPECT_EQ(loc.type().raw(), (TypeCode{TypeCategory::Integer, 4}.raw()));
+  EXPECT_EQ(loc.GetDimension(0).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(0).Extent(), 3);
+  EXPECT_EQ(loc.GetDimension(1).LowerBound(), 1);
+  EXPECT_EQ(loc.GetDimension(1).Extent(), 2);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(0), 4);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(1), 3);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(2), 4);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(3), 4);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(4), 3);
+  EXPECT_EQ(*loc.ZeroBasedIndexedElement<std::int32_t>(5), 2);
+  loc.Destroy();
+
   RTNAME(MaxvalDim)
   (scalarResult, *array1, /*DIM=*/1, __FILE__, __LINE__, /*MASK=*/nullptr);
   EXPECT_EQ(scalarResult.rank(), 0);
   EXPECT_EQ(*scalarResult.ZeroBasedIndexedElement<double>(0), 22.0);
   scalarResult.Destroy();
   RTNAME(MinvalDim)
   (scalarResult, *array1, /*DIM=*/1, __FILE__, __LINE__, /*MASK=*/nullptr);
   EXPECT_EQ(scalarResult.rank(), 0);