[flang][hlfir] Lower ordered elemental subroutine calls.

[vzakhari]

Summary

This patch sets unordered fir.do_loop attribute during lowering
of elemental subroutine calls to HLFIR, when it is safe to do so.
Proper handling of hlfir.elemental will be done in a separate patch.

[jeanPerier]

Looks good, thanks!

[tblah]

LGTM, thanks!

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

 /// Structure to describe a loop nest.
 struct LoopNest {
   fir::DoLoopOp outerLoop;
   fir::DoLoopOp innerLoop;
   llvm::SmallVector<mlir::Value> oneBasedIndices;
 };
 
 /// Generate a fir.do_loop nest looping from 1 to extents[i].
+/// \p isUnordered specifies whether the loops in the loop nest
+/// are unordered.
 LoopNest genLoopNest(mlir::Location loc, fir::FirOpBuilder &builder,
-                     mlir::ValueRange extents);
+                     mlir::ValueRange extents, bool isUnordered = false);
 inline LoopNest genLoopNest(mlir::Location loc, fir::FirOpBuilder &builder,
-                            mlir::Value shape) {
-  return genLoopNest(loc, builder, getIndexExtents(loc, builder, shape));
+                            mlir::Value shape, bool isUnordered = false) {
+  return genLoopNest(loc, builder, getIndexExtents(loc, builder, shape),
+                     isUnordered);
 }
 
 /// Inline the body of an hlfir.elemental at the current insertion point
 /// given a list of one based indices. This generates the computation
 /// 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,

flang/lib/Lower/ConvertCall.cpp

if (!shape && optionalWithShape) {
           hlfir::genShape(loc, builder, optionalWithShape->getOriginalActual());
       // TODO: There is an opportunity to add a runtime check here that
       // this array is present as required. Also, the optionality of all actual
       // could be checked and reset given the Fortran requirement.
       optionalWithShape->resetOptionalAspect();
     }
     assert(shape &&
            "elemental array calls must have at least one array arguments");
-    if (mustBeOrdered)
-      TODO(loc, "ordered elemental calls in HLFIR");
     // Push a new local scope so that any temps made inside the elemental
     // iterations are cleaned up inside the iterations.
     if (!callContext.resultType) {
       // Subroutine case. Generate call inside loop nest.
-      hlfir::LoopNest loopNest = hlfir::genLoopNest(loc, builder, shape);
+      hlfir::LoopNest loopNest =
+          hlfir::genLoopNest(loc, builder, shape, !mustBeOrdered);
       mlir::ValueRange oneBasedIndices = loopNest.oneBasedIndices;
       auto insPt = builder.saveInsertionPoint();
       builder.setInsertionPointToStart(loopNest.innerLoop.getBody());
       callContext.stmtCtx.pushScope();
       for (auto &preparedActual : loweredActuals)
         if (preparedActual)
           preparedActual->setElementalIndices(oneBasedIndices);
       impl().genElementalKernel(loweredActuals, callContext);
       callContext.stmtCtx.finalizeAndPop();
       builder.restoreInsertionPoint(insPt);
       return std::nullopt;
     }
     // Function case: generate call inside hlfir.elemental
