[flang] Fold MAXLOC and MINLOC

[klausler]

Summary

Generalize the code that folds FINDLOC to also handle
folding for MAXLOC and MINLOC.

[clementval]

LGTM

flang/include/flang/Evaluate/type.h

 using IntegerTypes = CategoryTypes<TypeCategory::Integer>;
 using RealTypes = CategoryTypes<TypeCategory::Real>;
 using ComplexTypes = CategoryTypes<TypeCategory::Complex>;
 using CharacterTypes = CategoryTypes<TypeCategory::Character>;
 using LogicalTypes = CategoryTypes<TypeCategory::Logical>;
 
 using FloatingTypes = common::CombineTuples<RealTypes, ComplexTypes>;
 using NumericTypes = common::CombineTuples<IntegerTypes, FloatingTypes>;
-using RelationalTypes = common::CombineTuples<NumericTypes, CharacterTypes>;
-using AllIntrinsicTypes = common::CombineTuples<RelationalTypes, LogicalTypes>;
+using RelationalTypes =
+    common::CombineTuples<IntegerTypes, RealTypes, CharacterTypes>;
+using AllIntrinsicTypes =
+    common::CombineTuples<NumericTypes, CharacterTypes, LogicalTypes>;
 using LengthlessIntrinsicTypes =
     common::CombineTuples<NumericTypes, LogicalTypes>;
 
 // Predicates: does a type represent a specific intrinsic type?
 template <typename T>
 constexpr bool IsSpecificIntrinsicType{common::HasMember<T, AllIntrinsicTypes>};
 
 // Predicate: is a type an intrinsic type that is completely characterized

flang/lib/Evaluate/fold-character.cpp

if (name == "achar" || name == "char") {
     }
   } else if (name == "trim") { // not elemental
     if (auto scalar{
             GetScalarConstantArguments<T>(context, funcRef.arguments())}) {
       return Expr<T>{Constant<T>{
           CharacterUtils<KIND>::TRIM(std::get<Scalar<T>>(*scalar))}};
     }
   }
-  // TODO: maxloc, minloc, transfer
+  // TODO: transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Character, KIND>> FoldOperation(
     FoldingContext &context, Concat<KIND> &&x) {
   if (auto array{ApplyElementwise(context, x)}) {
     return *array;

flang/lib/Evaluate/fold-integer.cpp

if (CheckReductionDIM(dim, context, arg, 1, mask->Rank())) {
         }
       }};
       return Expr<T>{DoReduction<T>(*mask, dim, Scalar<T>{}, accumulator)};
     }
   }
   return Expr<T>{std::move(ref)};
 }
 
