[flang-commits] [flang] [flang] Fix crash in reduction folding (PR #87201)

[Peter Klausler via flang-commits]

https://github.com/klausler created https://github.com/llvm/llvm-project/pull/87201

A reduction folding template assumed lower bounds were 1.

Fixes https://github.com/llvm/llvm-project/issues/86935.

>From 575da4bff12cf136dbdbb86bf8c59dfad5417528 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler at nvidia.com>
Date: Sun, 31 Mar 2024 12:58:23 -0700
Subject: [PATCH] [flang] Fix crash in reduction folding

A reduction folding template assumed lower bounds were 1.

Fixes https://github.com/llvm/llvm-project/issues/86935.
---
 flang/lib/Evaluate/fold-reduction.h | 24 +++++++++++++-----------
 flang/test/Evaluate/folding32.f90   |  6 ++++++
 2 files changed, 19 insertions(+), 11 deletions(-)
 create mode 100644 flang/test/Evaluate/folding32.f90

diff --git a/flang/lib/Evaluate/fold-reduction.h b/flang/lib/Evaluate/fold-reduction.h
index 1ee957c0faebd8..c84d35734ab5af 100644
--- a/flang/lib/Evaluate/fold-reduction.h
+++ b/flang/lib/Evaluate/fold-reduction.h
@@ -182,25 +182,27 @@ static Constant<T> DoReduction(const Constant<ARRAY> &array,
     ConstantSubscript &maskDimAt{maskAt[*dim - 1]};
     ConstantSubscript maskDimLbound{maskDimAt};
     for (auto n{GetSize(resultShape)}; n-- > 0;
-         IncrementSubscripts(at, array.shape()),
-         IncrementSubscripts(maskAt, mask.shape())) {
-      dimAt = dimLbound;
-      maskDimAt = maskDimLbound;
+         array.IncrementSubscripts(at), mask.IncrementSubscripts(maskAt)) {
       elements.push_back(identity);
-      bool firstUnmasked{true};
-      for (ConstantSubscript j{0}; j < dimExtent; ++j, ++dimAt, ++maskDimAt) {
-        if (mask.At(maskAt).IsTrue()) {
-          accumulator(elements.back(), at, firstUnmasked);
-          firstUnmasked = false;
+      if (dimExtent > 0) {
+        dimAt = dimLbound;
+        maskDimAt = maskDimLbound;
+        bool firstUnmasked{true};
+        for (ConstantSubscript j{0}; j < dimExtent; ++j, ++dimAt, ++maskDimAt) {
+          if (mask.At(maskAt).IsTrue()) {
+            accumulator(elements.back(), at, firstUnmasked);
+            firstUnmasked = false;
+          }
         }
+        --dimAt, --maskDimAt;
       }
       accumulator.Done(elements.back());
     }
   } else { // no DIM=, result is scalar
     elements.push_back(identity);
     bool firstUnmasked{true};
-    for (auto n{array.size()}; n-- > 0; IncrementSubscripts(at, array.shape()),
-         IncrementSubscripts(maskAt, mask.shape())) {
+    for (auto n{array.size()}; n-- > 0;
+         array.IncrementSubscripts(at), mask.IncrementSubscripts(maskAt)) {
       if (mask.At(maskAt).IsTrue()) {
         accumulator(elements.back(), at, firstUnmasked);
         firstUnmasked = false;
diff --git a/flang/test/Evaluate/folding32.f90 b/flang/test/Evaluate/folding32.f90
new file mode 100644
index 00000000000000..e4c8b26ca8fdc3
--- /dev/null
+++ b/flang/test/Evaluate/folding32.f90
@@ -0,0 +1,6 @@
+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Fold NORM2 reduction of array with non-default lower bound
+module m
+  real, parameter :: a(2:3) = 0.0
+  logical, parameter :: test1 = norm2(a) == 0.
+end