[flang][hlfir] Add codegen for vector subscripted LHS

[jeanPerier]

Summary

This patch adds support for vector subscripted assignment left-hand
side. It does not yet add support for the cases where the LHS must be
saved because its evaluation could be impacted by the assignment.

The implementation adds an hlfir::ElementalOpInterface to share the
elemental inlining utility and some other tools between
hlfir::ElementalOp and hlfir::ElelemntalAddrOp.

It adds generateYieldedLHS() to allow retrieving the LHS value
in lowering, whether or not it is vector subscripted. If it is vector
subscripted, this utility creates a loop nest iterating over the
elements and returns the address of an element.

[tblah]

One tiny nit. Otherwise this looks great to me :)

flang/test/HLFIR/order_assignments/vector-subscripts-codegen.fir

120

Should this be destroying VAL_16 (the result of the elemental)?

[jeanPerier]

flang/test/HLFIR/order_assignments/vector-subscripts-codegen.fir

120

Thanks for the review! Yes, this is indeed destroying VAL_16, for some reason the generate-test-checks.py utility is not making the link and is creating a new capture here. I updated it manually.

[jeanPerier]

Update LIT test to use the correct value instead of creating a new
capture with the hlfir.expr argument.

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

 namespace mlir {
 class IRMapping;
 }
 
 namespace hlfir {
 
 class AssociateOp;
 class ElementalOp;
+class ElementalOpInterface;
 class ElementalAddrOp;
 class YieldElementOp;
 
 /// Is this an SSA value type for the value of a Fortran procedure
 /// designator ?
 inline bool isFortranProcedureValue(mlir::Type type) {
   return type.isa<fir::BoxProcType>() ||
          (type.isa<mlir::TupleType>() &&

 /// of one element of the elemental expression. Return the YieldElementOp
 /// whose value argument is the element value.
 /// The original hlfir::ElementalOp is left untouched.
 hlfir::YieldElementOp inlineElementalOp(mlir::Location loc,
                                         fir::FirOpBuilder &builder,
                                         hlfir::ElementalOp elemental,
                                         mlir::ValueRange oneBasedIndices);
 
-/// Inline the body of an hlfir.elemental without cloning the resulting
-/// hlfir.yield_element, and return the cloned operand of the
-/// hlfir.yield_element. The mapper must be provided to cover complex cases
-/// where the inlined elemental is not defined in the current context and uses
-/// values that have been cloned already.
-/// A callback is provided to indicate if an hlfir.apply inside the
-/// hlfir.elemental must be immediately replaced by the inlining of the
-/// applied hlfir.elemental.
+/// Inline the body of an hlfir.elemental or hlfir.elemental_addr without
+/// cloning the resulting hlfir.yield_element/hlfir.yield, and return the cloned
+/// operand of the hlfir.yield_element/hlfir.yield. The mapper must be provided
+/// to cover complex cases where the inlined elemental is not defined in the
+/// current context and uses values that have been cloned already. A callback is
+/// provided to indicate if an hlfir.apply inside the hlfir.elemental must be
+/// immediately replaced by the inlining of the applied hlfir.elemental.
 mlir::Value inlineElementalOp(
     mlir::Location loc, fir::FirOpBuilder &builder,
-    hlfir::ElementalOp elemental, mlir::ValueRange oneBasedIndices,
+    hlfir::ElementalOpInterface elemental, mlir::ValueRange oneBasedIndices,
     mlir::IRMapping &mapper,
     const std::function<bool(hlfir::ElementalOp)> &mustRecursivelyInline);
 
 std::pair<fir::ExtendedValue, std::optional<hlfir::CleanupFunction>>
 convertToValue(mlir::Location loc, fir::FirOpBuilder &builder,
                const hlfir::Entity &entity);
 
 std::pair<fir::ExtendedValue, std::optional<hlfir::CleanupFunction>>

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

def hlfir_NoReassocOp : hlfir_Op<"no_reassoc", [NoMemoryEffect, SameOperandsAndResultType]> {
   }];
 
   let arguments = (ins AnyFortranEntity:$val);
   let results = (outs AnyFortranEntity);
 
   let assemblyFormat = "$val attr-dict `:` type($val)";
 }
 
-def hlfir_ElementalOp : hlfir_Op<"elemental", [RecursiveMemoryEffects]> {
+def hlfir_ElementalOpInterface : OpInterface<"ElementalOpInterface"> {
+  let description = [{
+    Interface for the operation holding a region with elemental computation.
+    It is used as a common interface bewteen hlfir.elemental and hlfir.elemental_addr.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/"Return the one based elemental indices.",
+      /*retTy=*/"mlir::Block::BlockArgListType",
+      /*methodName=*/"getIndices",
+      /*args=*/(ins),
+      /*methodBody=*/[{}]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the element entity being computed",
+      /*retTy=*/"mlir::Value",
+      /*methodName=*/"getElementEntity",
+      /*args=*/(ins),
+      /*methodBody=*/[{}]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Get element cleanup region, if any.",
+      /*retTy=*/"mlir::Region*",
+      /*methodName=*/"getElementCleanup",
+      /*args=*/(ins),
+      /*methodBody=*/[{}]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Get elemental region.",
+      /*retTy=*/"mlir::Region&",
+      /*methodName=*/"getElementalRegion",
+      /*args=*/(ins),
+      /*methodBody=*/[{}]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Must this elemental operation be evaluated in order?",
+      /*retTy=*/"bool",
+      /*methodName=*/"isOrdered",
+      /*args=*/(ins),
+      /*methodBody=*/[{}]
+    >,
+  ];
+
+  let cppNamespace = "hlfir";
+}
+
+def hlfir_ElementalOp : hlfir_Op<"elemental", [RecursiveMemoryEffects, hlfir_ElementalOpInterface]> {
   let summary = "elemental expression";
   let description = [{
     Represent an elemental expression as a function of the indices.
     This operation contain a region whose block arguments are one
     based indices iterating over the elemental expression shape.
     Given these indices, the element value for the given iteration
     can be computed in the region and yielded with the hlfir.yield_element
     operation.
def hlfir_ElementalOp : hlfir_Op<"elemental", [RecursiveMemoryEffects, hlfir_ElementalOpInterface]> {
 
   let extraClassDeclaration = [{
       mlir::Block *getBody() { return &getRegion().front(); }
 
       /// Get the indices iterating over the shape.
       mlir::Block::BlockArgListType getIndices() {
        return getBody()->getArguments();
       }
+      /// ElementalOpInterface implementation.
+
+      mlir::Region& getElementalRegion() {return getRegion();}
+      mlir::Value getElementEntity();
+      mlir::Region* getElementCleanup() {return nullptr;}
 
       /// Must this elemental be evaluated in order?
       /// TODO: add attribute and set it in lowering.
       bool isOrdered() {return true;}
   }];
 
   let skipDefaultBuilders = 1;
   let builders = [
def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp",
   }];
 
   let arguments = (ins AnyFortranEntity:$entity);
   let regions = (region  MaxSizedRegion<1>:$cleanup);
 
   let assemblyFormat = "$entity attr-dict `:` type($entity) custom<YieldOpCleanup>($cleanup)";
 }
 
-def hlfir_ElementalAddrOp : hlfir_Op<"elemental_addr", [Terminator, HasParent<"RegionAssignOp">, RecursiveMemoryEffects, RecursivelySpeculatable]> {
+def hlfir_ElementalAddrOp : hlfir_Op<"elemental_addr", [Terminator, HasParent<"RegionAssignOp">, RecursiveMemoryEffects, RecursivelySpeculatable, hlfir_ElementalOpInterface]> {
   let summary = "Yield the address of a vector subscripted variable inside an hlfir.region_assign";
   let description = [{
     Special terminator node for the left-hand side region of an hlfir.region_assign
     to a vector subscripted entity.
 
     It represents how the address of an element of such entity is computed given
     one based indices.
 
let extraClassDeclaration = [{
     mlir::Region::BlockArgListType getIndices() {
       return getBody().getArguments();
     }
 
     /// Return the hlfir::YieldOp terminator of the operation
     /// body. It yields the variable element address.
     /// This should only be called once the ElementalAddrOp has been built.
     hlfir::YieldOp getYieldOp();
+
+    /// ElementalOpInterface implementation.
+
+    mlir::Region& getElementalRegion() {return getBody();}
+    mlir::Value getElementEntity();
+    mlir::Region* getElementCleanup();
+
+    /// Must this elemental be evaluated in order?
+    /// TODO: add attribute and set it in lowering.
+    bool isOrdered() {return true;}
   }];
 
   let hasVerifier = 1;
 }
 
 /// Define ODS constraints to verify that a region ends with a yield of a
 /// certain type.
 def YieldIntegerOrEmpty : CPred<"yieldsIntegerOrEmpty($_self)">;

flang/lib/Optimizer/Builder/HLFIRTools.cpp

for (auto &op : elemental.getRegion().back().getOperations())
     newOp = builder.clone(op, mapper);
   auto yield = mlir::dyn_cast_or_null<hlfir::YieldElementOp>(newOp);
   assert(yield && "last ElementalOp operation must be am hlfir.yield_element");
   return yield;
 }
 
 mlir::Value hlfir::inlineElementalOp(
     mlir::Location loc, fir::FirOpBuilder &builder,
-    hlfir::ElementalOp elemental, mlir::ValueRange oneBasedIndices,
+    hlfir::ElementalOpInterface elemental, mlir::ValueRange oneBasedIndices,
     mlir::IRMapping &mapper,
     const std::function<bool(hlfir::ElementalOp)> &mustRecursivelyInline) {
-  mlir::Region &region = elemental.getRegion();
+  mlir::Region &region = elemental.getElementalRegion();
   // hlfir.elemental region is a SizedRegion<1>.
   assert(region.hasOneBlock() && "elemental region must have one block");
   mapper.map(elemental.getIndices(), oneBasedIndices);
-  mlir::Block::OpListType &ops = region.back().getOperations();
-  assert(!ops.empty() && "elemental block cannot be empty");
-  auto end = ops.end();
-  for (auto opIt = ops.begin(); std::next(opIt) != end; ++opIt) {
-    if (auto apply = mlir::dyn_cast<hlfir::ApplyOp>(*opIt))
+  for (auto &op : region.front().without_terminator()) {
+    if (auto apply = mlir::dyn_cast<hlfir::ApplyOp>(op))
       if (auto appliedElemental =
               apply.getExpr().getDefiningOp<hlfir::ElementalOp>())
         if (mustRecursivelyInline(appliedElemental)) {
           llvm::SmallVector<mlir::Value> clonedApplyIndices;
           for (auto indice : apply.getIndices())
             clonedApplyIndices.push_back(mapper.lookupOrDefault(indice));
-          mlir::Value inlined = inlineElementalOp(
-              loc, builder, appliedElemental, clonedApplyIndices, mapper,
+          hlfir::ElementalOpInterface elementalIface =
+              mlir::cast<hlfir::ElementalOpInterface>(
+                  appliedElemental.getOperation());
+          mlir::Value inlined = inlineElementalOp(loc, builder, elementalIface,
+                                                  clonedApplyIndices, mapper,
                                                   mustRecursivelyInline);
           mapper.map(apply.getResult(), inlined);
           continue;
         }
-    (void)builder.clone(*opIt, mapper);
+    (void)builder.clone(op, mapper);
   }
-  auto oldYield = mlir::dyn_cast_or_null<hlfir::YieldElementOp>(
-      region.back().getOperations().back());
-  assert(oldYield && "must terminate with yieldElementalOp");
-  return mapper.lookupOrDefault(oldYield.getElementValue());
+  return mapper.lookupOrDefault(elemental.getElementEntity());
 }
 
 hlfir::LoopNest hlfir::genLoopNest(mlir::Location loc,
                                    fir::FirOpBuilder &builder,
                                    mlir::ValueRange extents) {
   hlfir::LoopNest loopNest;
   assert(!extents.empty() && "must have at least one extent");
   auto insPt = builder.saveInsertionPoint();

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

void hlfir::ElementalOp::build(mlir::OpBuilder &builder,
   if (auto exprType = resultType.dyn_cast<hlfir::ExprType>()) {
     unsigned dim = exprType.getRank();
     mlir::Type indexType = builder.getIndexType();
     for (unsigned d = 0; d < dim; ++d)
       bodyRegion->front().addArgument(indexType, odsState.location);
   }
 }
 
+mlir::Value hlfir::ElementalOp::getElementEntity() {
+  return mlir::cast<hlfir::YieldElementOp>(getBody()->back()).getElementValue();
+}
+
 //===----------------------------------------------------------------------===//
 // ApplyOp
 //===----------------------------------------------------------------------===//
 
 void hlfir::ApplyOp::build(mlir::OpBuilder &builder,
                            mlir::OperationState &odsState, mlir::Value expr,
                            mlir::ValueRange indices,
                            mlir::ValueRange typeparams) {

 
 hlfir::YieldOp hlfir::ElementalAddrOp::getYieldOp() {
   hlfir::YieldOp yieldOp =
       mlir::dyn_cast_or_null<hlfir::YieldOp>(getTerminator(getBody()));
   assert(yieldOp && "element_addr is ill-formed");
   return yieldOp;
 }
 
+mlir::Value hlfir::ElementalAddrOp::getElementEntity() {
+  return getYieldOp().getEntity();
+}
+
+mlir::Region *hlfir::ElementalAddrOp::getElementCleanup() {
+  mlir::Region *cleanup = &getYieldOp().getCleanup();
+  return cleanup->empty() ? nullptr : cleanup;
+}
+
 //===----------------------------------------------------------------------===//
 // OrderedAssignmentTreeOpInterface
 //===----------------------------------------------------------------------===//
 
 mlir::LogicalResult hlfir::OrderedAssignmentTreeOpInterface::verifyImpl() {
   if (mlir::Region *body = getSubTreeRegion())
     if (!body->empty())
       for (mlir::Operation &op : body->front())

flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp

private:
   void post(hlfir::ForallOp);
   void post(hlfir::ForallMaskOp);
   void post(hlfir::WhereOp);
   void post(hlfir::ElseWhereOp);
   /// Enter (and maybe create) the fir.if else block of an ElseWhereOp,
   /// but do not generate the elswhere mask or the new fir.if.
   void enterElsewhere(hlfir::ElseWhereOp);
 
-  /// Is this an assignment to a vector subscripted entity?
-  static bool hasVectorSubscriptedLhs(hlfir::RegionAssignOp regionAssignOp);
   /// Are there any leaf region in the node that must be saved in the current
   /// run?
   bool mustSaveRegionIn(
       hlfir::OrderedAssignmentTreeOpInterface node,
       llvm::SmallVectorImpl<hlfir::SaveEntity> &saveEntities) const;
   /// Should this node be evaluated in the current run? Saving a region in a
   /// node does not imply the node needs to be evaluated.
   bool
private:
   /// optionally return the (uncloned) yield if there is any clean-up that
   /// should be done after using the entity. Like, generateYieldedScalarValue,
   /// this will return the saved value if the region was saved in a previous
   /// run.
   ValueAndCleanUp
   generateYieldedEntity(mlir::Region &region,
                         std::optional<mlir::Type> castToType = std::nullopt);
 
+  struct LhsValueAndCleanUp {
+    mlir::Value lhs;
+    std::optional<hlfir::YieldOp> elementalCleanup;
+    mlir::Region *nonElementalCleanup = nullptr;
+    std::optional<hlfir::LoopNest> vectorSubscriptLoopNest;
+  };
+
+  /// Generate the left-hand side. If the left-hand side is vector
+  /// subscripted (hlfir.elemental_addr), this will create a loop nest
+  /// (unless it was already created by a WHERE mask) and return the
+  /// element address.
+  LhsValueAndCleanUp generateYieldedLHS(mlir::Location loc,
+                                        mlir::Region &lhsRegion);
+
   /// If \p maybeYield is present and has a clean-up, generate the clean-up
   /// at the current insertion point (by cloning).
   void generateCleanupIfAny(std::optional<hlfir::YieldOp> maybeYield);
+  void generateCleanupIfAny(mlir::Region *cleanupRegion);
 
   /// Generate a masked entity. This can only be called when whereLoopNest was
   /// set (When an hlfir.where is being visited).
   /// This method returns the scalar element (that may have been previously
   /// saved) for the current indices inside the where loop.
   mlir::Value generateMaskedEntity(mlir::Location loc, mlir::Region &region) {
     MaskedArrayExpr maskedExpr(loc, region);
     return generateMaskedEntity(maskedExpr);
convertToMoldType(mlir::Location loc, fir::FirOpBuilder &builder,
   }
   // Variable to Variable mismatch (e.g., fir.heap<T> vs fir.ref<T>), or value
   // to Value mismatch (e.g. i1 vs fir.logical<4>).
   return hlfir::Entity{builder.createConvert(loc, mold.getType(), input)};
 }
 
 void OrderedAssignmentRewriter::pre(hlfir::RegionAssignOp regionAssignOp) {
   mlir::Location loc = regionAssignOp.getLoc();
-  auto [rhs, oldRhsYield] =
+  std::optional<hlfir::LoopNest> elementalLoopNest;
+  auto [rhsValue, oldRhsYield] =
       generateYieldedEntity(regionAssignOp.getRhsRegion());
-  if (hasVectorSubscriptedLhs(regionAssignOp))
-    TODO(loc, "assignment to vector subscripted entity");
-  auto [lhs, oldLhsYield] =
-      generateYieldedEntity(regionAssignOp.getLhsRegion());
+  LhsValueAndCleanUp loweredLhs =
+      generateYieldedLHS(loc, regionAssignOp.getLhsRegion());
+  hlfir::Entity rhsEntity{rhsValue};
+  hlfir::Entity lhsEntity{loweredLhs.lhs};
+  if (loweredLhs.vectorSubscriptLoopNest)
+    rhsEntity = hlfir::getElementAt(
+        loc, builder, rhsEntity,
+        loweredLhs.vectorSubscriptLoopNest->oneBasedIndices);
   if (!regionAssignOp.getUserDefinedAssignment().empty()) {
     hlfir::Entity userAssignLhs{regionAssignOp.getUserAssignmentLhs()};
     hlfir::Entity userAssignRhs{regionAssignOp.getUserAssignmentRhs()};
-    hlfir::Entity lhsEntity{lhs};
-    hlfir::Entity rhsEntity{rhs};
-    fir::DoLoopOp outerElementalLoop = nullptr;
+    std::optional<hlfir::LoopNest> elementalLoopNest;
     if (lhsEntity.isArray() && userAssignLhs.isScalar()) {
       // Elemental assignment with array argument (the RHS cannot be an array
       // if the LHS is not).
       mlir::Value shape = hlfir::genShape(loc, builder, lhsEntity);
-      hlfir::LoopNest elementalLoopNest =
-          hlfir::genLoopNest(loc, builder, shape);
-      outerElementalLoop = elementalLoopNest.outerLoop;
-      builder.setInsertionPointToStart(elementalLoopNest.innerLoop.getBody());
+      elementalLoopNest = hlfir::genLoopNest(loc, builder, shape);
+      builder.setInsertionPointToStart(elementalLoopNest->innerLoop.getBody());
       lhsEntity = hlfir::getElementAt(loc, builder, lhsEntity,
-                                      elementalLoopNest.oneBasedIndices);
+                                      elementalLoopNest->oneBasedIndices);
       rhsEntity = hlfir::getElementAt(loc, builder, rhsEntity,
-                                      elementalLoopNest.oneBasedIndices);
+                                      elementalLoopNest->oneBasedIndices);
     }
 
     llvm::SmallVector<hlfir::CleanupFunction, 2> argConversionCleanups;
     lhsEntity = convertToMoldType(loc, builder, lhsEntity, userAssignLhs,
                                   argConversionCleanups);
     rhsEntity = convertToMoldType(loc, builder, rhsEntity, userAssignRhs,
                                   argConversionCleanups);
     mapper.map(userAssignLhs, lhsEntity);
     mapper.map(userAssignRhs, rhsEntity);
     for (auto &op :
          regionAssignOp.getUserDefinedAssignment().front().without_terminator())
       (void)builder.clone(op, mapper);
     for (auto &cleanupConversion : argConversionCleanups)
       cleanupConversion();
-    if (outerElementalLoop)
-      builder.setInsertionPointAfter(outerElementalLoop);
+    if (elementalLoopNest)
+      builder.setInsertionPointAfter(elementalLoopNest->outerLoop);
   } else {
     // TODO: preserve allocatable assignment aspects for forall once
     // they are conveyed in hlfir.region_assign.
-    builder.create<hlfir::AssignOp>(loc, rhs, lhs);
+    builder.create<hlfir::AssignOp>(loc, rhsEntity, lhsEntity);
   }
+  generateCleanupIfAny(loweredLhs.elementalCleanup);
+  if (loweredLhs.vectorSubscriptLoopNest)
+    builder.setInsertionPointAfter(
+        loweredLhs.vectorSubscriptLoopNest->outerLoop);
   generateCleanupIfAny(oldRhsYield);
-  generateCleanupIfAny(oldLhsYield);
+  generateCleanupIfAny(loweredLhs.nonElementalCleanup);
 }
 
 void OrderedAssignmentRewriter::generateMaskIfOp(mlir::Value cdt) {
   mlir::Location loc = cdt.getLoc();
   cdt = hlfir::loadTrivialScalar(loc, builder, hlfir::Entity{cdt});
   cdt = builder.createConvert(loc, builder.getI1Type(), cdt);
   auto ifOp = builder.create<fir::IfOp>(cdt.getLoc(), std::nullopt, cdt,
                                         /*withElseRegion=*/false);
mlir::Value OrderedAssignmentRewriter::generateYieldedScalarValue(
   mlir::Location loc = region.getParentOp()->getLoc();
   auto [value, maybeYield] = generateYieldedEntity(region, castToType);
   value = hlfir::loadTrivialScalar(loc, builder, hlfir::Entity{value});
   assert(fir::isa_trivial(value.getType()) && "not a trivial scalar value");
   generateCleanupIfAny(maybeYield);
   return value;
 }
 
+OrderedAssignmentRewriter::LhsValueAndCleanUp
+OrderedAssignmentRewriter::generateYieldedLHS(mlir::Location loc,
+                                              mlir::Region &lhsRegion) {
+  LhsValueAndCleanUp loweredLhs;
+  hlfir::ElementalAddrOp elementalAddrLhs =
+      mlir::dyn_cast<hlfir::ElementalAddrOp>(lhsRegion.back().back());
+  if (elementalAddrLhs && !whereLoopNest) {
+    for (auto &op : lhsRegion.front().without_terminator())
+      (void)builder.clone(op, mapper);
+    mlir::Value newShape = mapper.lookupOrDefault(elementalAddrLhs.getShape());
+    loweredLhs.vectorSubscriptLoopNest =
+        hlfir::genLoopNest(loc, builder, newShape);
+    builder.setInsertionPointToStart(
+        loweredLhs.vectorSubscriptLoopNest->innerLoop.getBody());
+    mapper.map(elementalAddrLhs.getIndices(),
+               loweredLhs.vectorSubscriptLoopNest->oneBasedIndices);
+    for (auto &op : elementalAddrLhs.getBody().front().without_terminator())
+      (void)builder.clone(op, mapper);
+    loweredLhs.elementalCleanup = elementalAddrLhs.getYieldOp();
+    loweredLhs.lhs =
+        mapper.lookupOrDefault(loweredLhs.elementalCleanup->getEntity());
+  } else {
+    auto [lhs, yield] = generateYieldedEntity(lhsRegion);
+    loweredLhs.lhs = lhs;
+    if (yield && !yield->getCleanup().empty())
+      loweredLhs.nonElementalCleanup = &yield->getCleanup();
+  }
+  return loweredLhs;
+}
+
 mlir::Value
 OrderedAssignmentRewriter::generateMaskedEntity(MaskedArrayExpr &maskedExpr) {
   assert(whereLoopNest.has_value() && "must be inside WHERE loop nest");
   auto insertionPoint = builder.saveInsertionPoint();
   if (!maskedExpr.noneElementalPartWasGenerated) {
     // Generate none elemental part before the where loops (but inside the
     // current forall loops if any).
     builder.setInsertionPoint(whereLoopNest->outerLoop);
     maskedExpr.generateNoneElementalPart(builder, mapper);
   }
   // Generate the none elemental part cleanup after the where loops.
   builder.setInsertionPointAfter(whereLoopNest->outerLoop);
   maskedExpr.generateNoneElementalCleanupIfAny(builder, mapper);
   // Generate the value of the current element for the masked expression
   // at the current insertion point (inside the where loops, and any fir.if
   // generated for previous masks).
   builder.restoreInsertionPoint(insertionPoint);
   return maskedExpr.generateElementalParts(
       builder, whereLoopNest->oneBasedIndices, mapper);
 }
 
 void OrderedAssignmentRewriter::generateCleanupIfAny(
     std::optional<hlfir::YieldOp> maybeYield) {
   if (maybeYield.has_value())
-    if (!maybeYield->getCleanup().empty()) {
-      assert(maybeYield->getCleanup().hasOneBlock() &&
-             "region must contain one block");
-      for (auto &op : maybeYield->getCleanup().back().getOperations())
-        if (!mlir::isa<fir::FirEndOp>(op))
-          builder.clone(op, mapper);
-    }
+    generateCleanupIfAny(&maybeYield->getCleanup());
+void OrderedAssignmentRewriter::generateCleanupIfAny(
+    mlir::Region *cleanupRegion) {
+  if (cleanupRegion && !cleanupRegion->empty()) {
+    assert(cleanupRegion->hasOneBlock() && "region must contain one block");
+    for (auto &op : cleanupRegion->back().without_terminator())
+      builder.clone(op, mapper);
   }
-
-bool OrderedAssignmentRewriter::hasVectorSubscriptedLhs(
-    hlfir::RegionAssignOp regionAssignOp) {
-  return mlir::isa<hlfir::ElementalAddrOp>(
-      regionAssignOp.getLhsRegion().back().back());
 }
 
 bool OrderedAssignmentRewriter::mustSaveRegionIn(
     hlfir::OrderedAssignmentTreeOpInterface node,
     llvm::SmallVectorImpl<hlfir::SaveEntity> &saveEntities) const {
   for (auto &action : currentRun->actions)
     if (hlfir::SaveEntity *savedEntity =
             std::get_if<hlfir::SaveEntity>(&action))
if (hlfir::SaveEntity *savedEntity =
       auto assign = std::get<hlfir::RegionAssignOp>(action);
       if (node->isAncestor(assign.getOperation()))
         return true;
     }
   return false;
 }
 
 /// Is the apply using all the elemental indices in order?
-static bool isInOrderApply(hlfir::ApplyOp apply, hlfir::ElementalOp elemental) {
-  if (elemental.getIndices().size() != apply.getIndices().size())
+static bool isInOrderApply(hlfir::ApplyOp apply,
+                           hlfir::ElementalOpInterface elemental) {
+  mlir::Region::BlockArgListType elementalIndices = elemental.getIndices();
+  if (elementalIndices.size() != apply.getIndices().size())
     return false;
   for (auto [elementalIdx, applyIdx] :
-       llvm::zip(elemental.getIndices(), apply.getIndices()))
+       llvm::zip(elementalIndices, apply.getIndices()))
     if (elementalIdx != applyIdx)
       return false;
   return true;
 }
 
-/// Gather the chain of hlfir::ElementalOp, if any, that produced \p value.
+/// Gather the tree of hlfir::ElementalOpInterface use-def, if any, starting
+/// from \p elemental, which may be a nullptr.
 static void
-gatherElementalTree(mlir::Value value,
+gatherElementalTree(hlfir::ElementalOpInterface elemental,
                     llvm::SmallPtrSetImpl<mlir::Operation *> &elementalOps,
                     bool isOutOfOrder) {
-  if (auto elemental = value.getDefiningOp<hlfir::ElementalOp>()) {
+  if (elemental) {
     // Only inline an applied elemental that must be executed in order if the
     // applying indices are in order. An hlfir::Elemental may have been created
     // for a transformational like transpose, and Fortran 2018 standard
     // section 10.2.3.2, point 10 imply that impure elemental sub-expression
     // evaluations should not be masked if they are the arguments of
     // transformational expressions.
     if (isOutOfOrder && elemental.isOrdered())
       return;
     elementalOps.insert(elemental.getOperation());
-    for (mlir::Operation &op : elemental.getBody()->getOperations())
+    for (mlir::Operation &op : elemental.getElementalRegion().getOps())
       if (auto apply = mlir::dyn_cast<hlfir::ApplyOp>(op)) {
         bool isUnorderedApply =
             isOutOfOrder || !isInOrderApply(apply, elemental);
-        gatherElementalTree(apply.getExpr(), elementalOps, isUnorderedApply);
+        auto maybeElemental =
+            mlir::dyn_cast_or_null<hlfir::ElementalOpInterface>(
+                apply.getExpr().getDefiningOp());
+        gatherElementalTree(maybeElemental, elementalOps, isUnorderedApply);
       }
   }
 }
 
 MaskedArrayExpr::MaskedArrayExpr(mlir::Location loc, mlir::Region &region)
     : loc{loc}, region{region} {
   mlir::Operation &terminator = region.back().back();
-  // TODO: clarify if vector subscripts must be inlined or not here.
-  // In case of x(elemental(A), :), this could lead to more elemental(A)
-  // evaluation than needed, which is not OK if "elemental" is impure.
-  // The standard is not very clear here.
-  if (mlir::isa<hlfir::ElementalAddrOp>(terminator))
-    TODO(loc, "vector subscripted assignments inside WHERE");
+  if (auto elementalAddr =
+          mlir::dyn_cast<hlfir::ElementalOpInterface>(terminator)) {
+    // Vector subscripted designator (hlfir.elemental_addr terminator).
+    gatherElementalTree(elementalAddr, elementalParts, /*isOutOfOrder=*/false);
+    return;
+  }
+  // Try if elemental expression.
   mlir::Value entity = mlir::cast<hlfir::YieldOp>(terminator).getEntity();
-  gatherElementalTree(entity, elementalParts, /*isOutOfOrder=*/false);
+  auto maybeElemental = mlir::dyn_cast_or_null<hlfir::ElementalOpInterface>(
+      entity.getDefiningOp());
+  gatherElementalTree(maybeElemental, elementalParts, /*isOutOfOrder=*/false);
 }
 
 void MaskedArrayExpr::generateNoneElementalPart(fir::FirOpBuilder &builder,
                                                 mlir::IRMapping &mapper) {
   assert(!noneElementalPartWasGenerated &&
          "none elemental parts already generated");
   // Clone all operations, except the elemental and the final yield.
   mlir::Block::OpListType &ops = region.back().getOperations();

 
 mlir::Value
 MaskedArrayExpr::generateElementalParts(fir::FirOpBuilder &builder,
                                         mlir::ValueRange oneBasedIndices,
                                         mlir::IRMapping &mapper) {
   assert(noneElementalPartWasGenerated &&
          "non elemental part must have been generated");
   mlir::Operation &terminator = region.back().back();
-  if (mlir::isa<hlfir::ElementalAddrOp>(terminator))
-    TODO(loc, "vector subscripted assignments inside WHERE");
+  hlfir::ElementalOpInterface elemental =
+      mlir::dyn_cast<hlfir::ElementalAddrOp>(terminator);
+  if (!elemental) {
+    // If the terminator is not an hlfir.elemental_addr, try if the yielded
+    // entity was produced by an hlfir.elemental.
     mlir::Value entity = mlir::cast<hlfir::YieldOp>(terminator).getEntity();
-  auto elemental = entity.getDefiningOp<hlfir::ElementalOp>();
+    elemental = entity.getDefiningOp<hlfir::ElementalOp>();
     if (!elemental) {
+      // The yielded entity was not produced by an elemental operation,
+      // get its clone in the non elemental part evaluation and address it.
       hlfir::Entity clonedEntity{mapper.lookupOrDefault(entity)};
       return hlfir::getElementAt(loc, builder, clonedEntity, oneBasedIndices);
     }
+  }
+
   auto mustRecursivelyInline =
       [&](hlfir::ElementalOp appliedElemental) -> bool {
     return elementalParts.contains(appliedElemental.getOperation());
   };
   return inlineElementalOp(loc, builder, elemental, oneBasedIndices, mapper,
                            mustRecursivelyInline);
 }
 
 void MaskedArrayExpr::generateNoneElementalCleanupIfAny(
     fir::FirOpBuilder &builder, mlir::IRMapping &mapper) {
   mlir::Operation &terminator = region.back().back();
-  if (mlir::isa<hlfir::ElementalAddrOp>(terminator))
-    TODO(loc, "vector subscripted assignments inside WHERE");
+  mlir::Region *cleanupRegion = nullptr;
+  if (auto elementalAddr = mlir::dyn_cast<hlfir::ElementalAddrOp>(terminator)) {
+    cleanupRegion = &elementalAddr.getCleanup();
+  } else {
     auto yieldOp = mlir::cast<hlfir::YieldOp>(terminator);
-  if (yieldOp.getCleanup().empty())
+    cleanupRegion = &yieldOp.getCleanup();
+  }
+  if (cleanupRegion->empty())
     return;
-  for (mlir::Operation &op : yieldOp.getCleanup().getOps()) {
+  for (mlir::Operation &op : cleanupRegion->front().without_terminator()) {
     if (auto destroy = mlir::dyn_cast<hlfir::DestroyOp>(op))
       if (elementalParts.contains(destroy.getExpr().getDefiningOp()))
         continue;
-    if (!mlir::isa<fir::FirEndOp>(op))
-      (void)builder.clone(op, mapper);
   }
 }
 
 static bool isLeftHandSide(mlir::Region &region) {
   auto assign = mlir::dyn_cast<hlfir::RegionAssignOp>(region.getParentOp());
   return assign && (&assign.getLhsRegion() == &region);
 }
 
public:
   explicit RegionAssignConversion(mlir::MLIRContext *ctx)
       : OpRewritePattern{ctx} {}
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::RegionAssignOp regionAssignOp,
                   mlir::PatternRewriter &rewriter) const override {
     auto root = mlir::cast<hlfir::OrderedAssignmentTreeOpInterface>(
         regionAssignOp.getOperation());
-    if (!regionAssignOp.getUserDefinedAssignment().empty())
-      return ::rewrite(root, /*tryFusingAssignments=*/false, rewriter);
-    TODO(regionAssignOp.getLoc(),
-         "assignments to vector subscripted entity in HLFIR");
   }
 };
 
 class LowerHLFIROrderedAssignments
     : public hlfir::impl::LowerHLFIROrderedAssignmentsBase<
           LowerHLFIROrderedAssignments> {
 public:
   void runOnOperation() override {

flang/lib/Optimizer/HLFIR/Transforms/ScheduleOrderedAssignments.cpp

 
 /// Return the entity yielded by a region, or a null value if the region
 /// is not terminated by a yield.
 static mlir::Value getYieldedEntity(mlir::Region &region) {
   if (region.empty() || region.back().empty())
     return nullptr;
   if (auto yield = mlir::dyn_cast<hlfir::YieldOp>(region.back().back()))
     return yield.getEntity();
+  if (auto elementalAddr =
+          mlir::dyn_cast<hlfir::ElementalAddrOp>(region.back().back()))
+    return elementalAddr.getYieldOp().getEntity();
   return nullptr;
 }
 
 /// Gather the effect of an assignment. This is the implicit write to the LHS
 /// of an assignment. This also includes the effects of the user defined
 /// assignment, if any, but this does not include the effects of evaluating the
 /// RHS and LHS, which occur before the assignment effects in Fortran.
 static void gatherAssignEffects(
     hlfir::RegionAssignOp regionAssign,
     bool userDefAssignmentMayOnlyWriteToAssignedVariable,
     llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &assignEffects) {
   mlir::Value assignedVar = getYieldedEntity(regionAssign.getLhsRegion());
-  if (!assignedVar)
-    TODO(regionAssign.getLoc(),
-         "assignment to vector subscripted entity in HLFIR");
+  assert(assignedVar && "lhs cannot be an empty region");
   assignEffects.emplace_back(mlir::MemoryEffects::Write::get(), assignedVar);
 
   if (!regionAssign.getUserDefinedAssignment().empty()) {
     // The write effect on the INTENT(OUT) LHS argument is already taken
     // into account above.
     // This side effects are "defensive" and could be improved.
     // On top of the passed RHS argument, user defined assignments (even when
     // pure) may also read host/used/common variable. Impure user defined

flang/test/HLFIR/order_assignments/vector-subscripts-codegen.fir

+// Test code generation of hlfir.region_assign where the left-hand
+// side terminator is an hlfir.elemental_addr (Fortran assignments to
+// vector subscripted designators).
+// RUN: fir-opt %s --lower-hlfir-ordered-assignments | FileCheck %s
+
+func.func @simple(%arg0: !fir.ref<!fir.array<100xf32>> , %arg1: !fir.ref<!fir.array<10xi64>> , %arg2: !fir.ref<!fir.array<10xf32>> ) {
+  %c10 = arith.constant 10 : index
+  %c100 = arith.constant 100 : index
+  %0 = fir.shape %c100 : (index) -> !fir.shape<1>
+  %1:2 = hlfir.declare %arg0(%0) {uniq_name = "_QFsimpleEx"} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xf32>>, !fir.ref<!fir.array<100xf32>>)
+  %2 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %3:2 = hlfir.declare %arg1(%2) {uniq_name = "y"} : (!fir.ref<!fir.array<10xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi64>>, !fir.ref<!fir.array<10xi64>>)
+  %4:2 = hlfir.declare %arg2(%2) {uniq_name = "z"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
+  hlfir.region_assign {
+    hlfir.yield %4#0 : !fir.ref<!fir.array<10xf32>>
+  } to {
+    hlfir.elemental_addr %2 : !fir.shape<1> {
+    ^bb0(%arg3: index):
+      %5 = hlfir.designate %3#0 (%arg3)  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+      %6 = fir.load %5 : !fir.ref<i64>
+      %7 = hlfir.designate %1#0 (%6)  : (!fir.ref<!fir.array<100xf32>>, i64) -> !fir.ref<f32>
+      hlfir.yield %7 : !fir.ref<f32>
+    }
+  }
+  return
+}
+// CHECK-LABEL:   func.func @simple(
+// CHECK-SAME:                      %[[VAL_0:.*]]: !fir.ref<!fir.array<100xf32>>,
+// CHECK-SAME:                      %[[VAL_1:.*]]: !fir.ref<!fir.array<10xi64>>,
+// CHECK-SAME:                      %[[VAL_2:.*]]: !fir.ref<!fir.array<10xf32>>) {
+// CHECK:           %[[VAL_3:.*]] = arith.constant 10 : index
+// CHECK:           %[[VAL_4:.*]] = arith.constant 100 : index
+// CHECK:           %[[VAL_5:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
+// CHECK:           %[[VAL_6:.*]]:2 = hlfir.declare %[[VAL_0]](%[[VAL_5]]) {uniq_name = "_QFsimpleEx"} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xf32>>, !fir.ref<!fir.array<100xf32>>)
+// CHECK:           %[[VAL_7:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1>
+// CHECK:           %[[VAL_8:.*]]:2 = hlfir.declare %[[VAL_1]](%[[VAL_7]]) {uniq_name = "y"} : (!fir.ref<!fir.array<10xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi64>>, !fir.ref<!fir.array<10xi64>>)
+// CHECK:           %[[VAL_9:.*]]:2 = hlfir.declare %[[VAL_2]](%[[VAL_7]]) {uniq_name = "z"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
+// CHECK:           %[[VAL_10:.*]] = arith.constant 1 : index
+// CHECK:           fir.do_loop %[[VAL_11:.*]] = %[[VAL_10]] to %[[VAL_3]] step %[[VAL_10]] {
+// CHECK:             %[[VAL_12:.*]] = hlfir.designate %[[VAL_8]]#0 (%[[VAL_11]])  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+// CHECK:             %[[VAL_13:.*]] = fir.load %[[VAL_12]] : !fir.ref<i64>
+// CHECK:             %[[VAL_14:.*]] = hlfir.designate %[[VAL_6]]#0 (%[[VAL_13]])  : (!fir.ref<!fir.array<100xf32>>, i64) -> !fir.ref<f32>
+// CHECK:             %[[VAL_15:.*]] = hlfir.designate %[[VAL_9]]#0 (%[[VAL_11]])  : (!fir.ref<!fir.array<10xf32>>, index) -> !fir.ref<f32>
+// CHECK:             hlfir.assign %[[VAL_15]] to %[[VAL_14]] : !fir.ref<f32>, !fir.ref<f32>
+// CHECK:           }
+
+func.func @forall_vector_lhs(%arg0: !fir.ref<!fir.array<100x20xf32>> , %arg1: !fir.ref<!fir.array<10xi64>> ) {
+  %c20_i32 = arith.constant 20 : i32
+  %c1_i32 = arith.constant 1 : i32
+  %c10 = arith.constant 10 : index
+  %c100 = arith.constant 100 : index
+  %c20 = arith.constant 20 : index
+  %0 = fir.shape %c100, %c20 : (index, index) -> !fir.shape<2>
+  %1:2 = hlfir.declare %arg0(%0) {uniq_name = "x"} : (!fir.ref<!fir.array<100x20xf32>>, !fir.shape<2>) -> (!fir.ref<!fir.array<100x20xf32>>, !fir.ref<!fir.array<100x20xf32>>)
+  %2 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %3:2 = hlfir.declare %arg1(%2) {uniq_name = "y"} : (!fir.ref<!fir.array<10xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi64>>, !fir.ref<!fir.array<10xi64>>)
+  hlfir.forall lb {
+    hlfir.yield %c1_i32 : i32
+  } ub {
+    hlfir.yield %c20_i32 : i32
+  }  (%arg2: i32) {
+    hlfir.region_assign {
+      %4 = hlfir.elemental %2 : (!fir.shape<1>) -> !hlfir.expr<10xf32> {
+      ^bb0(%arg3: index):
+        %5 = hlfir.designate %3#0 (%arg3)  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+        %6 = fir.load %5 : !fir.ref<i64>
+        %7 = fir.convert %6 : (i64) -> f32
+        hlfir.yield_element %7 : f32
+      }
+      hlfir.yield %4 : !hlfir.expr<10xf32> cleanup {
+        hlfir.destroy %4 : !hlfir.expr<10xf32>
+      }
+    } to {
+      %4 = fir.convert %arg2 : (i32) -> i64
+      hlfir.elemental_addr %2 : !fir.shape<1> {
+      ^bb0(%arg3: index):
+        %5 = hlfir.designate %3#0 (%arg3)  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+        %6 = fir.load %5 : !fir.ref<i64>
+        %7 = hlfir.designate %1#0 (%6, %4)  : (!fir.ref<!fir.array<100x20xf32>>, i64, i64) -> !fir.ref<f32>
+        hlfir.yield %7 : !fir.ref<f32>
+      }
+    }
+  }
+  return
+}
+// CHECK-LABEL:   func.func @forall_vector_lhs(
+// CHECK-SAME:                                 %[[VAL_0:.*]]: !fir.ref<!fir.array<100x20xf32>>,
+// CHECK-SAME:                                 %[[VAL_1:.*]]: !fir.ref<!fir.array<10xi64>>) {
+// CHECK:           %[[VAL_2:.*]] = arith.constant 20 : i32
+// CHECK:           %[[VAL_3:.*]] = arith.constant 1 : i32
+// CHECK:           %[[VAL_4:.*]] = arith.constant 10 : index
+// CHECK:           %[[VAL_5:.*]] = arith.constant 100 : index
+// CHECK:           %[[VAL_6:.*]] = arith.constant 20 : index
+// CHECK:           %[[VAL_7:.*]] = fir.shape %[[VAL_5]], %[[VAL_6]] : (index, index) -> !fir.shape<2>
+// CHECK:           %[[VAL_8:.*]]:2 = hlfir.declare %[[VAL_0]](%[[VAL_7]]) {uniq_name = "x"} : (!fir.ref<!fir.array<100x20xf32>>, !fir.shape<2>) -> (!fir.ref<!fir.array<100x20xf32>>, !fir.ref<!fir.array<100x20xf32>>)
+// CHECK:           %[[VAL_9:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
+// CHECK:           %[[VAL_10:.*]]:2 = hlfir.declare %[[VAL_1]](%[[VAL_9]]) {uniq_name = "y"} : (!fir.ref<!fir.array<10xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi64>>, !fir.ref<!fir.array<10xi64>>)
+// CHECK:           %[[VAL_11:.*]] = fir.convert %[[VAL_3]] : (i32) -> index
+// CHECK:           %[[VAL_12:.*]] = fir.convert %[[VAL_2]] : (i32) -> index
+// CHECK:           %[[VAL_13:.*]] = arith.constant 1 : index
+// CHECK:           fir.do_loop %[[VAL_14:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_13]] {
+// CHECK:             %[[VAL_15:.*]] = fir.convert %[[VAL_14]] : (index) -> i32
+// CHECK:             %[[VAL_16:.*]] = hlfir.elemental %[[VAL_9]] : (!fir.shape<1>) -> !hlfir.expr<10xf32> {
+// CHECK:             ^bb0(%[[VAL_17:.*]]: index):
+// CHECK:               %[[VAL_18:.*]] = hlfir.designate %[[VAL_10]]#0 (%[[VAL_17]])  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+// CHECK:               %[[VAL_19:.*]] = fir.load %[[VAL_18]] : !fir.ref<i64>
+// CHECK:               %[[VAL_20:.*]] = fir.convert %[[VAL_19]] : (i64) -> f32
+// CHECK:               hlfir.yield_element %[[VAL_20]] : f32
+// CHECK:             }
+// CHECK:             %[[VAL_21:.*]] = fir.convert %[[VAL_15]] : (i32) -> i64
+// CHECK:             %[[VAL_22:.*]] = arith.constant 1 : index
+// CHECK:             fir.do_loop %[[VAL_23:.*]] = %[[VAL_22]] to %[[VAL_4]] step %[[VAL_22]] {
+// CHECK:               %[[VAL_24:.*]] = hlfir.designate %[[VAL_10]]#0 (%[[VAL_23]])  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+// CHECK:               %[[VAL_25:.*]] = fir.load %[[VAL_24]] : !fir.ref<i64>
+// CHECK:               %[[VAL_26:.*]] = hlfir.designate %[[VAL_8]]#0 (%[[VAL_25]], %[[VAL_21]])  : (!fir.ref<!fir.array<100x20xf32>>, i64, i64) -> !fir.ref<f32>
+// CHECK:               %[[VAL_27:.*]] = hlfir.apply %[[VAL_28:.*]], %[[VAL_23]] : (!hlfir.expr<10xf32>, index) -> f32
+// CHECK:               hlfir.assign %[[VAL_27]] to %[[VAL_26]] : f32, !fir.ref<f32>
+// CHECK:             }
+// CHECK:             hlfir.destroy %[[VAL_16]] : !hlfir.expr<10xf32>
+// CHECK:           }
+
+func.func @where_vector_subscripts(%arg0: !fir.ref<!fir.array<10x!fir.logical<4>>> , %arg1: !fir.ref<!fir.array<100xf32>> , %arg2: !fir.ref<!fir.array<10xi64>> ) {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c100 = arith.constant 100 : index
+  %c10 = arith.constant 10 : index
+  %0 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %1:2 = hlfir.declare %arg0(%0) {uniq_name = "l"} : (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.ref<!fir.array<10x!fir.logical<4>>>)
+  %2 = fir.shape %c100 : (index) -> !fir.shape<1>
+  %3:2 = hlfir.declare %arg1(%2) {uniq_name = "x"} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xf32>>, !fir.ref<!fir.array<100xf32>>)
+  %4:2 = hlfir.declare %arg2(%0) {uniq_name = "y"} : (!fir.ref<!fir.array<10xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi64>>, !fir.ref<!fir.array<10xi64>>)
+  hlfir.where {
+    hlfir.yield %1#0 : !fir.ref<!fir.array<10x!fir.logical<4>>>
+  } do {
+    hlfir.region_assign {
+      hlfir.yield %cst : f32
+    } to {
+      hlfir.elemental_addr %0 : !fir.shape<1> {
+      ^bb0(%arg3: index):
+        %5 = hlfir.designate %4#0 (%arg3)  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+        %6 = fir.load %5 : !fir.ref<i64>
+        %7 = hlfir.designate %3#0 (%6)  : (!fir.ref<!fir.array<100xf32>>, i64) -> !fir.ref<f32>
+        hlfir.yield %7 : !fir.ref<f32>
+      }
+    }
+  }
+  return
+}
+// CHECK-LABEL:   func.func @where_vector_subscripts(
+// CHECK-SAME:                                       %[[VAL_0:.*]]: !fir.ref<!fir.array<10x!fir.logical<4>>>,
+// CHECK-SAME:                                       %[[VAL_1:.*]]: !fir.ref<!fir.array<100xf32>>,
+// CHECK-SAME:                                       %[[VAL_2:.*]]: !fir.ref<!fir.array<10xi64>>) {
+// CHECK:           %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_4:.*]] = arith.constant 100 : index
+// CHECK:           %[[VAL_5:.*]] = arith.constant 10 : index
+// CHECK:           %[[VAL_6:.*]] = fir.shape %[[VAL_5]] : (index) -> !fir.shape<1>
+// CHECK:           %[[VAL_7:.*]]:2 = hlfir.declare %[[VAL_0]](%[[VAL_6]]) {uniq_name = "l"} : (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.ref<!fir.array<10x!fir.logical<4>>>)
+// CHECK:           %[[VAL_8:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
+// CHECK:           %[[VAL_9:.*]]:2 = hlfir.declare %[[VAL_1]](%[[VAL_8]]) {uniq_name = "x"} : (!fir.ref<!fir.array<100xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xf32>>, !fir.ref<!fir.array<100xf32>>)
+// CHECK:           %[[VAL_10:.*]]:2 = hlfir.declare %[[VAL_2]](%[[VAL_6]]) {uniq_name = "y"} : (!fir.ref<!fir.array<10xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi64>>, !fir.ref<!fir.array<10xi64>>)
+// CHECK:           %[[VAL_11:.*]] = arith.constant 1 : index
+// CHECK:           fir.do_loop %[[VAL_12:.*]] = %[[VAL_11]] to %[[VAL_5]] step %[[VAL_11]] {
+// CHECK:             %[[VAL_13:.*]] = hlfir.designate %[[VAL_7]]#0 (%[[VAL_12]])  : (!fir.ref<!fir.array<10x!fir.logical<4>>>, index) -> !fir.ref<!fir.logical<4>>
+// CHECK:             %[[VAL_14:.*]] = fir.load %[[VAL_13]] : !fir.ref<!fir.logical<4>>
+// CHECK:             %[[VAL_15:.*]] = fir.convert %[[VAL_14]] : (!fir.logical<4>) -> i1
+// CHECK:             fir.if %[[VAL_15]] {
+// CHECK:               %[[VAL_16:.*]] = hlfir.designate %[[VAL_10]]#0 (%[[VAL_12]])  : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+// CHECK:               %[[VAL_17:.*]] = fir.load %[[VAL_16]] : !fir.ref<i64>
+// CHECK:               %[[VAL_18:.*]] = hlfir.designate %[[VAL_9]]#0 (%[[VAL_17]])  : (!fir.ref<!fir.array<100xf32>>, i64) -> !fir.ref<f32>
+// CHECK:               hlfir.assign %[[VAL_3]] to %[[VAL_18]] : f32, !fir.ref<f32>
+// CHECK:             }
+// CHECK:           }