[flang] Fold CSHIFT

[klausler]

Summary

Implement folding of the transformational intrinsic function
CSHIFT for all types.

[PeteSteinfeld]

I had one comment in fold-implementation.h. Otherwise, all builds, tests, and looks good.

Thanks for doing this!

flang/lib/Evaluate/fold-implementation.h

557–559

Can you move this code down below after you get the "shift" variable below and avoid the "if" statement at line 563?

[klausler]

flang/lib/Evaluate/fold-implementation.h

557–559

If shiftExpr isn't available, I can't try to convert it.

flang/include/flang/Evaluate/tools.h

return std::visit([&](auto &&x) { return Expr<T>{Expand(std::move(x))}; },
         std::move(x.u));
   }
 
 private:
   ConstantSubscripts extents_;
   std::optional<ConstantSubscripts> lbounds_;
 };
 
+// Given a collection of element values, package them as a Constant.
+// If the type is Character or a derived type, take the length or type
+// (resp.) from a another Constant.
+template <typename T>
+Constant<T> PackageConstant(std::vector<Scalar<T>> &&elements,
+    const Constant<T> &reference, const ConstantSubscripts &shape) {
+  if constexpr (T::category == TypeCategory::Character) {
+    return Constant<T>{
+        reference.LEN(), std::move(elements), ConstantSubscripts{shape}};
+  } else if constexpr (T::category == TypeCategory::Derived) {
+    return Constant<T>{reference.GetType().GetDerivedTypeSpec(),
+        std::move(elements), ConstantSubscripts{shape}};
+  } else {
+    return Constant<T>{std::move(elements), ConstantSubscripts{shape}};
+  }
+}
+
 } // namespace Fortran::evaluate
 
 namespace Fortran::semantics {
 
 class Scope;
 
 // These functions are used in Evaluate so they are defined here rather than in
 // Semantics to avoid a link-time dependency on Semantics.

flang/lib/Evaluate/fold-character.cpp

if (name == "achar" || name == "char") {
     }
   } else if (name == "trim") { // not elemental
     if (auto scalar{
             GetScalarConstantArguments<T>(context, funcRef.arguments())}) {
       return Expr<T>{Constant<T>{
           CharacterUtils<KIND>::TRIM(std::get<Scalar<T>>(*scalar))}};
     }
   }
