Skip to content

Commit

Permalink
[flang] Correct folding of CSHIFT and EOSHIFT for DIM>1
Browse files Browse the repository at this point in the history 
The algorithm was wrong for higher dimensions, and so were
the expected test results.  Rework.

Differential Revision: https://reviews.llvm.org/D127018
  • Loading branch information
@klausler
klausler committed Jun 4, 2022
1 parent 47ec8b5 commit e0adee8
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 49 deletions.
112 changes: 67 additions & 45 deletions flang/lib/Evaluate/fold-implementation.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -613,26 +613,33 @@ template <typename T> Expr<T> Folder<T>::CSHIFT(FunctionRef<T> &&funcRef) {
}
if (ok) {
std::vector<Scalar<T>> resultElements;
ConstantSubscripts arrayAt{array->lbounds()};
ConstantSubscript dimLB{arrayAt[zbDim]};
ConstantSubscripts arrayLB{array->lbounds()};
ConstantSubscripts arrayAt{arrayLB};
ConstantSubscript &dimIndex{arrayAt[zbDim]};
ConstantSubscript dimLB{dimIndex}; // initial value
ConstantSubscript dimExtent{array->shape()[zbDim]};
ConstantSubscripts shiftAt{shift->lbounds()};
for (auto n{GetSize(array->shape())}; n > 0; n -= dimExtent) {
ConstantSubscript shiftCount{shift->At(shiftAt).ToInt64()};
ConstantSubscript zbDimIndex{shiftCount % dimExtent};
if (zbDimIndex < 0) {
zbDimIndex += dimExtent;
}
for (ConstantSubscript j{0}; j < dimExtent; ++j) {
arrayAt[zbDim] = dimLB + zbDimIndex;
resultElements.push_back(array->At(arrayAt));
if (++zbDimIndex == dimExtent) {
zbDimIndex = 0;
ConstantSubscripts shiftLB{shift->lbounds()};
for (auto n{GetSize(array->shape())}; n > 0; --n) {
ConstantSubscript origDimIndex{dimIndex};
ConstantSubscripts shiftAt;
if (shift->Rank() > 0) {
int k{0};
for (int j{0}; j < rank; ++j) {
if (j != zbDim) {
shiftAt.emplace_back(shiftLB[k++] + arrayAt[j] - arrayLB[j]);
}
}
}
arrayAt[zbDim] = dimLB + std::max<ConstantSubscript>(dimExtent, 1) - 1;
ConstantSubscript shiftCount{shift->At(shiftAt).ToInt64()};
dimIndex = dimLB + ((dimIndex - dimLB + shiftCount) % dimExtent);
if (dimIndex < dimLB) {
dimIndex += dimExtent;
} else if (dimIndex >= dimLB + dimExtent) {
dimIndex -= dimExtent;
}
resultElements.push_back(array->At(arrayAt));
dimIndex = origDimIndex;
array->IncrementSubscripts(arrayAt);
shift->IncrementSubscripts(shiftAt);
}
return Expr<T>{PackageConstant<T>(
std::move(resultElements), *array, array->shape())};
Expand Down Expand Up @@ -714,42 +721,57 @@ template <typename T> Expr<T> Folder<T>::EOSHIFT(FunctionRef<T> &&funcRef) {
}
if (ok) {
std::vector<Scalar<T>> resultElements;
ConstantSubscripts arrayAt{array->lbounds()};
ConstantSubscript dimLB{arrayAt[zbDim]};
ConstantSubscripts arrayLB{array->lbounds()};
ConstantSubscripts arrayAt{arrayLB};
ConstantSubscript &dimIndex{arrayAt[zbDim]};
ConstantSubscript dimLB{dimIndex}; // initial value
ConstantSubscript dimExtent{array->shape()[zbDim]};
ConstantSubscripts shiftAt{shift->lbounds()};
ConstantSubscripts boundaryAt;
ConstantSubscripts shiftLB{shift->lbounds()};
ConstantSubscripts boundaryLB;
if (boundary) {
boundaryAt = boundary->lbounds();
boundaryLB = boundary->lbounds();
}
for (auto n{GetSize(array->shape())}; n > 0; n -= dimExtent) {
for (auto n{GetSize(array->shape())}; n > 0; --n) {
ConstantSubscript origDimIndex{dimIndex};
ConstantSubscripts shiftAt;
if (shift->Rank() > 0) {
int k{0};
for (int j{0}; j < rank; ++j) {
if (j != zbDim) {
shiftAt.emplace_back(shiftLB[k++] + arrayAt[j] - arrayLB[j]);
}
}
}
ConstantSubscript shiftCount{shift->At(shiftAt).ToInt64()};
for (ConstantSubscript j{0}; j < dimExtent; ++j) {
ConstantSubscript zbAt{shiftCount + j};
if (zbAt >= 0 && zbAt < dimExtent) {
arrayAt[zbDim] = dimLB + zbAt;
resultElements.push_back(array->At(arrayAt));
} else if (boundary) {
resultElements.push_back(boundary->At(boundaryAt));
} else if constexpr (T::category == TypeCategory::Integer ||
T::category == TypeCategory::Real ||
T::category == TypeCategory::Complex ||
T::category == TypeCategory::Logical) {
resultElements.emplace_back();
} else if constexpr (T::category == TypeCategory::Character) {
auto len{static_cast<std::size_t>(array->LEN())};
typename Scalar<T>::value_type space{' '};
resultElements.emplace_back(len, space);
} else {
DIE("no derived type boundary");
dimIndex += shiftCount;
if (dimIndex >= dimLB && dimIndex < dimLB + dimExtent) {
resultElements.push_back(array->At(arrayAt));
} else if (boundary) {
ConstantSubscripts boundaryAt;
if (boundary->Rank() > 0) {
for (int j{0}; j < rank; ++j) {
int k{0};
if (j != zbDim) {
boundaryAt.emplace_back(
boundaryLB[k++] + arrayAt[j] - arrayLB[j]);
}
}
}
resultElements.push_back(boundary->At(boundaryAt));
} else if constexpr (T::category == TypeCategory::Integer ||
T::category == TypeCategory::Real ||
T::category == TypeCategory::Complex ||
T::category == TypeCategory::Logical) {
resultElements.emplace_back();
} else if constexpr (T::category == TypeCategory::Character) {
auto len{static_cast<std::size_t>(array->LEN())};
typename Scalar<T>::value_type space{' '};
resultElements.emplace_back(len, space);
} else {
DIE("no derived type boundary");
}
arrayAt[zbDim] = dimLB + std::max<ConstantSubscript>(dimExtent, 1) - 1;
dimIndex = origDimIndex;
array->IncrementSubscripts(arrayAt);
shift->IncrementSubscripts(shiftAt);
if (boundary) {
boundary->IncrementSubscripts(boundaryAt);
}
}
return Expr<T>{PackageConstant<T>(
std::move(resultElements), *array, array->shape())};
Expand Down
4 changes: 2 additions & 2 deletions flang/test/Evaluate/folding23.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ module m
logical, parameter :: test_eoshift_3 = all(eoshift([1., 2., 3.], 1) == [2., 3., 0.])
logical, parameter :: test_eoshift_4 = all(eoshift(['ab', 'cd', 'ef'], -1, 'x') == ['x ', 'ab', 'cd'])
logical, parameter :: test_eoshift_5 = all([eoshift(arr, 1, dim=1)] == [2, 0, 4, 0, 6, 0])
logical, parameter :: test_eoshift_6 = all([eoshift(arr, 1, dim=2)] == [3, 5, 0, 4, 6, 0])
logical, parameter :: test_eoshift_6 = all([eoshift(arr, 1, dim=2)] == [3, 4, 5, 6, 0, 0])
logical, parameter :: test_eoshift_7 = all([eoshift(arr, [1, -1, 0])] == [2, 0, 0, 3, 5, 6])
logical, parameter :: test_eoshift_8 = all([eoshift(arr, [1, -1], dim=2)] == [3, 5, 0, 0, 2, 4])
logical, parameter :: test_eoshift_8 = all([eoshift(arr, [1, -1], dim=2)] == [3, 0, 5, 2, 0, 4])
end module
4 changes: 2 additions & 2 deletions flang/test/Evaluate/folding27.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ module m
logical, parameter :: test_cshift_3 = all(cshift([1, 2, 3], 4) == [2, 3, 1])
logical, parameter :: test_cshift_4 = all(cshift([1, 2, 3], -1) == [3, 1, 2])
logical, parameter :: test_cshift_5 = all([cshift(arr, 1, dim=1)] == [2, 1, 4, 3, 6, 5])
logical, parameter :: test_cshift_6 = all([cshift(arr, 1, dim=2)] == [3, 5, 1, 4, 6, 2])
logical, parameter :: test_cshift_6 = all([cshift(arr, 1, dim=2)] == [3, 4, 5, 6, 1, 2])
logical, parameter :: test_cshift_7 = all([cshift(arr, [1, 2, 3])] == [2, 1, 3, 4, 6, 5])
logical, parameter :: test_cshift_8 = all([cshift(arr, [1, 2], dim=2)] == [3, 5, 1, 6, 2, 4])
logical, parameter :: test_cshift_8 = all([cshift(arr, [1, 2], dim=2)] == [3, 6, 5, 2, 1, 4])
end module

0 comments on commit e0adee8

Please sign in to comment.