[flang] Allow VariableAssignBufferization to handle hlfir::ExprType

[ergawy]

Given the following input:

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:

7.     a = [ i ] + b

In which case, ElementalAssignBufferization will kick in instead.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Kareem Ergawy (ergawy)

Changes

Assume you are given the following input:

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:

7.     a = [ i ] + b

---

Full diff: https://github.com/llvm/llvm-project/pull/115136.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp (+9-5)
• (modified) flang/test/HLFIR/opt-bufferization.fir (+42)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
index 7553e05b470634..c235e07596224f 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
@@ -594,15 +594,19 @@ class VariableAssignBufferization
 
 llvm::LogicalResult VariableAssignBufferization::matchAndRewrite(
     hlfir::AssignOp assign, mlir::PatternRewriter &rewriter) const {
+
   if (assign.isAllocatableAssignment())
     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..57ecc448e173a9 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:           }

[tblah]

This seems like a good solution to me. Please wait for approval from @vzakhari too

[ergawy]

Thanks for the review @tblah. Ping for other review 🔔! Anyone objecting to merging this PR (maybe @vzakhari)?

[ergawy]

One more ping 🔔! Please let me know if there are any objections to merging the PR.