-// FINDLOC()
-class FindlocHelper {
+// FINDLOC(), MAXLOC(), & MINLOC()
+enum class WhichLocation { Findloc, Maxloc, Minloc };
+template <WhichLocation WHICH> class LocationHelper {
 public:
-  FindlocHelper(
+  LocationHelper(
       DynamicType &&type, ActualArguments &arg, FoldingContext &context)
       : type_{type}, arg_{arg}, context_{context} {}
   using Result = std::optional<Constant<SubscriptInteger>>;
-  using Types = AllIntrinsicTypes;
+  using Types = std::conditional_t<WHICH == WhichLocation::Findloc,
+      AllIntrinsicTypes, RelationalTypes>;
 
   template <typename T> Result Test() const {
     if (T::category != type_.category() || T::kind != type_.kind()) {
       return std::nullopt;
     }
-    CHECK(arg_.size() == 6);
+    CHECK(arg_.size() == (WHICH == WhichLocation::Findloc ? 6 : 5));
     Folder<T> folder{context_};
     Constant<T> *array{folder.Folding(arg_[0])};
-    Constant<T> *value{folder.Folding(arg_[1])};
-    if (!array || !value) {
+    if (!array) {
       return std::nullopt;
     }
+    std::optional<Constant<T>> value;
+    if constexpr (WHICH == WhichLocation::Findloc) {
+      if (const Constant<T> *p{folder.Folding(arg_[1])}) {
+        value.emplace(*p);
+      } else {
+        return std::nullopt;
+      }
+    }
     std::optional<int> dim;
     Constant<LogicalResult> *mask{
-        GetReductionMASK(arg_[3], array->shape(), context_)};
-    if ((!mask && arg_[3]) ||
-        !CheckReductionDIM(dim, context_, arg_, 2, array->Rank())) {
+        GetReductionMASK(arg_[maskArg], array->shape(), context_)};
+    if ((!mask && arg_[maskArg]) ||
+        !CheckReductionDIM(dim, context_, arg_, dimArg, array->Rank())) {
       return std::nullopt;
     }
     bool back{false};
-    if (arg_[5]) {
-      const auto *backConst{Folder<LogicalResult>{context_}.Folding(arg_[5])};
+    if (arg_[backArg]) {
+      const auto *backConst{
+          Folder<LogicalResult>{context_}.Folding(arg_[backArg])};
       if (backConst) {
         back = backConst->GetScalarValue().value().IsTrue();
       } else {
         return std::nullopt;
       }
     }
+    const RelationalOperator relation{WHICH == WhichLocation::Findloc
+            ? RelationalOperator::EQ
+            : WHICH == WhichLocation::Maxloc
+            ? (back ? RelationalOperator::GE : RelationalOperator::GT)
+            : back ? RelationalOperator::LE
+                   : RelationalOperator::LT};
     // Use lower bounds of 1 exclusively.
     array->SetLowerBoundsToOne();
     ConstantSubscripts at{array->lbounds()}, maskAt, resultIndices, resultShape;
     if (mask) {
       mask->SetLowerBoundsToOne();
       maskAt = mask->lbounds();
     }
     if (dim) { // DIM=
       if (*dim < 1 || *dim > array->Rank()) {
         context_.messages().Say(
             "FINDLOC(DIM=%d) is out of range"_err_en_US, *dim);
         return std::nullopt;
       }
       int zbDim{*dim - 1};
       resultShape = array->shape();
       resultShape.erase(
           resultShape.begin() + zbDim); // scalar if array is vector
       ConstantSubscript dimLength{array->shape()[zbDim]};
       ConstantSubscript n{GetSize(resultShape)};
       for (ConstantSubscript j{0}; j < n; ++j) {
         ConstantSubscript hit{array->lbounds()[zbDim] - 1};
+        value.reset();
         for (ConstantSubscript k{0}; k < dimLength;
              ++k, ++at[zbDim], mask && ++maskAt[zbDim]) {
           if ((!mask || mask->At(maskAt).IsTrue()) &&
-              IsHit(array->At(at), *value)) {
+              IsHit(array->At(at), value, relation)) {
             hit = at[zbDim];
             if (!back) {
               break;
             }
           }
         }
         resultIndices.emplace_back(hit);
         at[zbDim] = array->lbounds()[zbDim] + dimLength - 1;
         array->IncrementSubscripts(at);
         at[zbDim] = array->lbounds()[zbDim];
         if (mask) {
           maskAt[zbDim] = mask->lbounds()[zbDim] + dimLength - 1;
           mask->IncrementSubscripts(maskAt);
           maskAt[zbDim] = mask->lbounds()[zbDim];
         }
       }
     } else { // no DIM=
       resultShape = ConstantSubscripts{array->Rank()}; // always a vector
       ConstantSubscript n{GetSize(array->shape())};
       resultIndices = ConstantSubscripts(array->Rank(), 0);
       for (ConstantSubscript j{0}; j < n; ++j, array->IncrementSubscripts(at),
            mask && mask->IncrementSubscripts(maskAt)) {
         if ((!mask || mask->At(maskAt).IsTrue()) &&
-            IsHit(array->At(at), *value)) {
+            IsHit(array->At(at), value, relation)) {
           resultIndices = at;
           if (!back) {
             break;
           }
         }
       }
     }
     std::vector<Scalar<SubscriptInteger>> resultElements;
     for (ConstantSubscript j : resultIndices) {
       resultElements.emplace_back(j);
     }
     return Constant<SubscriptInteger>{
         std::move(resultElements), std::move(resultShape)};
   }
 
 private:
   template <typename T>
-  bool IsHit(typename Constant<T>::Element element, Constant<T> value) const {
+  bool IsHit(typename Constant<T>::Element element,
+      std::optional<Constant<T>> &value,
+      [[maybe_unused]] RelationalOperator relation) const {
     std::optional<Expr<LogicalResult>> cmp;
+    if (value) {
       if constexpr (T::category == TypeCategory::Logical) {
         // array(at) .EQV. value?
+        static_assert(WHICH == WhichLocation::Findloc);
         cmp.emplace(
             ConvertToType<LogicalResult>(Expr<T>{LogicalOperation<T::kind>{
                 LogicalOperator::Eqv, Expr<T>{Constant<T>{std::move(element)}},
-              Expr<T>{std::move(value)}}}));
-    } else { // array(at) .EQ. value?
-      cmp.emplace(PackageRelation(RelationalOperator::EQ,
-          Expr<T>{Constant<T>{std::move(element)}}, Expr<T>{std::move(value)}));
+                Expr<T>{Constant<T>{*value}}}}));
+      } else { // compare array(at) to value
+        cmp.emplace(
+            PackageRelation(relation, Expr<T>{Constant<T>{std::move(element)}},
+                Expr<T>{Constant<T>{*value}}));
       }
       Expr<LogicalResult> folded{Fold(context_, std::move(*cmp))};
       return GetScalarConstantValue<LogicalResult>(folded).value().IsTrue();
+    } else { // first unmasked element seen for MAXLOC/MINLOC
+      value.emplace(std::move(element));
+      return true;
+    }
   }
 
+  static constexpr int dimArg{WHICH == WhichLocation::Findloc ? 2 : 1};
+  static constexpr int maskArg{dimArg + 1};
+  static constexpr int backArg{maskArg + 2};
+
   DynamicType type_;
   ActualArguments &arg_;
   FoldingContext &context_;
 };
 
-static std::optional<Constant<SubscriptInteger>> FoldFindlocCall(
+template <WhichLocation which>
+static std::optional<Constant<SubscriptInteger>> FoldLocationCall(
     ActualArguments &arg, FoldingContext &context) {
-  CHECK(arg.size() == 6);
   if (arg[0]) {
     if (auto type{arg[0]->GetType()}) {
-      return common::SearchTypes(FindlocHelper{std::move(*type), arg, context});
+      return common::SearchTypes(
+          LocationHelper<which>{std::move(*type), arg, context});
     }
   }
   return std::nullopt;
 }
 
-template <typename T>
-static Expr<T> FoldFindloc(FoldingContext &context, FunctionRef<T> &&ref) {
+template <WhichLocation which, typename T>
+static Expr<T> FoldLocation(FoldingContext &context, FunctionRef<T> &&ref) {
   static_assert(T::category == TypeCategory::Integer);
   if (std::optional<Constant<SubscriptInteger>> found{
-          FoldFindlocCall(ref.arguments(), context)}) {
+          FoldLocationCall<which>(ref.arguments(), context)}) {
     return Expr<T>{Fold(
         context, ConvertToType<T>(Expr<SubscriptInteger>{std::move(*found)}))};
   } else {
     return Expr<T>{std::move(ref)};
   }
 }
 
 // for IALL, IANY, & IPARITY
if (auto *sx{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
             return FoldElementalIntrinsic<T, TR>(context, std::move(funcRef),
                 &Scalar<TR>::template EXPONENT<Scalar<T>>);
           },
           sx->u);
     } else {
       DIE("exponent argument must be real");
     }
   } else if (name == "findloc") {
-    return FoldFindloc<T>(context, std::move(funcRef));
+    return FoldLocation<WhichLocation::Findloc, T>(context, std::move(funcRef));
   } else if (name == "huge") {
     return Expr<T>{Scalar<T>::HUGE()};
   } else if (name == "iachar" || name == "ichar") {
     auto *someChar{UnwrapExpr<Expr<SomeCharacter>>(args[0])};
     CHECK(someChar);
     if (auto len{ToInt64(someChar->LEN())}) {
       if (len.value() != 1) {
         // Do not die, this was not checked before
if (name == "abs") {
     if (auto *sx{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
       return std::visit(
           [](const auto &x) {
             using TR = typename std::decay_t<decltype(x)>::Result;
             return Expr<T>{Scalar<TR>::MAXEXPONENT};
           },
           sx->u);
     }
+  } else if (name == "maxloc") {
+    return FoldLocation<WhichLocation::Maxloc, T>(context, std::move(funcRef));
   } else if (name == "maxval") {
     return FoldMaxvalMinval<T>(context, std::move(funcRef),
         RelationalOperator::GT, T::Scalar::Least());
   } else if (name == "merge") {
     return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "merge_bits") {
     return FoldElementalIntrinsic<T, T, T, T>(
         context, std::move(funcRef), &Scalar<T>::MERGE_BITS);
+  } else if (name == "min") {
+    return FoldMINorMAX(context, std::move(funcRef), Ordering::Less);
+  } else if (name == "min0" || name == "min1") {
+    return RewriteSpecificMINorMAX(context, std::move(funcRef));
   } else if (name == "minexponent") {
     if (auto *sx{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
       return std::visit(
           [](const auto &x) {
             using TR = typename std::decay_t<decltype(x)>::Result;
             return Expr<T>{Scalar<TR>::MINEXPONENT};
           },
           sx->u);
     }
-  } else if (name == "min") {
-    return FoldMINorMAX(context, std::move(funcRef), Ordering::Less);
-  } else if (name == "min0" || name == "min1") {
-    return RewriteSpecificMINorMAX(context, std::move(funcRef));
+  } else if (name == "minloc") {
+    return FoldLocation<WhichLocation::Minloc, T>(context, std::move(funcRef));
   } else if (name == "minval") {
     return FoldMaxvalMinval<T>(
         context, std::move(funcRef), RelationalOperator::LT, T::Scalar::HUGE());
   } else if (name == "mod") {
     return FoldElementalIntrinsic<T, T, T>(context, std::move(funcRef),
         ScalarFuncWithContext<T, T, T>(
             [](FoldingContext &context, const Scalar<T> &x,
                 const Scalar<T> &y) -> Scalar<T> {
if (auto info{
             Fold(context, Expr<T>{8} * ConvertToType<T>(std::move(*bytes)))};
       }
     }
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   } else if (name == "ubound") {
     return UBOUND(context, std::move(funcRef));
   }
-  // TODO: dot_product, ibits, image_status, ishftc,
-  // matmul, maxloc, minloc, sign, transfer
+  // TODO: dot_product, ibits, ishftc, matmul, sign, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 // Substitutes a bare type parameter reference with its value if it has one now
 // in an instantiation.  Bare LEN type parameters are substituted only when
 // the known value is constant.
 Expr<TypeParamInquiry::Result> FoldOperation(
     FoldingContext &context, TypeParamInquiry &&inquiry) {

flang/lib/Evaluate/fold-real.cpp

return FoldElementalIntrinsic<T, T>(context, std::move(funcRef),
         ScalarFunc<T, T>(
             [](const Scalar<T> &x) -> Scalar<T> { return x.SQRT().value; }));
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   } else if (name == "tiny") {
     return Expr<T>{Scalar<T>::TINY()};
   }
   // TODO: dim, dot_product, fraction, matmul,
-  // maxloc, minloc, modulo, nearest, norm2, rrspacing, scale,
+  // modulo, nearest, norm2, rrspacing, scale,
   // set_exponent, spacing, transfer,
   // bessel_jn (transformational) and bessel_yn (transformational)
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Real, KIND>> FoldOperation(
     FoldingContext &context, ComplexComponent<KIND> &&x) {

flang/lib/Evaluate/fold-reduction.h

 //===-- lib/Evaluate/fold-reduction.h -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-// TODO: DOT_PRODUCT, NORM2, MAXLOC, MINLOC, PARITY
+// TODO: DOT_PRODUCT, NORM2, PARITY
 
 #ifndef FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 #define FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 
 #include "fold-implementation.h"
 
 namespace Fortran::evaluate {
 

flang/test/Evaluate/folding30.f90

 ! RUN: %python %S/test_folding.py %s %flang_fc1
-! Tests folding of FINDLOC
+! Tests folding of FINDLOC, MAXLOC, & MINLOC
 module m1
   integer, parameter :: ia1(2:6) = [1, 2, 3, 2, 1]
   integer, parameter :: ia2(2:3,2:4) = reshape([1, 2, 3, 3, 2, 1], shape(ia2))
 
-  logical, parameter :: ti1a = all(findloc(ia1, 1) == 1)
-  logical, parameter :: ti1ar = rank(findloc(ia1, 1)) == 1
-  logical, parameter :: ti1ak = kind(findloc(ia1, 1, kind=2)) == 2
-  logical, parameter :: ti1ad = findloc(ia1, 1, dim=1) == 1
-  logical, parameter :: ti1adr = rank(findloc(ia1, 1, dim=1)) == 0
-  logical, parameter :: ti1b = all(findloc(ia1, 1, back=.true.) == 5)
-  logical, parameter :: ti1c = all(findloc(ia1, 2, mask=[.false., .false., .true., .true., .true.]) == 4)
-
-  logical, parameter :: ti2a = all(findloc(ia2, 1) == [1, 1])
-  logical, parameter :: ti2ar = rank(findloc(ia2, 1)) == 1
-  logical, parameter :: ti2b = all(findloc(ia2, 1, back=.true.) == [2, 3])
-  logical, parameter :: ti2c = all(findloc(ia2, 2, mask=reshape([.false., .false., .true., .true., .true., .false.], shape(ia2))) == [1, 3])
-  logical, parameter :: ti2d = all(findloc(ia2, 1, dim=1) == [1, 0, 2])
-  logical, parameter :: ti2e = all(findloc(ia2, 1, dim=2) == [1, 3])
+  logical, parameter :: fi1a = all(findloc(ia1, 1) == 1)
+  logical, parameter :: fi1ar = rank(findloc(ia1, 1)) == 1
+  logical, parameter :: fi1ak = kind(findloc(ia1, 1, kind=2)) == 2
+  logical, parameter :: fi1ad = findloc(ia1, 1, dim=1) == 1
+  logical, parameter :: fi1adr = rank(findloc(ia1, 1, dim=1)) == 0
+  logical, parameter :: fi1b = all(findloc(ia1, 1, back=.true.) == 5)
+  logical, parameter :: fi1c = all(findloc(ia1, 2, mask=[.false., .false., .true., .true., .true.]) == 4)
+
+  logical, parameter :: fi2a = all(findloc(ia2, 1) == [1, 1])
+  logical, parameter :: fi2ar = rank(findloc(ia2, 1)) == 1
+  logical, parameter :: fi2b = all(findloc(ia2, 1, back=.true.) == [2, 3])
+  logical, parameter :: fi2c = all(findloc(ia2, 2, mask=reshape([.false., .false., .true., .true., .true., .false.], shape(ia2))) == [1, 3])
+  logical, parameter :: fi2d = all(findloc(ia2, 1, dim=1) == [1, 0, 2])
+  logical, parameter :: fi2e = all(findloc(ia2, 1, dim=2) == [1, 3])
+
+  logical, parameter :: xi1a = all(maxloc(ia1) == 3)
+  logical, parameter :: xi1ar = rank(maxloc(ia1)) == 1
+  logical, parameter :: xi1ak = kind(maxloc(ia1, kind=2)) == 2
+  logical, parameter :: xi1ad = maxloc(ia1, dim=1) == 3
+  logical, parameter :: xi1adr = rank(maxloc(ia1, dim=1)) == 0
+  logical, parameter :: xi1b = all(maxloc(ia1, back=.true.) == 3)
+  logical, parameter :: xi1c = all(maxloc(ia1, mask=[.false., .true., .false., .true., .true.]) == 2)
+  logical, parameter :: xi1d = all(maxloc(ia1, mask=[.false., .true., .false., .true., .true.], back=.true.) == 4)
+
+  logical, parameter :: xi2a = all(maxloc(ia2) == [1, 2])
+  logical, parameter :: xi2ar = rank(maxloc(ia2)) == 1
+  logical, parameter :: xi2b = all(maxloc(ia2, back=.true.) == [2, 2])
+  logical, parameter :: xi2c = all(maxloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .true.], shape(ia2))) == [2, 1])
+  logical, parameter :: xi2d = all(maxloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .true.], shape(ia2)), back=.true.) == [1, 3])
+  logical, parameter :: xi2e = all(maxloc(ia2, dim=1) == [2, 1, 1])
+  logical, parameter :: xi2f = all(maxloc(ia2, dim=1, back=.true.) == [2, 2, 1])
+  logical, parameter :: xi2g = all(maxloc(ia2, dim=2) == [2, 2])
+
+  logical, parameter :: ni1a = all(minloc(ia1) == 1)
+  logical, parameter :: ni1ar = rank(minloc(ia1)) == 1
+  logical, parameter :: ni1ak = kind(minloc(ia1, kind=2)) == 2
+  logical, parameter :: ni1ad = minloc(ia1, dim=1) == 1
+  logical, parameter :: ni1adr = rank(minloc(ia1, dim=1)) == 0
+  logical, parameter :: ni1b = all(minloc(ia1, back=.true.) == 5)
+  logical, parameter :: ni1c = all(minloc(ia1, mask=[.false., .true., .true., .true., .false.]) == 2)
+  logical, parameter :: ni1d = all(minloc(ia1, mask=[.false., .true., .true., .true., .false.], back=.true.) == 4)
+
+  logical, parameter :: ni2a = all(minloc(ia2) == [1, 1])
+  logical, parameter :: ni2ar = rank(minloc(ia2)) == 1
+  logical, parameter :: ni2b = all(minloc(ia2, back=.true.) == [2, 3])
+  logical, parameter :: ni2c = all(minloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .false.], shape(ia2))) == [2, 1])
+  logical, parameter :: ni2d = all(minloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .false.], shape(ia2)), back=.true.) == [1, 3])
+  logical, parameter :: ni2e = all(minloc(ia2, dim=1) == [1, 1, 2])
+  logical, parameter :: ni2f = all(minloc(ia2, dim=1, back=.true.) == [1, 2, 2])
+  logical, parameter :: ni2g = all(minloc(ia2, dim=2) == [1, 3])
 end module