[flang-commits] [flang] [Flang] - Add optional inlining of allocatable assignments with hlfir.expr RHS (PR #186880)

[via flang-commits]

================
@@ -0,0 +1,97 @@
+// Test inlining of hlfir.assign for allocatable LHS with hlfir.expr RHS.
+// This tests the -inline-hlfir-allocatable-expr-assign flag.
+
+// RUN: fir-opt --inline-hlfir-assign %s | FileCheck %s --check-prefix=DEFAULT
+// RUN: fir-opt -inline-hlfir-allocatable-expr-assign --inline-hlfir-assign %s | FileCheck %s --check-prefix=ENABLED
+
+// Test case: c = cos(a) where c is allocatable
+// This is derived from the flang-529628 test case.
+// The hlfir.elemental produces an hlfir.expr which doesn't alias with
+// memory, so the assignment can be safely inlined.
+
+func.func @test_allocatable_elemental_assign(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, %arg1: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  // Declare the allocatables
+  %a:2 = hlfir.declare %arg0 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
+  %c:2 = hlfir.declare %arg1 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFEc"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
+
+  // Load a to get its shape
+  %a_box = fir.load %a#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+  %dims:3 = fir.box_dims %a_box, %c0 : (!fir.box<!fir.heap<!fir.array<?xf64>>>, index) -> (index, index, index)
+  %shape = fir.shape %dims#1 : (index) -> !fir.shape<1>
+
+  // Create elemental: cos(a)
+  %elemental = hlfir.elemental %shape unordered : (!fir.shape<1>) -> !hlfir.expr<?xf64> {
+  ^bb0(%i: index):
+    %lb_offset = arith.subi %dims#0, %c1 : index
+    %idx = arith.addi %i, %lb_offset : index
+    %a_elem = hlfir.designate %a_box (%idx) : (!fir.box<!fir.heap<!fir.array<?xf64>>>, index) -> !fir.ref<f64>
+    %a_val = fir.load %a_elem : !fir.ref<f64>
+    %cos_val = math.cos %a_val fastmath<contract> : f64
+    hlfir.yield_element %cos_val : f64
+  }
+
+  // Assign elemental result to allocatable c
+  hlfir.assign %elemental to %c#0 realloc : !hlfir.expr<?xf64>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+  hlfir.destroy %elemental : !hlfir.expr<?xf64>
+  return
+}
+
+// DEFAULT-LABEL: func.func @test_allocatable_elemental_assign
+// By default (without the option), the allocatable assign should NOT be inlined
+// DEFAULT: hlfir.assign %{{.*}} to %{{.*}} realloc : !hlfir.expr<?xf64>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+
+// ENABLED-LABEL: func.func @test_allocatable_elemental_assign
+// With the option enabled, the assign should be inlined.
+// The generated code uses genReallocIfNeeded which creates fir.if for
+// allocation checking and a loop for assignment.
+// ENABLED-NOT: hlfir.assign %{{.*}} realloc
+// ENABLED: fir.if
+// ENABLED: fir.do_loop
+// ENABLED: hlfir.apply
+// ENABLED: hlfir.designate
+// ENABLED: hlfir.assign %{{.*}} : f64, !fir.ref<f64>
----------------
jeanPerier wrote:

genReallocIfNeeded/finalizeRealloc will soon or has already lost its test coverage as it is only used in the "legacy" lowering (without the HLFIR step). So I think it is important to test a few things that are hard to correctly test in end-to-end testing, like the fact that the freemem of the original storage happens after the assignment loop to cover the cases where the storage is used inside the expression.

All that to say, please add an `ENABLED: fir.if` and `ENABLED: fir.freemem` after the loops.

https://github.com/llvm/llvm-project/pull/186880