1 //===-- lib/Evaluate/fold-reduction.h -------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 // TODO: DOT_PRODUCT, NORM2, PARITY
10 
11 #ifndef FORTRAN_EVALUATE_FOLD_REDUCTION_H_
12 #define FORTRAN_EVALUATE_FOLD_REDUCTION_H_
13 
14 #include "fold-implementation.h"
15 
16 namespace Fortran::evaluate {
17 
18 // Fold and validate a DIM= argument.  Returns false on error.
19 bool CheckReductionDIM(std::optional<int> &dim, FoldingContext &,
20     ActualArguments &, std::optional<int> dimIndex, int rank);
21 
22 // Fold and validate a MASK= argument.  Return null on error, absent MASK=, or
23 // non-constant MASK=.
24 Constant<LogicalResult> *GetReductionMASK(
25     std::optional<ActualArgument> &maskArg, const ConstantSubscripts &shape,
26     FoldingContext &);
27 
28 // Common preprocessing for reduction transformational intrinsic function
29 // folding.  If the intrinsic can have DIM= &/or MASK= arguments, extract
30 // and check them.  If a MASK= is present, apply it to the array data and
31 // substitute identity values for elements corresponding to .FALSE. in
32 // the mask.  If the result is present, the intrinsic call can be folded.
33 template <typename T>
34 static std::optional<Constant<T>> ProcessReductionArgs(FoldingContext &context,
35     ActualArguments &arg, std::optional<int> &dim, const Scalar<T> &identity,
36     int arrayIndex, std::optional<int> dimIndex = std::nullopt,
37     std::optional<int> maskIndex = std::nullopt) {
38   if (arg.empty()) {
39     return std::nullopt;
40   }
41   Constant<T> *folded{Folder<T>{context}.Folding(arg[arrayIndex])};
42   if (!folded || folded->Rank() < 1) {
43     return std::nullopt;
44   }
45   if (!CheckReductionDIM(dim, context, arg, dimIndex, folded->Rank())) {
46     return std::nullopt;
47   }
48   if (maskIndex && static_cast<std::size_t>(*maskIndex) < arg.size() &&
49       arg[*maskIndex]) {
50     if (const Constant<LogicalResult> *mask{
51             GetReductionMASK(arg[*maskIndex], folded->shape(), context)}) {
52       // Apply the mask in place to the array
53       std::size_t n{folded->size()};
54       std::vector<typename Constant<T>::Element> elements;
55       if (auto scalarMask{mask->GetScalarValue()}) {
56         if (scalarMask->IsTrue()) {
57           return Constant<T>{*folded};
58         } else { // MASK=.FALSE.
59           elements = std::vector<typename Constant<T>::Element>(n, identity);
60         }
61       } else { // mask is an array; test its elements
62         elements = std::vector<typename Constant<T>::Element>(n, identity);
63         ConstantSubscripts at{folded->lbounds()};
64         for (std::size_t j{0}; j < n; ++j, folded->IncrementSubscripts(at)) {
65           if (mask->values()[j].IsTrue()) {
66             elements[j] = folded->At(at);
67           }
68         }
69       }
70       if constexpr (T::category == TypeCategory::Character) {
71         return Constant<T>{static_cast<ConstantSubscript>(identity.size()),
72             std::move(elements), ConstantSubscripts{folded->shape()}};
73       } else {
74         return Constant<T>{
75             std::move(elements), ConstantSubscripts{folded->shape()}};
76       }
77     } else {
78       return std::nullopt;
79     }
80   } else {
81     return Constant<T>{*folded};
82   }
83 }
84 
85 // Generalized reduction to an array of one dimension fewer (w/ DIM=)
86 // or to a scalar (w/o DIM=).
87 template <typename T, typename ACCUMULATOR, typename ARRAY>
DoReduction(const Constant<ARRAY> & array,std::optional<int> & dim,const Scalar<T> & identity,ACCUMULATOR & accumulator)88 static Constant<T> DoReduction(const Constant<ARRAY> &array,
89     std::optional<int> &dim, const Scalar<T> &identity,
90     ACCUMULATOR &accumulator) {
91   ConstantSubscripts at{array.lbounds()};
92   std::vector<typename Constant<T>::Element> elements;
93   ConstantSubscripts resultShape; // empty -> scalar
94   if (dim) { // DIM= is present, so result is an array
95     resultShape = array.shape();
96     resultShape.erase(resultShape.begin() + (*dim - 1));
97     ConstantSubscript dimExtent{array.shape().at(*dim - 1)};
98     ConstantSubscript &dimAt{at[*dim - 1]};
99     ConstantSubscript dimLbound{dimAt};
100     for (auto n{GetSize(resultShape)}; n-- > 0;
101          IncrementSubscripts(at, array.shape())) {
102       dimAt = dimLbound;
103       elements.push_back(identity);
104       for (ConstantSubscript j{0}; j < dimExtent; ++j, ++dimAt) {
105         accumulator(elements.back(), at);
106       }
107     }
108   } else { // no DIM=, result is scalar
109     elements.push_back(identity);
110     for (auto n{array.size()}; n-- > 0;
111          IncrementSubscripts(at, array.shape())) {
112       accumulator(elements.back(), at);
113     }
114   }
115   if constexpr (T::category == TypeCategory::Character) {
116     return {static_cast<ConstantSubscript>(identity.size()),
117         std::move(elements), std::move(resultShape)};
118   } else {
119     return {std::move(elements), std::move(resultShape)};
120   }
121 }
122 
123 // MAXVAL & MINVAL
124 template <typename T>
FoldMaxvalMinval(FoldingContext & context,FunctionRef<T> && ref,RelationalOperator opr,const Scalar<T> & identity)125 static Expr<T> FoldMaxvalMinval(FoldingContext &context, FunctionRef<T> &&ref,
126     RelationalOperator opr, const Scalar<T> &identity) {
127   static_assert(T::category == TypeCategory::Integer ||
128       T::category == TypeCategory::Real ||
129       T::category == TypeCategory::Character);
130   using Element = Scalar<T>;
131   std::optional<int> dim;
132   if (std::optional<Constant<T>> array{
133           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
134               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
135     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
136       Expr<LogicalResult> test{PackageRelation(opr,
137           Expr<T>{Constant<T>{array->At(at)}}, Expr<T>{Constant<T>{element}})};
138       auto folded{GetScalarConstantValue<LogicalResult>(
139           test.Rewrite(context, std::move(test)))};
140       CHECK(folded.has_value());
141       if (folded->IsTrue()) {
142         element = array->At(at);
143       }
144     }};
145     return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
146   }
147   return Expr<T>{std::move(ref)};
148 }
149 
150 // PRODUCT
151 template <typename T>
FoldProduct(FoldingContext & context,FunctionRef<T> && ref,Scalar<T> identity)152 static Expr<T> FoldProduct(
153     FoldingContext &context, FunctionRef<T> &&ref, Scalar<T> identity) {
154   static_assert(T::category == TypeCategory::Integer ||
155       T::category == TypeCategory::Real ||
156       T::category == TypeCategory::Complex);
157   using Element = typename Constant<T>::Element;
158   std::optional<int> dim;
159   if (std::optional<Constant<T>> array{
160           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
161               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
162     bool overflow{false};
163     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
164       if constexpr (T::category == TypeCategory::Integer) {
165         auto prod{element.MultiplySigned(array->At(at))};
166         overflow |= prod.SignedMultiplicationOverflowed();
167         element = prod.lower;
168       } else { // Real & Complex
169         auto prod{element.Multiply(array->At(at))};
170         overflow |= prod.flags.test(RealFlag::Overflow);
171         element = prod.value;
172       }
173     }};
174     if (overflow) {
175       context.messages().Say(
176           "PRODUCT() of %s data overflowed"_warn_en_US, T::AsFortran());
177     } else {
178       return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
179     }
180   }
181   return Expr<T>{std::move(ref)};
182 }
183 
184 // SUM
185 template <typename T>
FoldSum(FoldingContext & context,FunctionRef<T> && ref)186 static Expr<T> FoldSum(FoldingContext &context, FunctionRef<T> &&ref) {
187   static_assert(T::category == TypeCategory::Integer ||
188       T::category == TypeCategory::Real ||
189       T::category == TypeCategory::Complex);
190   using Element = typename Constant<T>::Element;
191   std::optional<int> dim;
192   Element identity{}, correction{};
193   if (std::optional<Constant<T>> array{
194           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
195               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
196     bool overflow{false};
197     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
198       if constexpr (T::category == TypeCategory::Integer) {
199         auto sum{element.AddSigned(array->At(at))};
200         overflow |= sum.overflow;
201         element = sum.value;
202       } else { // Real & Complex: use Kahan summation
203         auto next{array->At(at).Add(correction)};
204         overflow |= next.flags.test(RealFlag::Overflow);
205         auto sum{element.Add(next.value)};
206         overflow |= sum.flags.test(RealFlag::Overflow);
207         // correction = (sum - element) - next; algebraically zero
208         correction =
209             sum.value.Subtract(element).value.Subtract(next.value).value;
210         element = sum.value;
211       }
212     }};
213     if (overflow) {
214       context.messages().Say(
215           "SUM() of %s data overflowed"_warn_en_US, T::AsFortran());
216     } else {
217       return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
218     }
219   }
220   return Expr<T>{std::move(ref)};
221 }
222 
223 } // namespace Fortran::evaluate
224 #endif // FORTRAN_EVALUATE_FOLD_REDUCTION_H_
225