+    if (mustBeOrdered)
+      TODO(loc, "ordered elemental calls in HLFIR");
     mlir::Type elementType =
         hlfir::getFortranElementType(*callContext.resultType);
     // Get result length parameters.
     llvm::SmallVector<mlir::Value> typeParams;
     if (elementType.isa<fir::CharacterType>() ||
         fir::isRecordWithTypeParameters(elementType)) {
       auto charType = elementType.dyn_cast<fir::CharacterType>();
       if (charType && charType.hasConstantLen())

flang/lib/Optimizer/Builder/HLFIRTools.cpp

if (auto apply = mlir::dyn_cast<hlfir::ApplyOp>(op))
         }
     (void)builder.clone(op, mapper);
   }
   return mapper.lookupOrDefault(elemental.getElementEntity());
 }
 
 hlfir::LoopNest hlfir::genLoopNest(mlir::Location loc,
                                    fir::FirOpBuilder &builder,
-                                   mlir::ValueRange extents) {
+                                   mlir::ValueRange extents, bool isUnordered) {
   hlfir::LoopNest loopNest;
   assert(!extents.empty() && "must have at least one extent");
   auto insPt = builder.saveInsertionPoint();
   loopNest.oneBasedIndices.assign(extents.size(), mlir::Value{});
   // Build loop nest from column to row.
   auto one = builder.create<mlir::arith::ConstantIndexOp>(loc, 1);
   mlir::Type indexType = builder.getIndexType();
   unsigned dim = extents.size() - 1;
   for (auto extent : llvm::reverse(extents)) {
     auto ub = builder.createConvert(loc, indexType, extent);
-    loopNest.innerLoop = builder.create<fir::DoLoopOp>(loc, one, ub, one);
+    loopNest.innerLoop =
+        builder.create<fir::DoLoopOp>(loc, one, ub, one, isUnordered);
     builder.setInsertionPointToStart(loopNest.innerLoop.getBody());
     // Reverse the indices so they are in column-major order.
     loopNest.oneBasedIndices[dim--] = loopNest.innerLoop.getInductionVar();
     if (!loopNest.outerLoop)
       loopNest.outerLoop = loopNest.innerLoop;
   }
   builder.restoreInsertionPoint(insPt);
   return loopNest;

flang/test/HLFIR/extents-of-shape-of.f90

 
 ! CHECK-HLFIR-DAG:     %[[A_VAR:.*]]:2 = hlfir.declare %[[A_ARG]]
 ! CHECK-HLFIR-DAG:     %[[B_VAR:.*]]:2 = hlfir.declare %[[B_ARG]]
 ! CHECK-HLFIR-NEXT:    %[[MUL:.*]] = hlfir.matmul %[[A_VAR]]#0 %[[B_VAR]]#0
 ! CHECK-HLFIR-NEXT:    %[[SHAPE:.*]] = hlfir.shape_of %[[MUL]] : (!hlfir.expr<2x?xf32>) -> !fir.shape<2>
 ! CHECK-HLFIR-NEXT:    %[[EXT0:.*]] = arith.constant 2 : index
 ! CHECK-HLFIR-NEXT:    %[[EXT1:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 1 : index} : (!fir.shape<2>) -> index
 ! CHECK-HLFIR-NEXT:    %[[C1:.*]] = arith.constant 1 : index
-! CHECK-HLFIR-NEXT:    fir.do_loop %[[ARG2:.*]] = %[[C1]] to %[[EXT1]] step %[[C1]] {
-! CHECK-HLFIR-NEXT:      fir.do_loop %[[ARG3:.*]] = %[[C1]] to %[[EXT0]] step %[[C1]] {
+! CHECK-HLFIR-NEXT:    fir.do_loop %[[ARG2:.*]] = %[[C1]] to %[[EXT1]] step %[[C1]] unordered {
+! CHECK-HLFIR-NEXT:      fir.do_loop %[[ARG3:.*]] = %[[C1]] to %[[EXT0]] step %[[C1]] unordered {
 ! CHECK-HLFIR-NEXT:        %[[ELE:.*]] = hlfir.apply %[[MUL]], %[[ARG3]], %[[ARG2]] : (!hlfir.expr<2x?xf32>, index, index) -> f32
 ! CHECK-HLFIR-NEXT:        %[[ASSOC:.*]]:3 = hlfir.associate %[[ELE]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
 ! CHECK-HLFIR-NEXT:        fir.call
 ! CHECK-HLFIR-NEXT:        hlfir.end_associate
 ! CHECK-HLFIR-NEXT:      }
 ! CHECK-HLFIR-NEXT:    }
 ! CHECK-HLFIR-NEXT:    hlfir.destroy %[[MUL]]
 
 ! ...
 ! CHECK-FIR:           fir.call @_FortranAMatmul
 ! CHECK-FIR-NEXT:      %[[MUL:.*]] = fir.load %[[MUL_BOX:.*]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>
 ! CHECK-FIR-NEXT:      %[[C0:.*]] = arith.constant 0 : index
 ! CHECK-FIR-NEXT:      %[[DIMS0:.*]]:3 = fir.box_dims %[[MUL]], %[[C0]]
 ! CHECK-FIR-NEXT:      %[[C1:.*]] = arith.constant 1 : index
 ! CHECK-FIR-NEXT:      %[[DIMS1:.*]]:3 = fir.box_dims %[[MUL]], %[[C1]]
 ! ...
 ! CHECK-FIR:           %[[SHAPE:.*]] = fir.shape %[[DIMS0]]#1, %[[DIMS1]]#1
 ! CHECK-FIR-NEXT:      %[[C2:.*]] = arith.constant 2 : index
 ! CHECK-FIR-NEXT:      %[[C1_1:.*]] = arith.constant 1 : index
-! CHECK-FIR-NEXT:      fir.do_loop %[[ARG2:.*]] = %[[C1_1]] to %[[DIMS1]]#1 step %[[C1_1]] {
-! CHECK-FIR-NEXT:        fir.do_loop %[[ARG3:.*]] = %[[C1_1]] to %[[C2]] step %[[C1_1]] {
+! CHECK-FIR-NEXT:      fir.do_loop %[[ARG2:.*]] = %[[C1_1]] to %[[DIMS1]]#1 step %[[C1_1]] unordered {
+! CHECK-FIR-NEXT:        fir.do_loop %[[ARG3:.*]] = %[[C1_1]] to %[[C2]] step %[[C1_1]] unordered {
 ! ...
 
 ! CHECK-ALL:           return
 ! CHECK-ALL-NEXT:    }

flang/test/Lower/HLFIR/calls-optional.f90

 ! CHECK:  %[[VAL_5:.*]] = fir.box_addr %[[VAL_4]] : (!fir.box<!fir.ptr<!fir.array<?xf32>>>) -> !fir.ptr<!fir.array<?xf32>>
 ! CHECK:  %[[VAL_6:.*]] = fir.convert %[[VAL_5]] : (!fir.ptr<!fir.array<?xf32>>) -> i64
 ! CHECK:  %[[VAL_7:.*]] = arith.constant 0 : i64
 ! CHECK:  %[[VAL_8:.*]] = arith.cmpi ne, %[[VAL_6]], %[[VAL_7]] : i64
 ! CHECK:  %[[VAL_9:.*]] = arith.constant 0 : index
 ! CHECK:  %[[VAL_10:.*]]:3 = fir.box_dims %[[VAL_2]]#0, %[[VAL_9]] : (!fir.box<!fir.array<?xf32>>, index) -> (index, index, index)
 ! CHECK:  %[[VAL_11:.*]] = fir.load %[[VAL_3]]#0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf32>>>>
 ! CHECK:  %[[VAL_12:.*]] = arith.constant 1 : index
-! CHECK:  fir.do_loop %[[VAL_13:.*]] = %[[VAL_12]] to %[[VAL_10]]#1 step %[[VAL_12]] {
+! CHECK:  fir.do_loop %[[VAL_13:.*]] = %[[VAL_12]] to %[[VAL_10]]#1 step %[[VAL_12]] unordered {
 ! CHECK:    %[[VAL_14:.*]] = hlfir.designate %[[VAL_2]]#0 (%[[VAL_13]])  : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
 ! CHECK:    %[[VAL_15:.*]] = fir.if %[[VAL_8]] -> (!fir.ref<f32>) {
 ! CHECK:      %[[VAL_16:.*]] = arith.constant 0 : index
 ! CHECK:      %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_11]], %[[VAL_16]] : (!fir.box<!fir.ptr<!fir.array<?xf32>>>, index) -> (index, index, index)
 ! CHECK:      %[[VAL_18:.*]] = arith.constant 1 : index
 ! CHECK:      %[[VAL_19:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_18]] : index
 ! CHECK:      %[[VAL_20:.*]] = arith.addi %[[VAL_13]], %[[VAL_19]] : index
 ! CHECK:      %[[VAL_21:.*]] = hlfir.designate %[[VAL_11]] (%[[VAL_20]])  : (!fir.box<!fir.ptr<!fir.array<?xf32>>>, index) -> !fir.ref<f32>
subroutine optional_cannot_be_absent_optional(x)
   ! If all array arguments in an call are optional, they must be all present.
   call elem_takes_one_optional(x)
 end subroutine
 ! CHECK-LABEL: func.func @_QPoptional_cannot_be_absent_optional(
 ! CHECK:  %[[VAL_1:.*]]:2 = hlfir.declare %[[VAL_0:[a-z0-9]*]] {fortran_attrs = #fir.var_attrs<optional>, uniq_name = "_QFoptional_cannot_be_absent_optionalEx"} : (!fir.box<!fir.array<?xf32>>) -> (!fir.box<!fir.array<?xf32>>, !fir.box<!fir.array<?xf32>>)
 ! CHECK:  %[[VAL_2:.*]] = arith.constant 0 : index
 ! CHECK:  %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_1]]#0, %[[VAL_2]] : (!fir.box<!fir.array<?xf32>>, index) -> (index, index, index)
 ! CHECK:  %[[VAL_4:.*]] = arith.constant 1 : index
-! CHECK:  fir.do_loop %[[VAL_5:.*]] = %[[VAL_4]] to %[[VAL_3]]#1 step %[[VAL_4]] {
+! CHECK:  fir.do_loop %[[VAL_5:.*]] = %[[VAL_4]] to %[[VAL_3]]#1 step %[[VAL_4]] unordered {
 ! CHECK:    %[[VAL_6:.*]] = hlfir.designate %[[VAL_1]]#0 (%[[VAL_5]])  : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
 ! CHECK:    fir.call @_QPelem_takes_one_optional(%[[VAL_6]]) {{.*}} : (!fir.ref<f32>) -> ()
 ! CHECK:  }
 
 subroutine optional_elem_poly(x, y)
   interface
     elemental subroutine elem_optional_poly(x, y)
       class(*), optional, intent(in) :: x, y
     end subroutine
   end interface
   real :: x(:)
   real, optional :: y(:)
   call elem_optional_poly(x, y)
 end subroutine
 ! CHECK-LABEL: func.func @_QPoptional_elem_poly(
 ! CHECK:  %[[VAL_2:.*]]:2 = hlfir.declare %[[VAL_0:[a-z0-9]*]] {uniq_name = "_QFoptional_elem_polyEx"} : (!fir.box<!fir.array<?xf32>>) -> (!fir.box<!fir.array<?xf32>>, !fir.box<!fir.array<?xf32>>)
 ! CHECK:  %[[VAL_3:.*]]:2 = hlfir.declare %[[VAL_1:[a-z0-9]*]] {fortran_attrs = #fir.var_attrs<optional>, uniq_name = "_QFoptional_elem_polyEy"} : (!fir.box<!fir.array<?xf32>>) -> (!fir.box<!fir.array<?xf32>>, !fir.box<!fir.array<?xf32>>)
 ! CHECK:  %[[VAL_4:.*]] = fir.is_present %[[VAL_3]]#0 : (!fir.box<!fir.array<?xf32>>) -> i1
 ! CHECK:  %[[VAL_5:.*]] = arith.constant 0 : index
 ! CHECK:  %[[VAL_6:.*]]:3 = fir.box_dims %[[VAL_2]]#0, %[[VAL_5]] : (!fir.box<!fir.array<?xf32>>, index) -> (index, index, index)
 ! CHECK:  %[[VAL_7:.*]] = arith.constant 1 : index
-! CHECK:  fir.do_loop %[[VAL_8:.*]] = %[[VAL_7]] to %[[VAL_6]]#1 step %[[VAL_7]] {
+! CHECK:  fir.do_loop %[[VAL_8:.*]] = %[[VAL_7]] to %[[VAL_6]]#1 step %[[VAL_7]] unordered {
 ! CHECK:    %[[VAL_9:.*]] = hlfir.designate %[[VAL_2]]#0 (%[[VAL_8]])  : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
 ! CHECK:    %[[VAL_10:.*]] = fir.embox %[[VAL_9]] : (!fir.ref<f32>) -> !fir.box<f32>
 ! CHECK:    %[[VAL_11:.*]] = fir.rebox %[[VAL_10]] : (!fir.box<f32>) -> !fir.class<none>
 ! CHECK:    %[[VAL_12:.*]] = fir.if %[[VAL_4]] -> (!fir.class<none>) {
 ! CHECK:      %[[VAL_13:.*]] = hlfir.designate %[[VAL_3]]#0 (%[[VAL_8]])  : (!fir.box<!fir.array<?xf32>>, index) -> !fir.ref<f32>
 ! CHECK:      %[[VAL_14:.*]] = fir.embox %[[VAL_13]] : (!fir.ref<f32>) -> !fir.box<f32>
 ! CHECK:      %[[VAL_15:.*]] = fir.rebox %[[VAL_14]] : (!fir.box<f32>) -> !fir.class<none>
 ! CHECK:      fir.result %[[VAL_15]] : !fir.class<none>

flang/test/Lower/HLFIR/elemental-user-procedure-ref.f90

subroutine sub(x, y)
   call elem_sub(x, y)
 end subroutine
 ! CHECK-LABEL: func.func @_QPsub(
 ! CHECK:  %[[VAL_2:.*]]:2 = hlfir.declare %[[VAL_0:.*]] {{.*}}x
 ! CHECK:  %[[VAL_3:.*]] = arith.constant 10 : index
 ! CHECK:  %[[VAL_4:.*]] = arith.constant 20 : index
 ! CHECK:  %[[VAL_6:.*]]:2 = hlfir.declare %[[VAL_1:.*]](%[[VAL_5:[^)]*]]) {{.*}}y
 ! CHECK:  %[[VAL_7:.*]] = arith.constant 1 : index
-! CHECK:  fir.do_loop %[[VAL_8:.*]] = %[[VAL_7]] to %[[VAL_4]] step %[[VAL_7]] {
-! CHECK:    fir.do_loop %[[VAL_9:.*]] = %[[VAL_7]] to %[[VAL_3]] step %[[VAL_7]] {
+! CHECK:  fir.do_loop %[[VAL_8:.*]] = %[[VAL_7]] to %[[VAL_4]] step %[[VAL_7]] unordered {
+! CHECK:    fir.do_loop %[[VAL_9:.*]] = %[[VAL_7]] to %[[VAL_3]] step %[[VAL_7]] unordered {
 ! CHECK:      %[[VAL_10:.*]] = hlfir.designate %[[VAL_6]]#0 (%[[VAL_9]], %[[VAL_8]])  : (!fir.ref<!fir.array<10x20xf32>>, index, index) -> !fir.ref<f32>
 ! CHECK:      fir.call @_QPelem_sub(%[[VAL_2]]#1, %[[VAL_10]]) fastmath<contract> : (!fir.ref<i32>, !fir.ref<f32>) -> ()
 ! CHECK:    }
 ! CHECK:  }
+
+subroutine impure_elemental(x)
+  real :: x(10, 20)
+  interface
+    impure elemental subroutine impure_elem(a)
+      real, intent(in) :: a
+    end subroutine
+  end interface
+  call impure_elem(x)
+end subroutine
+! CHECK-LABEL:   func.func @_QPimpure_elemental(
+! CHECK-SAME:      %[[VAL_0:.*]]: !fir.ref<!fir.array<10x20xf32>> {fir.bindc_name = "x"}) {
+! CHECK:           %[[VAL_1:.*]] = arith.constant 10 : index
+! CHECK:           %[[VAL_2:.*]] = arith.constant 20 : index
+! CHECK:           %[[VAL_3:.*]] = fir.shape %[[VAL_1]], %[[VAL_2]] : (index, index) -> !fir.shape<2>
+! CHECK:           %[[VAL_4:.*]]:2 = hlfir.declare %[[VAL_0]](%[[VAL_3]]) {uniq_name = "_QFimpure_elementalEx"} : (!fir.ref<!fir.array<10x20xf32>>, !fir.shape<2>) -> (!fir.ref<!fir.array<10x20xf32>>, !fir.ref<!fir.array<10x20xf32>>)
+! CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
+! CHECK:           fir.do_loop %[[VAL_6:.*]] = %[[VAL_5]] to %[[VAL_2]] step %[[VAL_5]] {
+! CHECK:             fir.do_loop %[[VAL_7:.*]] = %[[VAL_5]] to %[[VAL_1]] step %[[VAL_5]] {
+! CHECK:               %[[VAL_8:.*]] = hlfir.designate %[[VAL_4]]#0 (%[[VAL_7]], %[[VAL_6]])  : (!fir.ref<!fir.array<10x20xf32>>, index, index) -> !fir.ref<f32>
+! CHECK:               fir.call @_QPimpure_elem(%[[VAL_8]]) fastmath<contract> : (!fir.ref<f32>) -> ()
+! CHECK:             }
+! CHECK:           }
+! CHECK:           return
+! CHECK:         }
+
+subroutine ordered_elemental(x)
+  real :: x(10, 20)
+  interface
+    elemental subroutine ordered_elem(a)
+      real, intent(inout) :: a
+    end subroutine
+  end interface
+  call ordered_elem(x)
+end subroutine
+! CHECK-LABEL:   func.func @_QPordered_elemental(
+! CHECK-SAME:      %[[VAL_0:.*]]: !fir.ref<!fir.array<10x20xf32>> {fir.bindc_name = "x"}) {
+! CHECK:           %[[VAL_1:.*]] = arith.constant 10 : index
+! CHECK:           %[[VAL_2:.*]] = arith.constant 20 : index
+! CHECK:           %[[VAL_3:.*]] = fir.shape %[[VAL_1]], %[[VAL_2]] : (index, index) -> !fir.shape<2>
+! CHECK:           %[[VAL_4:.*]]:2 = hlfir.declare %[[VAL_0]](%[[VAL_3]]) {uniq_name = "_QFordered_elementalEx"} : (!fir.ref<!fir.array<10x20xf32>>, !fir.shape<2>) -> (!fir.ref<!fir.array<10x20xf32>>, !fir.ref<!fir.array<10x20xf32>>)
+! CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
+! CHECK:           fir.do_loop %[[VAL_6:.*]] = %[[VAL_5]] to %[[VAL_2]] step %[[VAL_5]] {
+! CHECK:             fir.do_loop %[[VAL_7:.*]] = %[[VAL_5]] to %[[VAL_1]] step %[[VAL_5]] {
+! CHECK:               %[[VAL_8:.*]] = hlfir.designate %[[VAL_4]]#0 (%[[VAL_7]], %[[VAL_6]])  : (!fir.ref<!fir.array<10x20xf32>>, index, index) -> !fir.ref<f32>
+! CHECK:               fir.call @_QPordered_elem(%[[VAL_8]]) fastmath<contract> : (!fir.ref<f32>) -> ()
+! CHECK:             }
+! CHECK:           }
+! CHECK:           return
+! CHECK:         }