[flang][rt] Enable Count and CountDim for device build (#141684)

Author: clementval

flang-rt/lib/runtime/CMakeLists.txt

@@ -56,6 +56,7 @@ set(supported_sources
   product.cpp
   pseudo-unit.cpp
   ragged.cpp
+  reduction.cpp
   stat.cpp
   stop.cpp
   sum.cpp

flang-rt/lib/runtime/reduction.cpp

@@ -223,10 +223,13 @@ enum class LogicalReduction { All, Any, Parity };
 template <LogicalReduction REDUCTION> class LogicalAccumulator {
 public:
   using Type = bool;
-  explicit LogicalAccumulator(const Descriptor &array) : array_{array} {}
-  void Reinitialize() { result_ = REDUCTION == LogicalReduction::All; }
-  bool Result() const { return result_; }
-  bool Accumulate(bool x) {
+  RT_API_ATTRS explicit LogicalAccumulator(const Descriptor &array)
+      : array_{array} {}
+  RT_API_ATTRS void Reinitialize() {
+    result_ = REDUCTION == LogicalReduction::All;
+  }
+  RT_API_ATTRS bool Result() const { return result_; }
+  RT_API_ATTRS bool Accumulate(bool x) {
     if constexpr (REDUCTION == LogicalReduction::Parity) {
       result_ = result_ != x;
     } else if (x != (REDUCTION == LogicalReduction::All)) {
@@ -236,7 +239,7 @@ template <LogicalReduction REDUCTION> class LogicalAccumulator {
     return true;
   }
   template <typename IGNORED = void>
-  bool AccumulateAt(const SubscriptValue at[]) {
+  RT_API_ATTRS bool AccumulateAt(const SubscriptValue at[]) {
     return Accumulate(IsLogicalElementTrue(array_, at));
   }
 
@@ -246,9 +249,9 @@ template <LogicalReduction REDUCTION> class LogicalAccumulator {
 };
 
 template <typename ACCUMULATOR>
-inline auto GetTotalLogicalReduction(const Descriptor &x, const char *source,
-    int line, int dim, ACCUMULATOR &&accumulator, const char *intrinsic) ->
-    typename ACCUMULATOR::Type {
+RT_API_ATTRS inline auto GetTotalLogicalReduction(const Descriptor &x,
+    const char *source, int line, int dim, ACCUMULATOR &&accumulator,
+    const char *intrinsic) -> typename ACCUMULATOR::Type {
   Terminator terminator{source, line};
   if (dim < 0 || dim > 1) {
     terminator.Crash("%s: bad DIM=%d for ARRAY with rank=1", intrinsic, dim);
@@ -264,8 +267,9 @@ inline auto GetTotalLogicalReduction(const Descriptor &x, const char *source,
 }
 
 template <typename ACCUMULATOR>
-inline auto ReduceLogicalDimToScalar(const Descriptor &x, int zeroBasedDim,
-    SubscriptValue subscripts[]) -> typename ACCUMULATOR::Type {
+RT_API_ATTRS inline auto ReduceLogicalDimToScalar(
+    const Descriptor &x, int zeroBasedDim, SubscriptValue subscripts[]) ->
+    typename ACCUMULATOR::Type {
   ACCUMULATOR accumulator{x};
   SubscriptValue xAt[maxRank];
   GetExpandedSubscripts(xAt, x, zeroBasedDim, subscripts);
@@ -282,8 +286,8 @@ inline auto ReduceLogicalDimToScalar(const Descriptor &x, int zeroBasedDim,
 
 template <LogicalReduction REDUCTION> struct LogicalReduceHelper {
   template <int KIND> struct Functor {
-    void operator()(Descriptor &result, const Descriptor &x, int dim,
-        Terminator &terminator, const char *intrinsic) const {
+    RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,
+        int dim, Terminator &terminator, const char *intrinsic) const {
       // Standard requires result to have same LOGICAL kind as argument.
       CreatePartialReductionResult(
           result, x, x.ElementBytes(), dim, terminator, intrinsic, x.type());
@@ -301,8 +305,9 @@ template <LogicalReduction REDUCTION> struct LogicalReduceHelper {
 };
 
 template <LogicalReduction REDUCTION>
-inline void DoReduceLogicalDimension(Descriptor &result, const Descriptor &x,
-    int dim, Terminator &terminator, const char *intrinsic) {
+RT_API_ATTRS inline void DoReduceLogicalDimension(Descriptor &result,
+    const Descriptor &x, int dim, Terminator &terminator,
+    const char *intrinsic) {
   auto catKind{x.type().GetCategoryAndKind()};
   RUNTIME_CHECK(terminator, catKind && catKind->first == TypeCategory::Logical);
   ApplyLogicalKind<LogicalReduceHelper<REDUCTION>::template Functor, void>(
@@ -314,11 +319,12 @@ inline void DoReduceLogicalDimension(Descriptor &result, const Descriptor &x,
 class CountAccumulator {
 public:
   using Type = std::int64_t;
-  explicit CountAccumulator(const Descriptor &array) : array_{array} {}
-  void Reinitialize() { result_ = 0; }
-  Type Result() const { return result_; }
+  RT_API_ATTRS explicit CountAccumulator(const Descriptor &array)
+      : array_{array} {}
+  RT_API_ATTRS void Reinitialize() { result_ = 0; }
+  RT_API_ATTRS Type Result() const { return result_; }
   template <typename IGNORED = void>
-  bool AccumulateAt(const SubscriptValue at[]) {
+  RT_API_ATTRS bool AccumulateAt(const SubscriptValue at[]) {
     if (IsLogicalElementTrue(array_, at)) {
       ++result_;
     }
@@ -331,7 +337,7 @@ class CountAccumulator {
 };
 
 template <int KIND> struct CountDimension {
-  void operator()(Descriptor &result, const Descriptor &x, int dim,
+  RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x, int dim,
       Terminator &terminator) const {
     // Element size of the descriptor descriptor is the size
     // of {TypeCategory::Integer, KIND}.