[flang-commits] [flang] ca0525f - [flang] Fix crash in location reduction folding with DIM= and scalar MASK=

[Peter Klausler via flang-commits]

Author: Peter Klausler
Date: 2023-08-08T10:12:54-07:00
New Revision: ca0525fa16a4b3b003d3dac0278c2455abb38bee

URL: https://github.com/llvm/llvm-project/commit/ca0525fa16a4b3b003d3dac0278c2455abb38bee
DIFF: https://github.com/llvm/llvm-project/commit/ca0525fa16a4b3b003d3dac0278c2455abb38bee.diff

LOG: [flang] Fix crash in location reduction folding with DIM= and scalar MASK=

The shape of an expanded scalar MASK= was incorrect in the template
code that folds FINDLOC/MAXLOC/MINLOC, leading to a vector indexing
crash when a DIM= argument is also present.

Fixes https://github.com/llvm/llvm-project/issues/64286.

Differential Revision: https://reviews.llvm.org/D157329

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/flang/lib/Evaluate/fold-integer.cpp b/flang/lib/Evaluate/fold-integer.cpp
index 7f37577487d35d..02d5ea5a133ada 100644
--- a/flang/lib/Evaluate/fold-integer.cpp
+++ b/flang/lib/Evaluate/fold-integer.cpp
@@ -352,12 +352,12 @@ template <WhichLocation WHICH> class LocationHelper {
     if (mask) {
       if (auto scalarMask{mask->GetScalarValue()}) {
         // Convert into array in case of scalar MASK= (for
-        // MAXLOC/MINLOC/FINDLOC mask should be be conformable)
+        // MAXLOC/MINLOC/FINDLOC mask should 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}};
+            std::move(mask_elements), ConstantSubscripts{array->shape()}};
       }
       mask->SetLowerBoundsToOne();
       maskAt = mask->lbounds();

diff  --git a/flang/test/Evaluate/fold-findloc.f90 b/flang/test/Evaluate/fold-findloc.f90
index 7297081bdfae73..b8bb85af65dc6c 100644
--- a/flang/test/Evaluate/fold-findloc.f90
+++ b/flang/test/Evaluate/fold-findloc.f90
@@ -74,6 +74,8 @@ module m1
   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_fia2_mtd1 = all(findloc(ia2, 1, dim=1, mask=.true.) == [1, 0, 2])
+  logical, parameter:: test_fia2_mtd2 = all(findloc(ia2, 1, dim=2, mask=.true.) == [1, 3])
 
   logical, parameter:: test_mia1_mf  = all(minloc(ia1, mask=.false.) == 0)
   logical, parameter:: test_mia1_mfd = all(minloc(ia1, mask=.false., dim=1) == [0])
@@ -81,4 +83,6 @@ module m1
   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])
+  logical, parameter:: test_fia2_mfd1 = all(findloc(ia2, 1, dim=1, mask=.false.) == [0, 0, 0])
+  logical, parameter:: test_fia2_mfd2 = all(findloc(ia2, 1, dim=2, mask=.false.) == [0, 0])
 end module