[flang-commits] [flang] [flang] Check for overflows in RESHAPE folding (PR #68342)

Leandro Lupori via flang-commits flang-commits at lists.llvm.org
Tue Oct 17 12:32:35 PDT 2023


https://github.com/luporl updated https://github.com/llvm/llvm-project/pull/68342

>From daa8de4653bc9e4fd7f8b5827b3b0e0cc3be9213 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Wed, 4 Oct 2023 18:15:38 +0000
Subject: [PATCH 1/2] [flang] Check for overflows in RESHAPE folding

---
 flang/include/flang/Evaluate/constant.h  |  1 +
 flang/lib/Evaluate/constant.cpp          | 14 ++++++++++++++
 flang/lib/Evaluate/fold-implementation.h |  3 +++
 flang/test/Semantics/reshape.f90         |  5 +++++
 4 files changed, 23 insertions(+)

diff --git a/flang/include/flang/Evaluate/constant.h b/flang/include/flang/Evaluate/constant.h
index 04474e2f49a0f88..c5e08d38f793234 100644
--- a/flang/include/flang/Evaluate/constant.h
+++ b/flang/include/flang/Evaluate/constant.h
@@ -47,6 +47,7 @@ inline int GetRank(const ConstantSubscripts &s) {
 }
 
 std::size_t TotalElementCount(const ConstantSubscripts &);
+bool TotalElementCountOverflows(const ConstantSubscripts &);
 
 // Validate dimension re-ordering like ORDER in RESHAPE.
 // On success, return a vector that can be used as dimOrder in
diff --git a/flang/lib/Evaluate/constant.cpp b/flang/lib/Evaluate/constant.cpp
index 084836b4ec36773..76403364ce503e3 100644
--- a/flang/lib/Evaluate/constant.cpp
+++ b/flang/lib/Evaluate/constant.cpp
@@ -84,6 +84,20 @@ std::size_t TotalElementCount(const ConstantSubscripts &shape) {
   return static_cast<std::size_t>(GetSize(shape));
 }
 
+bool TotalElementCountOverflows(const ConstantSubscripts &shape) {
+  uint64_t size{1};
+  for (auto dim : shape) {
+    CHECK(dim >= 0);
+    uint64_t osize{size};
+    size = osize * dim;
+    if (size > std::numeric_limits<decltype(dim)>::max() ||
+        (dim != 0 && size / dim != osize)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 bool ConstantBounds::IncrementSubscripts(
     ConstantSubscripts &indices, const std::vector<int> *dimOrder) const {
   int rank{GetRank(shape_)};
diff --git a/flang/lib/Evaluate/fold-implementation.h b/flang/lib/Evaluate/fold-implementation.h
index 2a40018cd5a3865..0748a27268aee26 100644
--- a/flang/lib/Evaluate/fold-implementation.h
+++ b/flang/lib/Evaluate/fold-implementation.h
@@ -878,6 +878,9 @@ template <typename T> Expr<T> Folder<T>::RESHAPE(FunctionRef<T> &&funcRef) {
   } else if (HasNegativeExtent(shape.value())) {
     context_.messages().Say(
         "'shape=' argument must not have a negative extent"_err_en_US);
+  } else if (TotalElementCountOverflows(shape.value())) {
+    context_.messages().Say(
+        "'shape=' argument has too many elements"_err_en_US);
   } else {
     int rank{GetRank(shape.value())};
     std::size_t resultElements{TotalElementCount(shape.value())};
diff --git a/flang/test/Semantics/reshape.f90 b/flang/test/Semantics/reshape.f90
index 2e9b5adf3ff0e50..4f9be6510e5dcd4 100644
--- a/flang/test/Semantics/reshape.f90
+++ b/flang/test/Semantics/reshape.f90
@@ -44,6 +44,11 @@ program reshaper
   type(dType), parameter :: array19(*) = [dType::dType(field=[1,2])]
   logical, parameter :: lVar = all(array19(:)%field(1) == [2])
 
+  integer(8), parameter :: I64_MAX = INT(z'7fffffffffffffff', kind=8)
+  integer(8), parameter :: huge_shape(2) = [I64_MAX, I64_MAX]
+  !ERROR: 'shape=' argument has too many elements
+  integer :: array21(I64_MAX, I64_MAX) = RESHAPE([1, 2, 3], huge_shape)
+
   !ERROR: Size of 'shape=' argument must not be greater than 15
   CALL ext_sub(RESHAPE([(n, n=1,20)], &
     [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]))

>From c8449bd6978492692c3a999580d122851972a409 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Tue, 17 Oct 2023 19:26:53 +0000
Subject: [PATCH 2/2] Make TotalElementCount() return std::optional<uint64_t>

And change all of its callers, checking for overflows.
---
 flang/include/flang/Evaluate/constant.h      |  4 +-
 flang/include/flang/Evaluate/initial-image.h |  8 +-
 flang/lib/Evaluate/constant.cpp              | 27 +++---
 flang/lib/Evaluate/fold-designator.cpp       |  4 +-
 flang/lib/Evaluate/fold-implementation.h     | 87 ++++++++++++--------
 flang/lib/Evaluate/initial-image.cpp         | 10 ++-
 flang/lib/Semantics/data-to-inits.cpp        |  2 +
 7 files changed, 87 insertions(+), 55 deletions(-)

diff --git a/flang/include/flang/Evaluate/constant.h b/flang/include/flang/Evaluate/constant.h
index c5e08d38f793234..8c841918bccbe89 100644
--- a/flang/include/flang/Evaluate/constant.h
+++ b/flang/include/flang/Evaluate/constant.h
@@ -46,8 +46,8 @@ inline int GetRank(const ConstantSubscripts &s) {
   return static_cast<int>(s.size());
 }
 
-std::size_t TotalElementCount(const ConstantSubscripts &);
-bool TotalElementCountOverflows(const ConstantSubscripts &);
+// Returns the number of elements of shape, if no overflow occurs.
+std::optional<uint64_t> TotalElementCount(const ConstantSubscripts &shape);
 
 // Validate dimension re-ordering like ORDER in RESHAPE.
 // On success, return a vector that can be used as dimOrder in
diff --git a/flang/include/flang/Evaluate/initial-image.h b/flang/include/flang/Evaluate/initial-image.h
index 7c8f986f853a34b..753da15008a8985 100644
--- a/flang/include/flang/Evaluate/initial-image.h
+++ b/flang/include/flang/Evaluate/initial-image.h
@@ -22,7 +22,7 @@ namespace Fortran::evaluate {
 
 class InitialImage {
 public:
-  enum Result { Ok, NotAConstant, OutOfRange, SizeMismatch };
+  enum Result { Ok, NotAConstant, OutOfRange, SizeMismatch, TooManyElems };
 
   explicit InitialImage(std::size_t bytes) : data_(bytes) {}
   InitialImage(InitialImage &&that) = default;
@@ -60,7 +60,11 @@ class InitialImage {
     if (offset < 0 || offset + bytes > data_.size()) {
       return OutOfRange;
     } else {
-      auto elements{TotalElementCount(x.shape())};
+      auto optElements{TotalElementCount(x.shape())};
+      if (!optElements) {
+        return TooManyElems;
+      }
+      auto elements{*optElements};
       auto elementBytes{bytes > 0 ? bytes / elements : 0};
       if (elements * elementBytes != bytes) {
         return SizeMismatch;
diff --git a/flang/lib/Evaluate/constant.cpp b/flang/lib/Evaluate/constant.cpp
index 76403364ce503e3..f5041c1ae75feb1 100644
--- a/flang/lib/Evaluate/constant.cpp
+++ b/flang/lib/Evaluate/constant.cpp
@@ -80,11 +80,7 @@ ConstantSubscript ConstantBounds::SubscriptsToOffset(
   return offset;
 }
 
-std::size_t TotalElementCount(const ConstantSubscripts &shape) {
-  return static_cast<std::size_t>(GetSize(shape));
-}
-
-bool TotalElementCountOverflows(const ConstantSubscripts &shape) {
+std::optional<uint64_t> TotalElementCount(const ConstantSubscripts &shape) {
   uint64_t size{1};
   for (auto dim : shape) {
     CHECK(dim >= 0);
@@ -92,10 +88,10 @@ bool TotalElementCountOverflows(const ConstantSubscripts &shape) {
     size = osize * dim;
     if (size > std::numeric_limits<decltype(dim)>::max() ||
         (dim != 0 && size / dim != osize)) {
-      return true;
+      return std::nullopt;
     }
   }
-  return false;
+  return static_cast<uint64_t>(GetSize(shape));
 }
 
 bool ConstantBounds::IncrementSubscripts(
@@ -149,7 +145,7 @@ template <typename RESULT, typename ELEMENT>
 ConstantBase<RESULT, ELEMENT>::ConstantBase(
     std::vector<Element> &&x, ConstantSubscripts &&sh, Result res)
     : ConstantBounds(std::move(sh)), result_{res}, values_(std::move(x)) {
-  CHECK(size() == TotalElementCount(shape()));
+  CHECK(TotalElementCount(shape()) && size() == *TotalElementCount(shape()));
 }
 
 template <typename RESULT, typename ELEMENT>
@@ -163,7 +159,9 @@ bool ConstantBase<RESULT, ELEMENT>::operator==(const ConstantBase &that) const {
 template <typename RESULT, typename ELEMENT>
 auto ConstantBase<RESULT, ELEMENT>::Reshape(
     const ConstantSubscripts &dims) const -> std::vector<Element> {
-  std::size_t n{TotalElementCount(dims)};
+  std::optional<uint64_t> optN{TotalElementCount(dims)};
+  CHECK(optN);
+  uint64_t n{*optN};
   CHECK(!empty() || n == 0);
   std::vector<Element> elements;
   auto iter{values().cbegin()};
@@ -223,7 +221,8 @@ template <int KIND>
 Constant<Type<TypeCategory::Character, KIND>>::Constant(ConstantSubscript len,
     std::vector<Scalar<Result>> &&strings, ConstantSubscripts &&sh)
     : ConstantBounds(std::move(sh)), length_{len} {
-  CHECK(strings.size() == TotalElementCount(shape()));
+  CHECK(TotalElementCount(shape()) &&
+      strings.size() == *TotalElementCount(shape()));
   values_.assign(strings.size() * length_,
       static_cast<typename Scalar<Result>::value_type>(' '));
   ConstantSubscript at{0};
@@ -250,7 +249,9 @@ bool Constant<Type<TypeCategory::Character, KIND>>::empty() const {
 template <int KIND>
 std::size_t Constant<Type<TypeCategory::Character, KIND>>::size() const {
   if (length_ == 0) {
-    return TotalElementCount(shape());
+    std::optional<uint64_t> n{TotalElementCount(shape())};
+    CHECK(n);
+    return *n;
   } else {
     return static_cast<ConstantSubscript>(values_.size()) / length_;
   }
@@ -288,7 +289,9 @@ auto Constant<Type<TypeCategory::Character, KIND>>::Substring(
 template <int KIND>
 auto Constant<Type<TypeCategory::Character, KIND>>::Reshape(
     ConstantSubscripts &&dims) const -> Constant<Result> {
-  std::size_t n{TotalElementCount(dims)};
+  std::optional<uint64_t> optN{TotalElementCount(dims)};
+  CHECK(optN);
+  uint64_t n{*optN};
   CHECK(!empty() || n == 0);
   std::vector<Element> elements;
   ConstantSubscript at{0},
diff --git a/flang/lib/Evaluate/fold-designator.cpp b/flang/lib/Evaluate/fold-designator.cpp
index 7298b0a2fb10c56..6952436681f753c 100644
--- a/flang/lib/Evaluate/fold-designator.cpp
+++ b/flang/lib/Evaluate/fold-designator.cpp
@@ -373,7 +373,9 @@ ConstantObjectPointer ConstantObjectPointer::From(
     FoldingContext &context, const Expr<SomeType> &expr) {
   auto extents{GetConstantExtents(context, expr)};
   CHECK(extents);
-  std::size_t elements{TotalElementCount(*extents)};
+  std::optional<uint64_t> optElements{TotalElementCount(*extents)};
+  CHECK(optElements);
+  uint64_t elements{*optElements};
   CHECK(elements > 0);
   int rank{GetRank(*extents)};
   ConstantSubscripts at(rank, 1);
diff --git a/flang/lib/Evaluate/fold-implementation.h b/flang/lib/Evaluate/fold-implementation.h
index 0748a27268aee26..868b7b6990fd384 100644
--- a/flang/lib/Evaluate/fold-implementation.h
+++ b/flang/lib/Evaluate/fold-implementation.h
@@ -492,7 +492,13 @@ Expr<TR> FoldElementalIntrinsicHelper(FoldingContext &context,
     CHECK(rank == GetRank(shape));
     // Compute all the scalar values of the results
     std::vector<Scalar<TR>> results;
-    if (TotalElementCount(shape) > 0) {
+    std::optional<uint64_t> n{TotalElementCount(shape)};
+    if (!n) {
+      context.messages().Say(
+          "Too many elements in elemental intrinsic function result"_err_en_US);
+      return Expr<TR>{std::move(funcRef)};
+    }
+    if (*n > 0) {
       ConstantBounds bounds{shape};
       ConstantSubscripts resultIndex(rank, 1);
       ConstantSubscripts argIndex[]{std::get<I>(*args)->lbounds()...};
@@ -878,37 +884,41 @@ template <typename T> Expr<T> Folder<T>::RESHAPE(FunctionRef<T> &&funcRef) {
   } else if (HasNegativeExtent(shape.value())) {
     context_.messages().Say(
         "'shape=' argument must not have a negative extent"_err_en_US);
-  } else if (TotalElementCountOverflows(shape.value())) {
-    context_.messages().Say(
-        "'shape=' argument has too many elements"_err_en_US);
   } else {
-    int rank{GetRank(shape.value())};
-    std::size_t resultElements{TotalElementCount(shape.value())};
-    std::optional<std::vector<int>> dimOrder;
-    if (order) {
-      dimOrder = ValidateDimensionOrder(rank, *order);
-    }
-    std::vector<int> *dimOrderPtr{dimOrder ? &dimOrder.value() : nullptr};
-    if (order && !dimOrder) {
-      context_.messages().Say("Invalid 'order=' argument in RESHAPE"_err_en_US);
-    } else if (resultElements > source->size() && (!pad || pad->empty())) {
+    std::optional<uint64_t> optResultElement{TotalElementCount(shape.value())};
+    if (!optResultElement) {
       context_.messages().Say(
-          "Too few elements in 'source=' argument and 'pad=' "
-          "argument is not present or has null size"_err_en_US);
+          "'shape=' argument has too many elements"_err_en_US);
     } else {
-      Constant<T> result{!source->empty() || !pad
-              ? source->Reshape(std::move(shape.value()))
-              : pad->Reshape(std::move(shape.value()))};
-      ConstantSubscripts subscripts{result.lbounds()};
-      auto copied{result.CopyFrom(*source,
-          std::min(source->size(), resultElements), subscripts, dimOrderPtr)};
-      if (copied < resultElements) {
-        CHECK(pad);
-        copied += result.CopyFrom(
-            *pad, resultElements - copied, subscripts, dimOrderPtr);
+      int rank{GetRank(shape.value())};
+      uint64_t resultElements{*optResultElement};
+      std::optional<std::vector<int>> dimOrder;
+      if (order) {
+        dimOrder = ValidateDimensionOrder(rank, *order);
+      }
+      std::vector<int> *dimOrderPtr{dimOrder ? &dimOrder.value() : nullptr};
+      if (order && !dimOrder) {
+        context_.messages().Say(
+            "Invalid 'order=' argument in RESHAPE"_err_en_US);
+      } else if (resultElements > source->size() && (!pad || pad->empty())) {
+        context_.messages().Say(
+            "Too few elements in 'source=' argument and 'pad=' "
+            "argument is not present or has null size"_err_en_US);
+      } else {
+        Constant<T> result{!source->empty() || !pad
+                ? source->Reshape(std::move(shape.value()))
+                : pad->Reshape(std::move(shape.value()))};
+        ConstantSubscripts subscripts{result.lbounds()};
+        auto copied{result.CopyFrom(*source,
+            std::min(source->size(), resultElements), subscripts, dimOrderPtr)};
+        if (copied < resultElements) {
+          CHECK(pad);
+          copied += result.CopyFrom(
+              *pad, resultElements - copied, subscripts, dimOrderPtr);
+        }
+        CHECK(copied == resultElements);
+        return Expr<T>{std::move(result)};
       }
-      CHECK(copied == resultElements);
-      return Expr<T>{std::move(result)};
     }
   }
   // Invalid, prevent re-folding
@@ -947,14 +957,19 @@ template <typename T> Expr<T> Folder<T>::SPREAD(FunctionRef<T> &&funcRef) {
     ConstantSubscripts shape{source->shape()};
     shape.insert(shape.begin() + *dim - 1, *ncopies);
     Constant<T> spread{source->Reshape(std::move(shape))};
-    std::vector<int> dimOrder;
-    for (int j{0}; j < sourceRank; ++j) {
-      dimOrder.push_back(j < *dim - 1 ? j : j + 1);
-    }
-    dimOrder.push_back(*dim - 1);
-    ConstantSubscripts at{spread.lbounds()}; // all 1
-    spread.CopyFrom(*source, TotalElementCount(spread.shape()), at, &dimOrder);
-    return Expr<T>{std::move(spread)};
+    std::optional<uint64_t> n{TotalElementCount(spread.shape())};
+    if (!n) {
+      context_.messages().Say("Too many elements in SPREAD result"_err_en_US);
+    } else {
+      std::vector<int> dimOrder;
+      for (int j{0}; j < sourceRank; ++j) {
+        dimOrder.push_back(j < *dim - 1 ? j : j + 1);
+      }
+      dimOrder.push_back(*dim - 1);
+      ConstantSubscripts at{spread.lbounds()}; // all 1
+      spread.CopyFrom(*source, *n, at, &dimOrder);
+      return Expr<T>{std::move(spread)};
+    }
   }
   // Invalid, prevent re-folding
   return MakeInvalidIntrinsic(std::move(funcRef));
diff --git a/flang/lib/Evaluate/initial-image.cpp b/flang/lib/Evaluate/initial-image.cpp
index a0fe4ec95da94d7..3b0d738c422d4f4 100644
--- a/flang/lib/Evaluate/initial-image.cpp
+++ b/flang/lib/Evaluate/initial-image.cpp
@@ -18,7 +18,11 @@ auto InitialImage::Add(ConstantSubscript offset, std::size_t bytes,
   if (offset < 0 || offset + bytes > data_.size()) {
     return OutOfRange;
   } else {
-    auto elements{TotalElementCount(x.shape())};
+    auto optElements{TotalElementCount(x.shape())};
+    if (!optElements) {
+      return TooManyElems;
+    }
+    auto elements{*optElements};
     auto elementBytes{bytes > 0 ? bytes / elements : 0};
     if (elements * elementBytes != bytes) {
       return SizeMismatch;
@@ -89,7 +93,9 @@ class AsConstantHelper {
     }
     using Const = Constant<T>;
     using Scalar = typename Const::Element;
-    std::size_t elements{TotalElementCount(extents_)};
+    std::optional<uint64_t> optElements{TotalElementCount(extents_)};
+    CHECK(optElements);
+    uint64_t elements{*optElements};
     std::vector<Scalar> typedValue(elements);
     auto elemBytes{ToInt64(type_.MeasureSizeInBytes(
         context_, GetRank(extents_) > 0, charLength_))};
diff --git a/flang/lib/Semantics/data-to-inits.cpp b/flang/lib/Semantics/data-to-inits.cpp
index bc0355a2c597a6f..4ff0ed1fbe49fe8 100644
--- a/flang/lib/Semantics/data-to-inits.cpp
+++ b/flang/lib/Semantics/data-to-inits.cpp
@@ -456,6 +456,8 @@ bool DataInitializationCompiler<DSV>::InitElement(
         exprAnalyzer_.Say(
             "DATA statement value '%s' for '%s' has the wrong length"_warn_en_US,
             folded.AsFortran(), DescribeElement());
+      } else if (status == evaluate::InitialImage::TooManyElems) {
+        exprAnalyzer_.Say("DATA statement has too many elements"_err_en_US);
       } else {
         CHECK(exprAnalyzer_.context().AnyFatalError());
       }



More information about the flang-commits mailing list