[flang-commits] [PATCH] D124408: [flang] Support FINDLOC/MAXLOC/MINLOC with scalar mask

[Mike K via Phabricator via flang-commits]

FruitClover updated this revision to Diff 425173.
FruitClover added a comment.

clang-format


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D124408/new/

https://reviews.llvm.org/D124408

Files:
  flang/lib/Evaluate/fold-integer.cpp
  flang/test/Evaluate/fold-findloc.f90


Index: flang/test/Evaluate/fold-findloc.f90
===================================================================
--- flang/test/Evaluate/fold-findloc.f90
+++ flang/test/Evaluate/fold-findloc.f90
@@ -66,4 +66,19 @@
   logical, parameter :: test_char2 = all(minloc(a).eq.[1])
   logical, parameter :: test_char3 = all(maxloc(a, back=.true.).eq.[4])
   logical, parameter :: test_char4 = all(minloc(a, back=.true.).eq.[3])
+
+  ! Check with scalar MASK=
+  logical, parameter:: test_mia1_mt  = all(minloc(ia1, mask=.true.) == 1)
+  logical, parameter:: test_mia1_mtd = all(minloc(ia1, mask=.true., dim=1) == [1])
+  logical, parameter:: test_xia1_mt  = all(maxloc(ia1, mask=.true.) == 3)
+  logical, parameter:: test_xia1_mtd = all(maxloc(ia1, mask=.true., dim=1) == [3])
+  logical, parameter:: test_fia1_mt  = all(findloc(ia1, 1, mask=.true.) == 1)
+  logical, parameter:: test_fia1_mtd = all(findloc(ia1, 1, mask=.true., dim=1) == [1])
+
+  logical, parameter:: test_mia1_mf  = all(minloc(ia1, mask=.false.) == 0)
+  logical, parameter:: test_mia1_mfd = all(minloc(ia1, mask=.false., dim=1) == [0])
+  logical, parameter:: test_xia1_mf  = all(maxloc(ia1, mask=.false.) == 0)
+  logical, parameter:: test_xia1_mfd = all(maxloc(ia1, mask=.false., dim=1) == [0])
+  logical, parameter:: test_fia1_mf  = all(findloc(ia1, 1, mask=.false.) == 0)
+  logical, parameter:: test_fia1_mfd = all(findloc(ia1, 1, mask=.false., dim=1) == [0])
 end module
Index: flang/lib/Evaluate/fold-integer.cpp
===================================================================
--- flang/lib/Evaluate/fold-integer.cpp
+++ flang/lib/Evaluate/fold-integer.cpp
@@ -272,6 +272,15 @@
     array->SetLowerBoundsToOne();
     ConstantSubscripts at{array->lbounds()}, maskAt, resultIndices, resultShape;
     if (mask) {
+      if (auto scalarMask{mask->GetScalarValue()}) {
+        // Convert into array in case of scalar MASK= (for
+        // MAXLOC/MINLOC/FINDLOC mask should be be conformable)
+        ConstantSubscript n{GetSize(array->shape())};
+        std::vector<Scalar<LogicalResult>> mask_elements(
+            n, Scalar<LogicalResult>{scalarMask.value()});
+        *mask = Constant<LogicalResult>{
+            std::move(mask_elements), ConstantSubscripts{n}};
+      }
       mask->SetLowerBoundsToOne();
       maskAt = mask->lbounds();
     }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D124408.425173.patch
Type: text/x-patch
Size: 2329 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/flang-commits/attachments/20220426/29af69bc/attachment-0001.bin>