[flang][lowering] delay stack save/restor emission in elemental calls

[jeanPerier]

stack save/restore emitted for character elemental function result allocation inside hlfir.elemental in lowering created memory bugs because result memory is actually still used after the stack restore when lowering the elemental into a loop where the result element is copied into the array result storage.

Instead of adding special handling for stack save/restore in lowering, just avoid emitting those since the stack reclaim pass is able to emit them in the generated loop. Not having those stack save/restore will also help optimizations that want to elide the temporary allocation for the element result when that is possible.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: None (jeanPerier)

Changes

stack save/restore emitted for character elemental function result allocation inside hlfir.elemental in lowering created memory bugs because result memory is actually still used after the stack restore when lowering the elemental into a loop where the result element is copied into the array result storage.

Instead of adding special handling for stack save/restore in lowering, just avoid emitting those since the stack reclaim pass is able to emit them in the generated loop. Not having those stack save/restore will also help optimizations that want to elide the temporary allocation for the element result when that is possible.

---

Full diff: https://github.com/llvm/llvm-project/pull/109142.diff

3 Files Affected:

• (modified) flang/lib/Lower/ConvertCall.cpp (+5-1)
• (modified) flang/test/Lower/HLFIR/elemental-array-ops.f90 (-2)
• (added) flang/test/Lower/HLFIR/elemental-user-procedure-stacksave.f90 (+22)

diff --git a/flang/lib/Lower/ConvertCall.cpp b/flang/lib/Lower/ConvertCall.cpp
index 2fedc01bc77fc1..017bfd049d3dc5 100644
--- a/flang/lib/Lower/ConvertCall.cpp
+++ b/flang/lib/Lower/ConvertCall.cpp
@@ -366,7 +366,11 @@ std::pair<fir::ExtendedValue, bool> Fortran::lower::genCallOpAndResult(
       resultLengths = lengths;
     }
 
-    if (!extents.empty() || !lengths.empty()) {
+    if ((!extents.empty() || !lengths.empty()) && !isElemental) {
+      // Note: in the elemental context, the alloca ownership inside the
+      // elemental region is implicit, and later pass in lowering (stack
+      // reclaim) fir.do_loop will be in charge of emitting any stack
+      // save/restore if needed.
       auto *bldr = &converter.getFirOpBuilder();
       mlir::Value sp = bldr->genStackSave(loc);
       stmtCtx.attachCleanup(
diff --git a/flang/test/Lower/HLFIR/elemental-array-ops.f90 b/flang/test/Lower/HLFIR/elemental-array-ops.f90
index 9929c17ec33994..18e1fb0a787e73 100644
--- a/flang/test/Lower/HLFIR/elemental-array-ops.f90
+++ b/flang/test/Lower/HLFIR/elemental-array-ops.f90
@@ -182,12 +182,10 @@ end subroutine char_return
 ! CHECK:             %[[VAL_23:.*]] = arith.constant 0 : index
 ! CHECK:             %[[VAL_24:.*]] = arith.cmpi sgt, %[[VAL_22]], %[[VAL_23]] : index
 ! CHECK:             %[[VAL_25:.*]] = arith.select %[[VAL_24]], %[[VAL_22]], %[[VAL_23]] : index
-! CHECK:             %[[VAL_26:.*]] = llvm.intr.stacksave : !llvm.ptr
 ! CHECK:             %[[VAL_27:.*]] = fir.call @_QPcallee(%[[VAL_2]], %[[VAL_25]], %[[VAL_20]]) fastmath<contract> : (!fir.ref<!fir.char<1,3>>, index, !fir.boxchar<1>) -> !fir.boxchar<1>
 ! CHECK:             %[[VAL_28:.*]]:2 = hlfir.declare %[[VAL_2]] typeparams %[[VAL_25]] {uniq_name = ".tmp.func_result"} : (!fir.ref<!fir.char<1,3>>, index) -> (!fir.ref<!fir.char<1,3>>, !fir.ref<!fir.char<1,3>>)
 ! CHECK:             %[[MustFree:.*]] = arith.constant false
 ! CHECK:             %[[ResultTemp:.*]] = hlfir.as_expr %[[VAL_28]]#0 move %[[MustFree]] : (!fir.ref<!fir.char<1,3>>, i1) -> !hlfir.expr<!fir.char<1,3>>
-! CHECK:             llvm.intr.stackrestore %[[VAL_26]] : !llvm.ptr
 ! CHECK:             hlfir.yield_element %[[ResultTemp]] : !hlfir.expr<!fir.char<1,3>>
 ! CHECK:           }
 ! CHECK:           %[[VAL_29:.*]] = arith.constant 0 : index
diff --git a/flang/test/Lower/HLFIR/elemental-user-procedure-stacksave.f90 b/flang/test/Lower/HLFIR/elemental-user-procedure-stacksave.f90
new file mode 100644
index 00000000000000..839342f2da6d85
--- /dev/null
+++ b/flang/test/Lower/HLFIR/elemental-user-procedure-stacksave.f90
@@ -0,0 +1,22 @@
+! Check that stack save and restore needed for elemental function result
+! allocation inside loops are not emitted directly in lowering, but inserted if
+! needed in the stack-reclaim pass.
+
+! RUN: %flang_fc1 -emit-hlfir %s -o - | FileCheck %s --check-prefix=CHECK-HLFIR
+! RUN: %flang_fc1 -emit-llvm %s -o - | FileCheck %s --check-prefix=CHECK-LLVM
+subroutine foo(c1, c2)
+  character(*), dimension(100) :: c1, c2
+  interface
+    elemental pure function func(c)
+      character(*), intent(in) :: c
+      character(len(c)) :: func
+    end function
+  end interface
+  c1 = func(c2)
+end subroutine
+
+! CHECK-HLFIR-NOT: stacksave
+! CHECK: return
+
+! CHECK-LLVM: stacksave
+! CHECK-LLVM: stackrestore

[clementval]

LGTM

[tblah]

Thanks!