[flang][hlfir] Codegen for unordered elemental operations.

[vzakhari]

Summary

Depends on D154031, D154032

[jeanPerier]

LGTM

[tblah]

LGTM, thanks

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

matchAndRewrite(hlfir::ElementalOp elemental, OpAdaptor adaptor,
 
     mlir::Value shape = adaptor.getShape();
     auto extents = hlfir::getIndexExtents(loc, builder, shape);
     auto [temp, cleanup] =
         createArrayTemp(loc, builder, elemental.getType(), shape, extents,
                         adaptor.getTypeparams());
     // Generate a loop nest looping around the fir.elemental shape and clone
     // fir.elemental region inside the inner loop.
-    hlfir::LoopNest loopNest = hlfir::genLoopNest(loc, builder, extents);
+    hlfir::LoopNest loopNest =
+        hlfir::genLoopNest(loc, builder, extents, !elemental.isOrdered());
     auto insPt = builder.saveInsertionPoint();
     builder.setInsertionPointToStart(loopNest.innerLoop.getBody());
     auto yield = hlfir::inlineElementalOp(loc, builder, elemental,
                                           loopNest.oneBasedIndices);
     hlfir::Entity elementValue(yield.getElementValue());
     // Skip final AsExpr if any. It would create an element temporary,
     // which is no needed since the element will be assigned right away in
     // the array temporary. An hlfir.as_expr may have been added if the

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

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);
+    loweredLhs.vectorSubscriptLoopNest = hlfir::genLoopNest(
+        loc, builder, newShape, !elementalAddrLhs.isOrdered());
     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 =

flang/test/HLFIR/elemental-codegen.fir

 // CHECK:        %[[VAL_10:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_9]], %[[VAL_8]])  : (!fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>, index, index) -> !fir.box<!fir.type<t{field:f32}>>
 // CHECK:        %[[VAL_11:.*]] = hlfir.designate %[[VAL_5]]#0 (%[[VAL_9]], %[[VAL_8]])  : (!fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>, index, index) -> !fir.ref<!fir.type<t{field:f32}>>
 // CHECK:        hlfir.assign %[[VAL_10]] to %[[VAL_11]] temporary_lhs : !fir.box<!fir.type<t{field:f32}>>, !fir.ref<!fir.type<t{field:f32}>>
 // CHECK:      }
 // CHECK:    }
 // CHECK:    %[[VAL_12:.*]] = fir.undefined tuple<!fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>, i1>
 // CHECK:    %[[VAL_13:.*]] = fir.insert_value %[[VAL_12]], %[[VAL_6]], [1 : index] : (tuple<!fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>, i1>, i1) -> tuple<!fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>, i1>
 // CHECK:    %[[VAL_14:.*]] = fir.insert_value %[[VAL_13]], %[[VAL_5]]#0, [0 : index] : (tuple<!fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>, i1>, !fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>) -> tuple<!fir.box<!fir.array<?x?x!fir.type<t{field:f32}>>>, i1>
+
+func.func @unordered() {
+  %c10 = arith.constant 10 : index
+  %c20 = arith.constant 20 : index
+  %0 = fir.shape %c10, %c20 : (index, index) -> !fir.shape<2>
+  %3 = hlfir.elemental %0 unordered : (!fir.shape<2>) -> !hlfir.expr<10x20xi32> {
+  ^bb0(%i: index, %j: index):
+    %c0 = arith.constant 0 : i32
+    hlfir.yield_element %c0 : i32
+  }
+  return
+}
+// CHECK-LABEL:   func.func @unordered() {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 10 : index
+// CHECK:           %[[VAL_1:.*]] = arith.constant 20 : index
+// CHECK:           %[[VAL_2:.*]] = fir.shape %[[VAL_0]], %[[VAL_1]] : (index, index) -> !fir.shape<2>
+// CHECK:           %[[VAL_3:.*]] = fir.allocmem !fir.array<10x20xi32> {bindc_name = ".tmp.array", uniq_name = ""}
+// CHECK:           %[[VAL_4:.*]]:2 = hlfir.declare %[[VAL_3]](%[[VAL_2]]) {uniq_name = ".tmp.array"} : (!fir.heap<!fir.array<10x20xi32>>, !fir.shape<2>) -> (!fir.heap<!fir.array<10x20xi32>>, !fir.heap<!fir.array<10x20xi32>>)
+// CHECK:           %[[VAL_5:.*]] = arith.constant true
+// CHECK:           %[[VAL_6:.*]] = arith.constant 1 : index
+// CHECK:           fir.do_loop %[[VAL_7:.*]] = %[[VAL_6]] to %[[VAL_1]] step %[[VAL_6]] unordered {
+// CHECK:             fir.do_loop %[[VAL_8:.*]] = %[[VAL_6]] to %[[VAL_0]] step %[[VAL_6]] unordered {
+// CHECK:               %[[VAL_9:.*]] = arith.constant 0 : i32
+// CHECK:               %[[VAL_10:.*]] = hlfir.designate %[[VAL_4]]#0 (%[[VAL_8]], %[[VAL_7]])  : (!fir.heap<!fir.array<10x20xi32>>, index, index) -> !fir.ref<i32>
+// CHECK:               hlfir.assign %[[VAL_9]] to %[[VAL_10]] temporary_lhs : i32, !fir.ref<i32>
+// CHECK:             }
+// CHECK:           }
+// CHECK:           %[[VAL_11:.*]] = fir.undefined tuple<!fir.heap<!fir.array<10x20xi32>>, i1>
+// CHECK:           %[[VAL_12:.*]] = fir.insert_value %[[VAL_11]], %[[VAL_5]], [1 : index] : (tuple<!fir.heap<!fir.array<10x20xi32>>, i1>, i1) -> tuple<!fir.heap<!fir.array<10x20xi32>>, i1>
+// CHECK:           %[[VAL_13:.*]] = fir.insert_value %[[VAL_12]], %[[VAL_4]]#0, [0 : index] : (tuple<!fir.heap<!fir.array<10x20xi32>>, i1>, !fir.heap<!fir.array<10x20xi32>>) -> tuple<!fir.heap<!fir.array<10x20xi32>>, i1>
+// CHECK:           return
+// CHECK:         }

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

 // 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:           }
+
+func.func @unordered(%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 unordered : !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 @unordered(
+// 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]] unordered {
+// 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:           }
+// CHECK:           return
+// CHECK:         }