From 3abcd5f282fd17be679fa1775154672feead4ead Mon Sep 17 00:00:00 2001
From: jeanPerier <jperier@nvidia.com>
Date: Wed, 24 Apr 2024 17:37:40 +0200
Subject: [PATCH] [flang][openmp] fix OMPFunctionFiltering pass after #87796
 (#89776)

The pass assumed that all fun.func symbol usages could be safely
replaced by undef, that is not true after #87796 that added a back link
from internal procedure back to the parent procedure. This caused the
internal procedures to be erased and then processed (segfault).

Also set visibility of such internal procedures so that MLIR do not remove them before
the target function is generated for the target region.
---
 .../Transforms/OMPFunctionFiltering.cpp       | 11 ++++++
 .../Lower/OpenMP/function-filtering-3.f90     | 34 +++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 flang/test/Lower/OpenMP/function-filtering-3.f90

diff --git a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
index 959099d039a5..005e84cb8e9f 100644
--- a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
+++ b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Transforms/Passes.h"
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -66,6 +67,16 @@ public:
         SymbolTable::UseRange funcUses = *funcOp.getSymbolUses(op);
         for (SymbolTable::SymbolUse use : funcUses) {
           Operation *callOp = use.getUser();
+          if (auto internalFunc = mlir::dyn_cast<func::FuncOp>(callOp)) {
+            // Do not delete internal procedures holding the symbol of their
+            // Fortran host procedure as attribute.
+            internalFunc->removeAttr(fir::getHostSymbolAttrName());
+            // Set public visibility so that the function is not deleted by MLIR
+            // because unused. Changing it is OK here because the function will
+            // be deleted anyway in the second filtering phase.
+            internalFunc.setVisibility(mlir::SymbolTable::Visibility::Public);
+            continue;
+          }
           // If the callOp has users then replace them with Undef values.
           if (!callOp->use_empty()) {
             SmallVector<Value> undefResults;
diff --git a/flang/test/Lower/OpenMP/function-filtering-3.f90 b/flang/test/Lower/OpenMP/function-filtering-3.f90
new file mode 100644
index 000000000000..a277c06d6206
--- /dev/null
+++ b/flang/test/Lower/OpenMP/function-filtering-3.f90
@@ -0,0 +1,34 @@
+! RUN: %flang_fc1 -fopenmp -flang-experimental-hlfir -emit-llvm %s -o - | FileCheck --check-prefixes=LLVM-HOST,LLVM-ALL %s
+! RUN: %flang_fc1 -fopenmp -emit-hlfir %s -o - | FileCheck --check-prefixes=MLIR-HOST,MLIR-ALL %s
+! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -flang-experimental-hlfir -emit-llvm %s -o - | FileCheck --check-prefixes=LLVM-DEVICE,LLVM-ALL %s
+! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-hlfir %s -o - | FileCheck --check-prefixes=MLIR-DEVICE,MLIR-ALL %s
+! RUN: bbc -fopenmp -emit-hlfir %s -o - | FileCheck --check-prefixes=MLIR-HOST,MLIR-ALL %s
+! RUN: bbc -fopenmp -fopenmp-is-target-device -emit-hlfir %s -o - | FileCheck --check-prefixes=MLIR-DEVICE,MLIR-ALL %s
+
+! Check that the correct LLVM IR functions are kept for the host and device
+! after running the whole set of translation and transformation passes from
+! Fortran.
+
+! MLIR-HOST: func.func @{{.*}}host_parent_procedure(
+! MLIR-HOST: return
+! MLIR-DEVICE-NOT: func.func {{.*}}host_parent_procedure(
+
+! LLVM-HOST: define {{.*}} @host_parent_procedure{{.*}}(
+! LLVM-DEVICE-NOT: {{.*}} @{{.*}}_host_parent_procedure{{.*}}(
+subroutine host_parent_procedure(x)
+  integer, intent(out) :: x
+  call target_internal_proc(x)
+contains
+! MLIR-ALL: func.func {{.*}}@_QFhost_parent_procedurePtarget_internal_proc(
+
+! LLVM-HOST: define {{.*}} @_QFhost_parent_procedurePtarget_internal_proc(
+! LLVM-DEVICE-NOT: define {{.*}} @_QFhost_parent_procedurePtarget_internal_proc(
+! LLVM-ALL: define {{.*}} @__omp_offloading_{{.*}}QFhost_parent_procedurePtarget_internal_proc{{.*}}(
+
+subroutine target_internal_proc(x)
+  integer, intent(out) :: x
+  !$omp target map(from:x)
+    x = 10
+  !$omp end target
+end subroutine
+end subroutine