-  // TODO: cshift, eoshift, maxloc, minloc, pack, spread, transfer,
-  // transpose, unpack
+  // TODO: findloc, maxloc, minloc, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Character, KIND>> FoldOperation(
     FoldingContext &context, Concat<KIND> &&x) {
   if (auto array{ApplyElementwise(context, x)}) {
     return *array;

flang/lib/Evaluate/fold-complex.cpp

Expr<Type<TypeCategory::Complex, KIND>> FoldIntrinsicFunction(
   } else if (name == "merge") {
     return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "product") {
     auto one{Scalar<Part>::FromInteger(value::Integer<8>{1}).value};
     return FoldProduct<T>(context, std::move(funcRef), Scalar<T>{one});
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   }
-  // TODO: cshift, dot_product, eoshift, matmul, pack, spread, transfer,
-  // transpose, unpack
+  // TODO: dot_product, matmul, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Complex, KIND>> FoldOperation(
     FoldingContext &context, ComplexConstructor<KIND> &&x) {
   if (auto array{ApplyElementwise(context, x)}) {
     return *array;

flang/lib/Evaluate/fold-implementation.h

public:
   std::optional<Constant<T>> ApplyComponent(Constant<SomeDerived> &&,
       const Symbol &component,
       const std::vector<Constant<SubscriptInteger>> * = nullptr);
   std::optional<Constant<T>> GetConstantComponent(
       Component &, const std::vector<Constant<SubscriptInteger>> * = nullptr);
   std::optional<Constant<T>> Folding(ArrayRef &);
   Expr<T> Folding(Designator<T> &&);
   Constant<T> *Folding(std::optional<ActualArgument> &);
-  Expr<T> Reshape(FunctionRef<T> &&);
+
+  Expr<T> CSHIFT(FunctionRef<T> &&);
+  Expr<T> RESHAPE(FunctionRef<T> &&);
 
 private:
   FoldingContext &context_;
 };
 
 std::optional<Constant<SubscriptInteger>> GetConstantSubscript(
     FoldingContext &, Subscript &, const NamedEntity &, int dim);
 

 // gets re-folded.
 template <typename T> Expr<T> MakeInvalidIntrinsic(FunctionRef<T> &&funcRef) {
   SpecificIntrinsic invalid{std::get<SpecificIntrinsic>(funcRef.proc().u)};
   invalid.name = IntrinsicProcTable::InvalidName;
   return Expr<T>{FunctionRef<T>{ProcedureDesignator{std::move(invalid)},
       ActualArguments{std::move(funcRef.arguments())}}};
 }
 
-template <typename T> Expr<T> Folder<T>::Reshape(FunctionRef<T> &&funcRef) {
+template <typename T> Expr<T> Folder<T>::CSHIFT(FunctionRef<T> &&funcRef) {
+  auto args{funcRef.arguments()};
+  CHECK(args.size() == 3);
+  const auto *array{UnwrapConstantValue<T>(args[0])};
+  const auto *shiftExpr{UnwrapExpr<Expr<SomeInteger>>(args[1])};
+  auto dim{GetInt64ArgOr(args[2], 1)};
+  if (!array || !shiftExpr || !dim) {
+    return Expr<T>{std::move(funcRef)};
+  }
+  auto convertedShift{Fold(context_,
+      ConvertToType<SubscriptInteger>(Expr<SomeInteger>{*shiftExpr}))};
+  const auto *shift{UnwrapConstantValue<SubscriptInteger>(convertedShift)};
+  if (!shift) {
+    return Expr<T>{std::move(funcRef)};
+  }
+  // Arguments are constant
+  if (*dim < 1 || *dim > array->Rank()) {
+    context_.messages().Say("Invalid 'dim=' argument (%jd) in CSHIFT"_err_en_US,
+        static_cast<std::intmax_t>(*dim));
+  } else if (shift->Rank() > 0 && shift->Rank() != array->Rank() - 1) {
+    // message already emitted from intrinsic look-up
+  } else {
+    int rank{array->Rank()};
+    int zbDim{static_cast<int>(*dim) - 1};
+    bool ok{true};
+    if (shift->Rank() > 0) {
+      int k{0};
+      for (int j{0}; j < rank; ++j) {
+        if (j != zbDim) {
+          if (array->shape()[j] != shift->shape()[k]) {
+            context_.messages().Say(
+                "Invalid 'shift=' argument in CSHIFT; extent on dimension %d is %jd but must be %jd"_err_en_US,
+                k + 1, static_cast<std::intmax_t>(shift->shape()[k]),
+                static_cast<std::intmax_t>(array->shape()[j]));
+            ok = false;
+          }
+          ++k;
+        }
+      }
+    }
+    if (ok) {
+      std::vector<Scalar<T>> resultElements;
+      ConstantSubscripts arrayAt{array->lbounds()};
+      ConstantSubscript dimLB{arrayAt[zbDim]};
+      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;
+          }
+        }
+        arrayAt[zbDim] = dimLB + dimExtent - 1;
+        array->IncrementSubscripts(arrayAt);
+        shift->IncrementSubscripts(shiftAt);
+      }
+      return Expr<T>{PackageConstant<T>(
+          std::move(resultElements), *array, array->shape())};
+    }
+  }
+  // Invalid, prevent re-folding
+  return MakeInvalidIntrinsic(std::move(funcRef));
+}
+
+template <typename T> Expr<T> Folder<T>::RESHAPE(FunctionRef<T> &&funcRef) {
   auto args{funcRef.arguments()};
   CHECK(args.size() == 4);
   const auto *source{UnwrapConstantValue<T>(args[0])};
   const auto *pad{UnwrapConstantValue<T>(args[2])};
   std::optional<std::vector<ConstantSubscript>> shape{
       GetIntegerVector<ConstantSubscript>(args[1])};
   std::optional<std::vector<int>> order{GetIntegerVector<int>(args[3])};
   if (!source || !shape || (args[2] && !pad) || (args[3] && !order)) {
Expr<T> FoldOperation(FoldingContext &context, FunctionRef<T> &&funcRef) {
   ActualArguments &args{funcRef.arguments()};
   for (std::optional<ActualArgument> &arg : args) {
     if (auto *expr{UnwrapExpr<Expr<SomeType>>(arg)}) {
       *expr = Fold(context, std::move(*expr));
     }
   }
   if (auto *intrinsic{std::get_if<SpecificIntrinsic>(&funcRef.proc().u)}) {
     const std::string name{intrinsic->name};
-    if (name == "reshape") {
-      return Folder<T>{context}.Reshape(std::move(funcRef));
+    if (name == "cshift") {
+      return Folder<T>{context}.CSHIFT(std::move(funcRef));
+    } else if (name == "reshape") {
+      return Folder<T>{context}.RESHAPE(std::move(funcRef));
     }
-    // TODO: other type independent transformationals
+    // TODO: eoshift, pack, spread, unpack, transpose
+    // TODO: extends_type_of, same_type_as
     if constexpr (!std::is_same_v<T, SomeDerived>) {
       return FoldIntrinsicFunction(context, std::move(funcRef));
     }
   }
   return Expr<T>{std::move(funcRef)};
 }
 
 template <typename T>

flang/lib/Evaluate/fold-integer.cpp

if (auto info{
             Fold(context, Expr<T>{8} * ConvertToType<T>(std::move(*bytes)))};
       }
     }
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   } else if (name == "ubound") {
     return UBOUND(context, std::move(funcRef));
   }
-  // TODO:
-  // cshift, dot_product, eoshift, findloc, ibits, image_status, ishftc,
-  // matmul, maxloc, minloc, not, pack, sign, spread, transfer, transpose,
-  // unpack
+  // TODO: count(w/ dim), dot_product, findloc, ibits, image_status, ishftc,
+  // matmul, maxloc, minloc, sign, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 // Substitutes a bare type parameter reference with its value if it has one now
 // in an instantiation.  Bare LEN type parameters are substituted only when
 // the known value is constant.
 Expr<TypeParamInquiry::Result> FoldOperation(
     FoldingContext &context, TypeParamInquiry &&inquiry) {

flang/lib/Evaluate/fold-logical.cpp

return FoldMerge<T>(context, std::move(funcRef));
       name == "__builtin_ieee_support_io" ||
       name == "__builtin_ieee_support_nan" ||
       name == "__builtin_ieee_support_sqrt" ||
       name == "__builtin_ieee_support_standard" ||
       name == "__builtin_ieee_support_subnormal" ||
       name == "__builtin_ieee_support_underflow_control") {
     return Expr<T>{true};
   }
-  // TODO: btest, cshift, dot_product, eoshift, is_iostat_end,
+  // TODO: btest, dot_product, eoshift, is_iostat_end,
   // is_iostat_eor, lge, lgt, lle, llt, logical, matmul, out_of_range,
-  // pack, parity, spread, transfer, transpose, unpack, extends_type_of,
-  // same_type_as
+  // parity, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 template <typename T>
 Expr<LogicalResult> FoldOperation(
     FoldingContext &context, Relational<T> &&relation) {
   if (auto array{ApplyElementwise(context, relation,
           std::function<Expr<LogicalResult>(Expr<T> &&, Expr<T> &&)>{

flang/lib/Evaluate/fold-real.cpp

Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
   } else if (name == "sign") {
     return FoldElementalIntrinsic<T, T, T>(
         context, std::move(funcRef), &Scalar<T>::SIGN);
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   } else if (name == "tiny") {
     return Expr<T>{Scalar<T>::TINY()};
   }
-  // TODO: cshift, dim, dot_product, eoshift, fraction, matmul,
-  // maxloc, minloc, modulo, nearest, norm2, pack, rrspacing, scale,
-  // set_exponent, spacing, spread, transfer, transpose, unpack,
+  // TODO: dim, dot_product, fraction, matmul,
+  // maxloc, minloc, modulo, nearest, norm2, rrspacing, scale,
+  // set_exponent, spacing, transfer,
   // bessel_jn (transformational) and bessel_yn (transformational)
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Real, KIND>> FoldOperation(
     FoldingContext &context, ComplexComponent<KIND> &&x) {
   using Operand = Type<TypeCategory::Complex, KIND>;

flang/test/Evaluate/folding22.f90

 logical, parameter :: test_array_array = all(array_array.eq.(['1a', '2b']))
 
 
 character(1), parameter :: input(2) = ['x', 'y']
 character(*), parameter :: zero_sized(*) = input(2:1:1) // 'abcde'
 logical, parameter :: test_zero_sized = len(zero_sized).eq.6
 
 end
- 

flang/test/Evaluate/folding27.f90

+! RUN: %S/test_folding.sh %s %t %flang_fc1
+! REQUIRES: shell
+! Tests folding of CSHIFT (valid cases)
+module m
+  integer, parameter :: arr(2,3) = reshape([1, 2, 3, 4, 5, 6], shape(arr))
+  logical, parameter :: test_sanity = all([arr] == [1, 2, 3, 4, 5, 6])
+  logical, parameter :: test_cshift_0 = all(cshift([1, 2, 3], 0) == [1, 2, 3])
+  logical, parameter :: test_cshift_1 = all(cshift([1, 2, 3], 1) == [2, 3, 1])
+  logical, parameter :: test_cshift_2 = all(cshift([1, 2, 3], 3) == [1, 2, 3])
+  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_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])
+end module