[flang][AliasAnalysis] don't crash on load from blockarg

[tblah]

Values can have no defining operation when the value is a blockarg. I wrote this as a test for hlfir bufferization rather than alias analysis because I couldn't find a way to add the test.ptr attribute to a block argument.

[llvmbot]

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

Author: Tom Eccles (tblah)

Changes

Values can have no defining operation when the value is a blockarg. I wrote this as a test for hlfir bufferization rather than alias analysis because I couldn't find a way to add the test.ptr attribute to a block argument.

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/Analysis/AliasAnalysis.cpp (+3-2)
• (added) flang/test/HLFIR/opt-variable-assign-omp.fir (+43)

diff --git a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
index 0b0f83d024ce33..611f212269fb7c 100644
--- a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
+++ b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
@@ -599,8 +599,9 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
           // If load is inside target and it points to mapped item,
           // continue tracking.
           Operation *loadMemrefOp = op.getMemref().getDefiningOp();
-          bool isDeclareOp = llvm::isa<fir::DeclareOp>(loadMemrefOp) ||
-                             llvm::isa<hlfir::DeclareOp>(loadMemrefOp);
+          bool isDeclareOp =
+              llvm::isa_and_present<fir::DeclareOp>(loadMemrefOp) ||
+              llvm::isa_and_present<hlfir::DeclareOp>(loadMemrefOp);
           if (isDeclareOp &&
               llvm::isa<omp::TargetOp>(loadMemrefOp->getParentOp())) {
             v = op.getMemref();
diff --git a/flang/test/HLFIR/opt-variable-assign-omp.fir b/flang/test/HLFIR/opt-variable-assign-omp.fir
new file mode 100755
index 00000000000000..10cb2b4408fb86
--- /dev/null
+++ b/flang/test/HLFIR/opt-variable-assign-omp.fir
@@ -0,0 +1,43 @@
+// RUN: fir-opt %s --opt-bufferization | FileCheck %s
+
+// Test that alias analysis doesn't crash determining if the arguments to
+// hlfir.assign alias.
+// CHECK: omp.private {type = firstprivate} @_QFFbEl_firstprivate_box_Uxi32
+
+// TODO: we can't currently optimize this assign because alias analysis doesn't
+// know that the block arguments of the copy region cannot alias.
+
+omp.private {type = firstprivate} @_QFFbEl_firstprivate_box_Uxi32 : !fir.ref<!fir.box<!fir.array<?xi32>>> alloc {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<?xi32>>>):
+  %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.array<?xi32>>>
+  %c0 = arith.constant 0 : index
+  %1:3 = fir.box_dims %0, %c0 : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
+  %2 = fir.shape %1#1 : (index) -> !fir.shape<1>
+  %3 = fir.allocmem !fir.array<?xi32>, %1#1 {bindc_name = ".tmp", uniq_name = ""}
+  %true = arith.constant true
+  %4:2 = hlfir.declare %3(%2) {uniq_name = ".tmp"} : (!fir.heap<!fir.array<?xi32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.heap<!fir.array<?xi32>>)
+  %c0_0 = arith.constant 0 : index
+  %5:3 = fir.box_dims %0, %c0_0 : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
+  %6 = fir.shape_shift %5#0, %5#1 : (index, index) -> !fir.shapeshift<1>
+  %7 = fir.rebox %4#0(%6) : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<?xi32>>
+  %8 = fir.alloca !fir.box<!fir.array<?xi32>>
+  fir.store %7 to %8 : !fir.ref<!fir.box<!fir.array<?xi32>>>
+  omp.yield(%8 : !fir.ref<!fir.box<!fir.array<?xi32>>>)
+} copy {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<?xi32>>>, %arg1 : !fir.ref<!fir.box<!fir.array<?xi32>>>):
+  %0 = fir.load %arg0 {test.ptr = "load_from_block_arg"} : !fir.ref<!fir.box<!fir.array<?xi32>>>
+  hlfir.assign %0 to %arg1 : !fir.box<!fir.array<?xi32>>, !fir.ref<!fir.box<!fir.array<?xi32>>>
+  omp.yield(%arg1 : !fir.ref<!fir.box<!fir.array<?xi32>>>)
+} dealloc {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<?xi32>>>):
+  %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.array<?xi32>>>
+  %1 = fir.box_addr %0 : (!fir.box<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+  %2 = fir.convert %1 : (!fir.ref<!fir.array<?xi32>>) -> i64
+  %c0_i64 = arith.constant 0 : i64
+  %3 = arith.cmpi ne, %2, %c0_i64 : i64
+  fir.if %3 {
+    %4 = fir.convert %1 : (!fir.ref<!fir.array<?xi32>>) -> !fir.heap<!fir.array<?xi32>>
+    fir.freemem %4 : !fir.heap<!fir.array<?xi32>>
+  }
+  omp.yield
+}

[Renaud-K]

LGTM