[flang] Support FINDLOC/MAXLOC/MINLOC with scalar mask

Previously MASK= elements were accessed in assumption that mask is an array of
input argument rank (and in combination with explicit DIM= argument we had
out-of-bounds access), but for MAXLOC/MINLOC/FINDLOC mask should be be
conformable and could be scalar.

Add new regression tests with scalar mask for verification.

Reviewed By: klausler

Differential Revision: https://reviews.llvm.org/D124408
This commit is contained in:
Mike Kashkarov 2022-04-28 18:57:51 +09:00
parent 92e22c97e9
commit 35cc2ec4ed
2 changed files with 24 additions and 0 deletions

View File

@ -313,6 +313,15 @@ public:
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();
}

View File

@ -66,4 +66,19 @@ module m1
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