366 lines · cpp
1//===-- lib/runtime/findloc.cpp ---------------------------------*- C++ -*-===//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-exception6//7//===----------------------------------------------------------------------===//8 9// Implements FINDLOC for all required operand types and shapes and result10// integer kinds.11 12#include "flang-rt/runtime/reduction-templates.h"13#include "flang/Runtime/character.h"14#include "flang/Runtime/reduction.h"15#include <cinttypes>16#include <complex>17 18namespace Fortran::runtime {19 20template <TypeCategory CAT1, int KIND1, TypeCategory CAT2, int KIND2>21struct Equality {22 using Type1 = CppTypeFor<CAT1, KIND1>;23 using Type2 = CppTypeFor<CAT2, KIND2>;24 RT_API_ATTRS bool operator()(const Descriptor &array,25 const SubscriptValue at[], const Descriptor &target) const {26 if constexpr (KIND1 >= KIND2) {27 return *array.Element<Type1>(at) ==28 static_cast<Type1>(*target.OffsetElement<Type2>());29 } else {30 return static_cast<Type2>(*array.Element<Type1>(at)) ==31 *target.OffsetElement<Type2>();32 }33 }34};35 36template <int KIND1, int KIND2>37struct Equality<TypeCategory::Complex, KIND1, TypeCategory::Complex, KIND2> {38 using Type1 = CppTypeFor<TypeCategory::Complex, KIND1>;39 using Type2 = CppTypeFor<TypeCategory::Complex, KIND2>;40 RT_API_ATTRS bool operator()(const Descriptor &array,41 const SubscriptValue at[], const Descriptor &target) const {42 const Type1 &xz{*array.Element<Type1>(at)};43 const Type2 &tz{*target.OffsetElement<Type2>()};44 return xz.real() == tz.real() && xz.imag() == tz.imag();45 }46};47 48template <int KIND1, TypeCategory CAT2, int KIND2>49struct Equality<TypeCategory::Complex, KIND1, CAT2, KIND2> {50 using Type1 = CppTypeFor<TypeCategory::Complex, KIND1>;51 using Type2 = CppTypeFor<CAT2, KIND2>;52 RT_API_ATTRS bool operator()(const Descriptor &array,53 const SubscriptValue at[], const Descriptor &target) const {54 const Type1 &z{*array.Element<Type1>(at)};55 return z.imag() == 0 && z.real() == *target.OffsetElement<Type2>();56 }57};58 59template <TypeCategory CAT1, int KIND1, int KIND2>60struct Equality<CAT1, KIND1, TypeCategory::Complex, KIND2> {61 using Type1 = CppTypeFor<CAT1, KIND1>;62 using Type2 = CppTypeFor<TypeCategory::Complex, KIND2>;63 RT_API_ATTRS bool operator()(const Descriptor &array,64 const SubscriptValue at[], const Descriptor &target) const {65 const Type2 &z{*target.OffsetElement<Type2>()};66 return *array.Element<Type1>(at) == z.real() && z.imag() == 0;67 }68};69 70template <int KIND> struct CharacterEquality {71 using Type = CppTypeFor<TypeCategory::Character, KIND>;72 RT_API_ATTRS bool operator()(const Descriptor &array,73 const SubscriptValue at[], const Descriptor &target) const {74 return CharacterScalarCompare<Type>(array.Element<Type>(at),75 target.OffsetElement<Type>(),76 array.ElementBytes() / static_cast<unsigned>(KIND),77 target.ElementBytes() / static_cast<unsigned>(KIND)) == 0;78 }79};80 81struct LogicalEquivalence {82 RT_API_ATTRS bool operator()(const Descriptor &array,83 const SubscriptValue at[], const Descriptor &target) const {84 return IsLogicalElementTrue(array, at) ==85 IsLogicalElementTrue(target, at /*ignored*/);86 }87};88 89template <typename EQUALITY> class LocationAccumulator {90public:91 RT_API_ATTRS LocationAccumulator(92 const Descriptor &array, const Descriptor &target, bool back)93 : array_{array}, target_{target}, back_{back} {}94 RT_API_ATTRS void Reinitialize() { gotAnything_ = false; }95 template <typename A>96 RT_API_ATTRS void GetResult(A *p, int zeroBasedDim = -1) {97 if (zeroBasedDim >= 0) {98 *p = gotAnything_ ? location_[zeroBasedDim] -99 array_.GetDimension(zeroBasedDim).LowerBound() + 1100 : 0;101 } else if (gotAnything_) {102 for (int j{0}; j < rank_; ++j) {103 p[j] = location_[j] - array_.GetDimension(j).LowerBound() + 1;104 }105 } else {106 // no unmasked hits? result is all zeroes107 for (int j{0}; j < rank_; ++j) {108 p[j] = 0;109 }110 }111 }112 template <typename IGNORED>113 RT_API_ATTRS bool AccumulateAt(const SubscriptValue at[]) {114 if (equality_(array_, at, target_)) {115 gotAnything_ = true;116 for (int j{0}; j < rank_; ++j) {117 location_[j] = at[j];118 }119 return back_;120 } else {121 return true;122 }123 }124 125private:126 const Descriptor &array_;127 const Descriptor &target_;128 const bool back_{false};129 const int rank_{array_.rank()};130 bool gotAnything_{false};131 SubscriptValue location_[maxRank];132 const EQUALITY equality_{};133};134 135template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>136struct TotalNumericFindlocHelper {137 template <int TARGET_KIND> struct Functor {138 RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,139 const Descriptor &target, int kind, int dim, const Descriptor *mask,140 bool back, Terminator &terminator) const {141 using Eq = Equality<XCAT, XKIND, TARGET_CAT, TARGET_KIND>;142 using Accumulator = LocationAccumulator<Eq>;143 Accumulator accumulator{x, target, back};144 DoTotalReduction<void>(x, dim, mask, accumulator, "FINDLOC", terminator);145 ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor,146 void>(kind, terminator, accumulator, result);147 }148 };149};150 151template <TypeCategory CAT,152 template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>153 class HELPER>154struct NumericFindlocHelper {155 template <int KIND> struct Functor {156 // NVCC inlines more aggressively which causes too many specializations of157 // this function to be inlined causing compiler timeouts. Set as158 // noinline to allow compilation to complete.159 RT_API_ATTRS RT_DEVICE_NOINLINE void operator()(TypeCategory targetCat,160 int targetKind, Descriptor &result, const Descriptor &x,161 const Descriptor &target, int kind, int dim, const Descriptor *mask,162 bool back, Terminator &terminator) const {163 switch (targetCat) {164 case TypeCategory::Integer:165 case TypeCategory::Unsigned:166 ApplyIntegerKind<167 HELPER<CAT, KIND, TypeCategory::Integer>::template Functor, void>(168 targetKind, terminator, result, x, target, kind, dim, mask, back,169 terminator);170 break;171 case TypeCategory::Real:172 ApplyFloatingPointKind<173 HELPER<CAT, KIND, TypeCategory::Real>::template Functor, void>(174 targetKind, terminator, result, x, target, kind, dim, mask, back,175 terminator);176 break;177 case TypeCategory::Complex:178 ApplyFloatingPointKind<179 HELPER<CAT, KIND, TypeCategory::Complex>::template Functor, void>(180 targetKind, terminator, result, x, target, kind, dim, mask, back,181 terminator);182 break;183 default:184 terminator.Crash(185 "FINDLOC: bad target category %d for array category %d",186 static_cast<int>(targetCat), static_cast<int>(CAT));187 }188 }189 };190};191 192template <int KIND> struct CharacterFindlocHelper {193 RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,194 const Descriptor &target, int kind, const Descriptor *mask, bool back,195 Terminator &terminator) {196 using Accumulator = LocationAccumulator<CharacterEquality<KIND>>;197 Accumulator accumulator{x, target, back};198 DoTotalReduction<void>(x, 0, mask, accumulator, "FINDLOC", terminator);199 ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor, void>(200 kind, terminator, accumulator, result);201 }202};203 204static RT_API_ATTRS void LogicalFindlocHelper(Descriptor &result,205 const Descriptor &x, const Descriptor &target, int kind,206 const Descriptor *mask, bool back, Terminator &terminator) {207 using Accumulator = LocationAccumulator<LogicalEquivalence>;208 Accumulator accumulator{x, target, back};209 DoTotalReduction<void>(x, 0, mask, accumulator, "FINDLOC", terminator);210 ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor, void>(211 kind, terminator, accumulator, result);212}213 214extern "C" {215RT_EXT_API_GROUP_BEGIN216 217void RTDEF(Findloc)(Descriptor &result, const Descriptor &x,218 const Descriptor &target, int kind, const char *source, int line,219 const Descriptor *mask, bool back) {220 int rank{x.rank()};221 SubscriptValue extent[1]{rank};222 result.Establish(TypeCategory::Integer, kind, nullptr, 1, extent,223 CFI_attribute_allocatable);224 result.GetDimension(0).SetBounds(1, extent[0]);225 Terminator terminator{source, line};226 if (int stat{result.Allocate(kNoAsyncObject)}) {227 terminator.Crash(228 "FINDLOC: could not allocate memory for result; STAT=%d", stat);229 }230 CheckIntegerKind(terminator, kind, "FINDLOC");231 auto xType{x.type().GetCategoryAndKind()};232 auto targetType{target.type().GetCategoryAndKind()};233 RUNTIME_CHECK(terminator, xType.has_value() && targetType.has_value());234 switch (xType->first) {235 case TypeCategory::Integer:236 case TypeCategory::Unsigned:237 ApplyIntegerKind<NumericFindlocHelper<TypeCategory::Integer,238 TotalNumericFindlocHelper>::template Functor,239 void>(xType->second, terminator, targetType->first, targetType->second,240 result, x, target, kind, 0, mask, back, terminator);241 break;242 case TypeCategory::Real:243 ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Real,244 TotalNumericFindlocHelper>::template Functor,245 void>(xType->second, terminator, targetType->first, targetType->second,246 result, x, target, kind, 0, mask, back, terminator);247 break;248 case TypeCategory::Complex:249 ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Complex,250 TotalNumericFindlocHelper>::template Functor,251 void>(xType->second, terminator, targetType->first, targetType->second,252 result, x, target, kind, 0, mask, back, terminator);253 break;254 case TypeCategory::Character:255 RUNTIME_CHECK(terminator,256 targetType->first == TypeCategory::Character &&257 targetType->second == xType->second);258 ApplyCharacterKind<CharacterFindlocHelper, void>(xType->second, terminator,259 result, x, target, kind, mask, back, terminator);260 break;261 case TypeCategory::Logical:262 RUNTIME_CHECK(terminator, targetType->first == TypeCategory::Logical);263 LogicalFindlocHelper(result, x, target, kind, mask, back, terminator);264 break;265 default:266 terminator.Crash(267 "FINDLOC: bad data type code (%d) for array", x.type().raw());268 }269}270 271RT_EXT_API_GROUP_END272} // extern "C"273 274// FINDLOC with DIM=275 276template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>277struct PartialNumericFindlocHelper {278 template <int TARGET_KIND> struct Functor {279 RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,280 const Descriptor &target, int kind, int dim, const Descriptor *mask,281 bool back, Terminator &terminator) const {282 using Eq = Equality<XCAT, XKIND, TARGET_CAT, TARGET_KIND>;283 using Accumulator = LocationAccumulator<Eq>;284 Accumulator accumulator{x, target, back};285 ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor,286 void>(kind, terminator, result, x, dim, mask, terminator, "FINDLOC",287 accumulator);288 }289 };290};291 292template <int KIND> struct PartialCharacterFindlocHelper {293 RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,294 const Descriptor &target, int kind, int dim, const Descriptor *mask,295 bool back, Terminator &terminator) {296 using Accumulator = LocationAccumulator<CharacterEquality<KIND>>;297 Accumulator accumulator{x, target, back};298 ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor,299 void>(kind, terminator, result, x, dim, mask, terminator, "FINDLOC",300 accumulator);301 }302};303 304static RT_API_ATTRS void PartialLogicalFindlocHelper(Descriptor &result,305 const Descriptor &x, const Descriptor &target, int kind, int dim,306 const Descriptor *mask, bool back, Terminator &terminator) {307 using Accumulator = LocationAccumulator<LogicalEquivalence>;308 Accumulator accumulator{x, target, back};309 ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor, void>(310 kind, terminator, result, x, dim, mask, terminator, "FINDLOC",311 accumulator);312}313 314extern "C" {315RT_EXT_API_GROUP_BEGIN316 317void RTDEF(FindlocDim)(Descriptor &result, const Descriptor &x,318 const Descriptor &target, int kind, int dim, const char *source, int line,319 const Descriptor *mask, bool back) {320 Terminator terminator{source, line};321 CheckIntegerKind(terminator, kind, "FINDLOC");322 auto xType{x.type().GetCategoryAndKind()};323 auto targetType{target.type().GetCategoryAndKind()};324 RUNTIME_CHECK(terminator, xType.has_value() && targetType.has_value());325 switch (xType->first) {326 case TypeCategory::Integer:327 case TypeCategory::Unsigned:328 ApplyIntegerKind<NumericFindlocHelper<TypeCategory::Integer,329 PartialNumericFindlocHelper>::template Functor,330 void>(xType->second, terminator, targetType->first, targetType->second,331 result, x, target, kind, dim, mask, back, terminator);332 break;333 case TypeCategory::Real:334 ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Real,335 PartialNumericFindlocHelper>::template Functor,336 void>(xType->second, terminator, targetType->first, targetType->second,337 result, x, target, kind, dim, mask, back, terminator);338 break;339 case TypeCategory::Complex:340 ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Complex,341 PartialNumericFindlocHelper>::template Functor,342 void>(xType->second, terminator, targetType->first, targetType->second,343 result, x, target, kind, dim, mask, back, terminator);344 break;345 case TypeCategory::Character:346 RUNTIME_CHECK(terminator,347 targetType->first == TypeCategory::Character &&348 targetType->second == xType->second);349 ApplyCharacterKind<PartialCharacterFindlocHelper, void>(xType->second,350 terminator, result, x, target, kind, dim, mask, back, terminator);351 break;352 case TypeCategory::Logical:353 RUNTIME_CHECK(terminator, targetType->first == TypeCategory::Logical);354 PartialLogicalFindlocHelper(355 result, x, target, kind, dim, mask, back, terminator);356 break;357 default:358 terminator.Crash(359 "FINDLOC: bad data type code (%d) for array", x.type().raw());360 }361}362 363RT_EXT_API_GROUP_END364} // extern "C"365} // namespace Fortran::runtime366