[Flang] Wrap array constructors within a hlfir.exactly_once op

[ceseo]

When inside a WHERE construct, the array constructor should be generated within an hlfir.exactly_once region.

Fixes #130532

[llvmbot]

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

Author: Carlos Seo (ceseo)

Changes

When inside a WHERE construct, the array constructor should be generated within an hlfir.exactly_once region.

Fixes #130532

---

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

2 Files Affected:

• (modified) flang/lib/Lower/ConvertArrayConstructor.cpp (+47)
• (added) flang/test/Lower/array-constructor-exactly-once.f90 (+20)

diff --git a/flang/lib/Lower/ConvertArrayConstructor.cpp b/flang/lib/Lower/ConvertArrayConstructor.cpp
index 006f022b5379a..558aad1685739 100644
--- a/flang/lib/Lower/ConvertArrayConstructor.cpp
+++ b/flang/lib/Lower/ConvertArrayConstructor.cpp
@@ -20,6 +20,18 @@
 #include "flang/Optimizer/Builder/Todo.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 
+namespace {
+/// Check if we are inside a WHERE construct's masked expression region.
+/// Array constructors inside WHERE statements must be evaluated exactly once
+/// without mask control, similar to non-elemental function calls.
+
+static bool isInWhereMaskedExpression(fir::FirOpBuilder &builder) {
+  mlir::Operation *op = builder.getRegion().getParentOp();
+  return op && op->getParentOfType<hlfir::WhereOp>();
+}
+
+} // namespace
+
 // Array constructors are lowered with three different strategies.
 // All strategies are not possible with all array constructors.
 //
@@ -780,6 +792,41 @@ hlfir::EntityWithAttributes Fortran::lower::ArrayConstructorBuilder<T>::gen(
     const Fortran::evaluate::ArrayConstructor<T> &arrayCtorExpr,
     Fortran::lower::SymMap &symMap, Fortran::lower::StatementContext &stmtCtx) {
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+
+  // Array constructors inside a where-assignment-stmt must be executed
+  // exactly once without mask control, per Fortran 2023 section 10.2.3.2.
+  // Lower them in a special region so that this can be enforced when
+  // scheduling forall/where expression evaluations.
+  if (isInWhereMaskedExpression(builder) &&
+      !builder.getRegion().getParentOfType<hlfir::ExactlyOnceOp>()) {
+    Fortran::lower::StatementContext localStmtCtx;
+    mlir::Type bogusType = builder.getIndexType();
+    auto exactlyOnce = hlfir::ExactlyOnceOp::create(builder, loc, bogusType);
+    mlir::Block *block = builder.createBlock(&exactlyOnce.getBody());
+    builder.setInsertionPointToStart(block);
+
+    // Recursively generate the array constructor inside the exactly_once region
+    hlfir::EntityWithAttributes res = ArrayConstructorBuilder<T>::gen(
+        loc, converter, arrayCtorExpr, symMap, localStmtCtx);
+
+    auto yield = hlfir::YieldOp::create(builder, loc, res);
+    Fortran::lower::genCleanUpInRegionIfAny(loc, builder, yield.getCleanup(),
+                                            localStmtCtx);
+    builder.setInsertionPointAfter(exactlyOnce);
+    exactlyOnce->getResult(0).setType(res.getType());
+
+    if (hlfir::isFortranValue(exactlyOnce.getResult()))
+      return hlfir::EntityWithAttributes{exactlyOnce.getResult()};
+
+    // Create hlfir.declare for the result to satisfy
+    // hlfir::EntityWithAttributes requirements.
+    auto [exv, cleanup] = hlfir::translateToExtendedValue(
+        loc, builder, hlfir::Entity{exactlyOnce});
+    assert(!cleanup && "result is a variable");
+    return hlfir::genDeclare(loc, builder, exv, ".arrayctor.result",
+                             fir::FortranVariableFlagsAttr{});
+  }
+
   // Select the lowering strategy given the array constructor.
   auto arrayBuilder = selectArrayCtorLoweringStrategy(
       loc, converter, arrayCtorExpr, symMap, stmtCtx);
diff --git a/flang/test/Lower/array-constructor-exactly-once.f90 b/flang/test/Lower/array-constructor-exactly-once.f90
new file mode 100644
index 0000000000000..04ab6a0ff3fc5
--- /dev/null
+++ b/flang/test/Lower/array-constructor-exactly-once.f90
@@ -0,0 +1,20 @@
+! RUN: flang -fc1 -emit-hlfir %s -o - | FileCheck %s
+
+program main
+  call test06()
+  print *,'pass'
+end program main
+
+subroutine test06()
+  type ty1
+     integer ,allocatable :: a(:,:,:)
+  end type ty1
+  type(ty1) :: str(1)
+  integer ,allocatable :: b(:,:,:)
+  allocate(str(1)%a(1,1,1),b(1,1,1))
+  b=1
+  write(6,*) "b                                 = ", b
+  write(6,*) "reshape((/(b,jj=1,1)/),(/1,1,1/)) = ", reshape((/(b,jj=1,1)/),(/1,1,1/))
+  where ((/.true./)) str=(/(ty1(reshape((/(b,jj=1,1)/),(/1,1,1/))),ii=1,1)/)
+  ! CHECK: hlfir.exactly_once : !hlfir.expr<1x!fir.type<_QFtest06Tty1{a:!fir.box<!fir.heap<!fir.array<?x?x?xi32>>>}>>
+end subroutine test06

[jeanPerier]

Thanks!