[flang][HLFIR] Fix crash in WHERE with exactly_once inside elemental (#194443)

Fix a segfault in LowerHLFIROrderedAssignments when compiling a WHERE
statement whose mask contains an array constructor with an implied-do
loop (e.g. WHERE([(f(J), J=1,N)]) ...). The hlfir.exactly_once op inside
the hlfir.elemental has live-in values that are block arguments from the
enclosing elemental, which canonicalizeExactlyOnceInsideWhere cannot
pull into the exactly_once region.

The fix has two parts:

1. In canonicalizeExactlyOnceInsideWhere, skip exactly_once ops nested
inside hlfir.elemental and skip block argument live-ins, since these
cannot be relocated.

2. In both overloads of inlineElementalOp, handle hlfir.exactly_once by
inlining its body and cleanup operations instead of cloning the op
verbatim (which left an illegal op after lowering).

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Co-Authored-By: @jeanPerier

---------

Co-authored-by: Claude Opus 4 (1M context) <noreply@anthropic.com>
Co-authored-by: Jean Perier <jperier@nvidia.com>

flang/include/flang/Optimizer/Builder/HLFIRTools.h

@@ -614,6 +614,10 @@ bool designatePreservesContinuity(hlfir::DesignateOp op);
 /// and it can be used on pure FIR representation as well as on HLFIR.
 bool isSimplyContiguous(mlir::Value base, bool checkWhole = true);
 
+/// Return true if \p region belongs to a sub-expression of an hlfir.where
+/// that is subject to the WHERE mask control.
+bool isInsideHlfirWhereMaskedExpression(mlir::Region &region);
+
 } // namespace hlfir
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_HLFIRTOOLS_H

flang/lib/Lower/ConvertArrayConstructor.cpp

@@ -20,18 +20,6 @@
 #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.
 //
@@ -797,7 +785,7 @@ hlfir::EntityWithAttributes Fortran::lower::ArrayConstructorBuilder<T>::gen(
   // 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) &&
+  if (hlfir::isInsideHlfirWhereMaskedExpression(builder.getRegion()) &&
       !builder.getRegion().getParentOfType<hlfir::ExactlyOnceOp>()) {
     Fortran::lower::StatementContext localStmtCtx;
     mlir::Type bogusType = builder.getIndexType();

flang/lib/Lower/ConvertCall.cpp

@@ -3262,19 +3262,13 @@ bool Fortran::lower::isIntrinsicModuleProcRef(
   return module && module->attrs().test(Fortran::semantics::Attr::INTRINSIC);
 }
 
-static bool isInWhereMaskedExpression(fir::FirOpBuilder &builder) {
-  // The MASK of the outer WHERE is not masked itself.
-  mlir::Operation *op = builder.getRegion().getParentOp();
-  return op && op->getParentOfType<hlfir::WhereOp>();
-}
-
 std::optional<hlfir::EntityWithAttributes> Fortran::lower::convertCallToHLFIR(
     mlir::Location loc, Fortran::lower::AbstractConverter &converter,
     const evaluate::ProcedureRef &procRef, std::optional<mlir::Type> resultType,
     Fortran::lower::SymMap &symMap, Fortran::lower::StatementContext &stmtCtx) {
   auto &builder = converter.getFirOpBuilder();
   if (resultType && !procRef.IsElemental() &&
-      isInWhereMaskedExpression(builder) &&
+      hlfir::isInsideHlfirWhereMaskedExpression(builder.getRegion()) &&
       !builder.getRegion().getParentOfType<hlfir::ExactlyOnceOp>()) {
     // Non elemental calls inside a where-assignment-stmt must be executed
     // exactly once without mask control. Lower them in a special region so that

flang/lib/Optimizer/Builder/HLFIRTools.cpp

@@ -1811,3 +1811,16 @@ bool hlfir::isSimplyContiguous(mlir::Value base, bool checkWhole) {
           [&](fir::ConvertOp op) { return isSimplyContiguous(op.getValue()); })
       .Default([](auto &&) { return false; });
 }
+
+bool hlfir::isInsideHlfirWhereMaskedExpression(mlir::Region &region) {
+  hlfir::WhereOp whereOp = region.getParentOfType<hlfir::WhereOp>();
+  if (!whereOp)
+    return false;
+  // If the where is nested inside another where, even its mask region must
+  // be evaluated masked.
+  if (whereOp->getParentOfType<hlfir::WhereOp>())
+    return true;
+  // The top-level where mask is not itself controlled by any other where mask,
+  // all other expressions nested under the where must be evaluated masked.
+  return !whereOp.getMaskRegion().isAncestor(&region);
+}

flang/test/Lower/HLFIR/where-implied-do-mask.f90

@@ -0,0 +1,22 @@
+! RUN: bbc -emit-hlfir -o - %s | FileCheck %s
+! RUN: %flang_fc1 -emit-llvm -o /dev/null %s
+
+module m
+  implicit none
+contains
+  subroutine test(arr, mask, default)
+    real, intent(inout) :: arr(:)
+    integer, intent(in) :: mask(:)
+    real, intent(in) :: default
+    integer :: j
+    where([(.not.any(mask == j), j = 1, size(arr))]) arr = default
+  end subroutine
+end module
+
+! CHECK-LABEL: func.func @_QMmPtest(
+! CHECK: hlfir.where
+! CHECK:   hlfir.elemental
+! CHECK:     hlfir.any
+! CHECK:   hlfir.yield
+! CHECK: hlfir.region_assign
+! CHECK: return