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

This commit is contained in:
Valentin Clement (バレンタイン クレメン)
2025-05-28 09:55:49 -07:00
committed by GitHub
parent a41e20a53c
commit fc9ce037ef
2 changed files with 26 additions and 19 deletions

View File

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

View File

@@ -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 @@ public:
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 @@ private:
};
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 @@ private:
};
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}.