[flang-commits] [flang] ef4fb18 - [flang][hlfir] Resolve shape_of users when bufferizing eval_in_mem (#201214)

[via flang-commits]

Author: khaki3
Date: 2026-06-03T06:51:20-07:00
New Revision: ef4fb183a8d5c0d1910f7b5b6776ba669d6795ef

URL: https://github.com/llvm/llvm-project/commit/ef4fb183a8d5c0d1910f7b5b6776ba669d6795ef
DIFF: https://github.com/llvm/llvm-project/commit/ef4fb183a8d5c0d1910f7b5b6776ba669d6795ef.diff

LOG: [flang][hlfir] Resolve shape_of users when bufferizing eval_in_mem (#201214)

A follow-up to #197814.

Example:

```fortran
bmat = matmul(mat, mat)   ! bmat is allocatable
```

In this code, `SeparateAllocatableAssign` sizes the reallocation with an
`hlfir.shape_of` of the RHS. Once the `matmul` is lowered to
`hlfir.eval_in_mem`, that `shape_of` is an extra user, so
`EvaluateIntoMemoryAssignBufferization` erases the `eval_in_mem` while
it's still used, hitting a `use-after-erase` assertion at `-O2`.

Fix: in `OptimizedBufferization`, redirect a `shape_of` user to the
`eval_in_mem`'s shape operand before erasing it.

Added: 
    

Modified: 
    flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
    flang/test/HLFIR/opt-bufferization-eval_in_mem.fir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
index 51af673406b4a..d717b39479380 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/OptimizedBufferization.cpp
@@ -638,16 +638,30 @@ tryUsingAssignLhsDirectly(hlfir::EvaluateInMemoryOp evalInMem,
   mlir::Location loc = evalInMem.getLoc();
   hlfir::DestroyOp destroy;
   hlfir::AssignOp assign;
-  for (auto user : llvm::enumerate(evalInMem->getUsers())) {
-    if (user.index() > 2)
+  // To evaluate the hlfir.eval_in_mem directly into the LHS, its result must
+  // only be used in the assignment, in a destroy, and in hlfir.shape_of (which
+  // can be replaced by a direct use of the shape operand).
+  llvm::SmallVector<hlfir::ShapeOfOp> shapeOfs;
+  for (mlir::Operation *user : evalInMem->getUsers()) {
+    if (auto op = mlir::dyn_cast<hlfir::AssignOp>(user)) {
+      if (assign)
+        return mlir::failure();
+      assign = op;
+    } else if (auto op = mlir::dyn_cast<hlfir::DestroyOp>(user)) {
+      if (destroy)
+        return mlir::failure();
+      destroy = op;
+    } else if (auto op = mlir::dyn_cast<hlfir::ShapeOfOp>(user)) {
+      shapeOfs.push_back(op);
+    } else {
       return mlir::failure();
-    mlir::TypeSwitch<mlir::Operation *, void>(user.value())
-        .Case([&](hlfir::AssignOp op) { assign = op; })
-        .Case([&](hlfir::DestroyOp op) { destroy = op; });
+    }
   }
   if (!assign || !destroy || destroy.mustFinalizeExpr() ||
       assign.isAllocatableAssignment())
     return mlir::failure();
+  if (!shapeOfs.empty() && !evalInMem.getShape())
+    return mlir::failure();
 
   hlfir::Entity lhs{assign.getLhs()};
   // EvaluateInMemoryOp memory is contiguous, so in general, it can only be
@@ -690,6 +704,10 @@ tryUsingAssignLhsDirectly(hlfir::EvaluateInMemoryOp evalInMem,
   fir::FirOpBuilder builder(rewriter, evalInMem.getOperation());
   mlir::Value rawLhs = hlfir::genVariableRawAddress(loc, builder, lhs);
   hlfir::computeEvaluateOpIn(loc, builder, evalInMem, rawLhs);
+  // Redirect shape_of users to the shape operand so the eval_in_mem can be
+  // erased without leaving dangling uses.
+  for (hlfir::ShapeOfOp shapeOf : shapeOfs)
+    rewriter.replaceOp(shapeOf, evalInMem.getShape());
   rewriter.eraseOp(assign);
   rewriter.eraseOp(destroy);
   rewriter.eraseOp(evalInMem);

diff  --git a/flang/test/HLFIR/opt-bufferization-eval_in_mem.fir b/flang/test/HLFIR/opt-bufferization-eval_in_mem.fir
index fdfb2ce6979dc..9ef8115284514 100644
--- a/flang/test/HLFIR/opt-bufferization-eval_in_mem.fir
+++ b/flang/test/HLFIR/opt-bufferization-eval_in_mem.fir
@@ -60,3 +60,33 @@ func.func @_QPnegative_test_is_target(%arg0: !fir.ref<!fir.array<10xf32>> {fir.b
 // CHECK:         hlfir.destroy %[[VAL_10]] : !hlfir.expr<10xf32>
 // CHECK:         return
 // CHECK:       }
+
+// A hlfir.shape_of user of the eval_in_mem (such as the one left behind by
+// SeparateAllocatableAssign when sizing a reallocation) must not block the
+// in-place rewrite: it is redirected to the eval_in_mem shape operand.
+func.func @_QPtest_shape_of_user(%arg0: !fir.ref<!fir.array<10xf32>> {fir.bindc_name = "x"}, %arg1: !fir.ref<index>) {
+  %c10 = arith.constant 10 : index
+  %0 = fir.dummy_scope : !fir.dscope
+  %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %2:2 = hlfir.declare %arg0(%1) dummy_scope %0 {uniq_name = "_QFtestEx"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>, !fir.dscope) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
+  %3 = hlfir.eval_in_mem shape %1 : (!fir.shape<1>) -> !hlfir.expr<10xf32> {
+  ^bb0(%arg2: !fir.ref<!fir.array<10xf32>>):
+    %4 = fir.call @_QPfoo() fastmath<contract> : () -> !fir.array<10xf32>
+    fir.save_result %4 to %arg2(%1) : !fir.array<10xf32>, !fir.ref<!fir.array<10xf32>>, !fir.shape<1>
+  }
+  %shp = hlfir.shape_of %3 : (!hlfir.expr<10xf32>) -> !fir.shape<1>
+  %ext = hlfir.get_extent %shp {dim = 0 : index} : (!fir.shape<1>) -> index
+  hlfir.assign %3 to %2#0 : !hlfir.expr<10xf32>, !fir.ref<!fir.array<10xf32>>
+  hlfir.destroy %3 : !hlfir.expr<10xf32>
+  fir.store %ext to %arg1 : !fir.ref<index>
+  return
+}
+// CHECK-LABEL: func.func @_QPtest_shape_of_user(
+// CHECK:         %[[SHAPE:.*]] = fir.shape %{{.*}} : (index) -> !fir.shape<1>
+// The shape_of is redirected to the eval_in_mem shape operand...
+// CHECK:         hlfir.get_extent %[[SHAPE]]
+// ...and the result is still evaluated directly into the LHS (no temporary).
+// CHECK:         %[[CALL:.*]] = fir.call @_QPfoo() fastmath<contract> : () -> !fir.array<10xf32>
+// CHECK:         fir.save_result %[[CALL]] to %{{.*}}#0(%[[SHAPE]])
+// CHECK-NOT:     hlfir.eval_in_mem
+// CHECK-NOT:     hlfir.shape_of