[flang-commits] [flang] [flang][AliasAnalysis] don't crash on load from blockarg (PR #120760)

[Tom Eccles via flang-commits]

https://github.com/tblah created https://github.com/llvm/llvm-project/pull/120760

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.

>From ba7fe5fe9d1b5f0247d0321e9891f507b15725e1 Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Mon, 16 Dec 2024 14:43:03 +0000
Subject: [PATCH] [flang][AliasAnalysis] don't crash on load from blockarg

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.
---
 .../lib/Optimizer/Analysis/AliasAnalysis.cpp  |  5 ++-
 flang/test/HLFIR/opt-variable-assign-omp.fir  | 43 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)
 create mode 100755 flang/test/HLFIR/opt-variable-assign-omp.fir

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
+}