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

[Eugene Epshteyn via flang-commits]

================
@@ -0,0 +1,101 @@
+// 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: hlfir.elemental
+// ENABLED: %[[ORIG_MEM:.*]] = fir.box_addr %6 : (!fir.box<!fir.heap<!fir.array<?xf64>>>) -> !fir.heap<!fir.array<?xf64>>
----------------
eugeneepshteyn wrote:

Recommendation: replace `%6` with `%{{.*}}`

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