[flang-commits] [flang] [flang] Allow `VariableAssignBufferization` to handle `hlfir::ExprType` (PR #115136)

[Kareem Ergawy via flang-commits]

https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/115136

>From 37f581f36511dcb3a32e8c16847e0589bf1fd187 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Wed, 6 Nov 2024 02:41:35 -0600
Subject: [PATCH] [flang] Allow `VariableAssignBufferization` to handle
 `hlfir::ExprType`

Assume you are given the following input:
```fortran
1. subroutine ComputeDifferencesKernel
2.   implicit none
3.   integer :: i
4.   integer, dimension(1) :: a
5.
6.   do i = 1, 10
7.     a = [ i ]
8.   end do
9. end subroutine ComputeDifferencesKernel
```

Currently, the assignment in line 7 ends up as a call to the Fortran
runtime function: `AAssign` since the corresponding `hlfir.assign` is
not optimized away by `VariableAssignBufferization`. The reason this
assignment is not optimized away is that `VariableAssignBufferization`
does not match whenever the RHS of the assignment is a `hlfir.expr`
value. However, this behavior is introduced only to prevent clashes
between `VariableAssignBufferization` and `ElementalAssignBufferization`
which optimizes away assignemnts that result from `hlfir.elemental` ops.

This patch relaxes that restriction by checking whether the RHS of an
`hlfir.assign` is the result of `hlfir.elemental` or not. If not, we can
safely proceed with `VariableAssignBufferization`.

Note that in the above example, we won't get a `hlfir.elemental` in the
IR. We would get if we changed line 7 to something like:
```fortran
7.     a = [ i ] + b
```
---
 .../Transforms/OptimizedBufferization.cpp     | 13 +++---
 flang/test/HLFIR/opt-bufferization.fir        | 42 +++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
index 7553e05b470634..d05a3258cf293c 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
@@ -598,11 +598,14 @@ llvm::LogicalResult VariableAssignBufferization::matchAndRewrite(
     return rewriter.notifyMatchFailure(assign, "AssignOp may imply allocation");
 
   hlfir::Entity rhs{assign.getRhs()};
-  // TODO: ExprType check is here to avoid conflicts with
-  // ElementalAssignBufferization pattern. We need to combine
-  // these matchers into a single one that applies to AssignOp.
-  if (mlir::isa<hlfir::ExprType>(rhs.getType()))
-    return rewriter.notifyMatchFailure(assign, "RHS is not in memory");
+
+  // To avoid conflicts with ElementalAssignBufferization pattern, we avoid
+  // matching RHS when it is an `ExprType` defined by an `ElementalOp`; which is
+  // among the main criteria matched by ElementalAssignBufferization.
+  if (mlir::isa<hlfir::ExprType>(rhs.getType()) &&
+      mlir::isa<hlfir::ElementalOp>(rhs.getDefiningOp()))
+    return rewriter.notifyMatchFailure(
+        assign, "RHS is an ExprType defined by ElementalOp");
 
   if (!rhs.isArray())
     return rewriter.notifyMatchFailure(assign,
diff --git a/flang/test/HLFIR/opt-bufferization.fir b/flang/test/HLFIR/opt-bufferization.fir
index faa8f4bcdb7789..87afb3cc924533 100644
--- a/flang/test/HLFIR/opt-bufferization.fir
+++ b/flang/test/HLFIR/opt-bufferization.fir
@@ -796,3 +796,45 @@ func.func @_QPddx(%arg0: !fir.box<!fir.array<?x?xf64>> {fir.bindc_name = "array"
 // CHECK:           %[[VAL_61:.*]] = fir.load %[[VAL_26]]#1 : !fir.ref<!fir.array<?x?xf64>>
 // CHECK:           return %[[VAL_61]] : !fir.array<?x?xf64>
 // CHECK:         }
+
+// `hlfir.expr` bufferization (when the expresion is not the result of
+// `hlfir.elemental`)
+func.func @_QPfoo() {
+  %c1 = arith.constant 1 : index
+  %0 = fir.alloca !fir.array<1xi32> {bindc_name = "iavs", uniq_name = "_QFfooEiavs"}
+  %1 = fir.shape %c1 : (index) -> !fir.shape<1>
+  %2:2 = hlfir.declare %0(%1) {uniq_name = "_QFfooEiavs"} : (!fir.ref<!fir.array<1xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<1xi32>>, !fir.ref<!fir.array<1xi32>>)
+  %3 = fir.alloca i32 {bindc_name = "iv", uniq_name = "_QFfooEiv"}
+  %4:2 = hlfir.declare %3 {uniq_name = "_QFfooEiv"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+  %c10_i32 = arith.constant 10 : i32
+  %6 = fir.convert %c10_i32 : (i32) -> index
+  %7 = fir.convert %c1 : (index) -> i32
+  %8:2 = fir.do_loop %arg0 = %c1 to %6 step %c1 iter_args(%arg1 = %7) -> (index, i32) {
+    fir.store %arg1 to %4#1 : !fir.ref<i32>
+    %9 = fir.allocmem !fir.array<1xi32> {bindc_name = ".tmp.arrayctor", uniq_name = ""}
+    %10 = fir.shape %c1 : (index) -> !fir.shape<1>
+    %11:2 = hlfir.declare %9(%10) {uniq_name = ".tmp.arrayctor"} : (!fir.heap<!fir.array<1xi32>>, !fir.shape<1>) -> (!fir.heap<!fir.array<1xi32>>, !fir.heap<!fir.array<1xi32>>)
+    %12 = fir.load %4#0 : !fir.ref<i32>
+    %13 = hlfir.designate %11#0 (%c1)  : (!fir.heap<!fir.array<1xi32>>, index) -> !fir.ref<i32>
+    hlfir.assign %12 to %13 : i32, !fir.ref<i32>
+    %true = arith.constant true
+    %14 = hlfir.as_expr %11#0 move %true : (!fir.heap<!fir.array<1xi32>>, i1) -> !hlfir.expr<1xi32>
+    hlfir.assign %14 to %2#0 : !hlfir.expr<1xi32>, !fir.ref<!fir.array<1xi32>>
+    hlfir.destroy %14 : !hlfir.expr<1xi32>
+    %15 = arith.addi %arg0, %c1 : index
+    %16 = fir.convert %c1 : (index) -> i32
+    %17 = fir.load %4#1 : !fir.ref<i32>
+    %18 = arith.addi %17, %16 : i32
+    fir.result %15, %18 : index, i32
+  }
+  fir.store %8#1 to %4#1 : !fir.ref<i32>
+  return
+}
+
+// CHECK-LABEL:   func.func @_QPfoo
+// CHECK:           %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           fir.do_loop {{.*}} {
+// CHECK-NOT:         hlfir.assign %{{.*}} to %{{.*}}#0 : !hlfir.expr<1xi32>, !fir.ref<!fir.array<1xi32>>
+// CHECK:             fir.do_loop %{{.*}} = %[[C1]] to %[[C1]] step %[[C1]] unordered {
+// CHECK:             }
+// CHECK:           }