[Flang][MLIR][OpenMP] Improve device-only function filtering

This patch improves the implementation of a recent function filtering
workaround to address problems uncovered by D154247.

In particular, the problem was related to the removal of functions called from
within target regions. Since target regions have to remain until LLVM IR is
generated, removing these functions from MLIR results in undefined references
any time there are calls to them in a target region. This patch modifies the
MLIR function filtering pass to make these functions "external" rather than
removing them. This way, the processing and lowering of MLIR functions that
will eventually be discarded is still prevented, but no calls to undefined
functions remain either.

Additionally, the approach of just filtering host-only functions during device
compilation, and not filtering device-only functions during host compilation,
is maintained. This is because code generation for device-only functions is
required for host fallback to work.

Depends on D156988

Differential Revision: https://reviews.llvm.org/D155827

flang/include/flang/Optimizer/Transforms/Passes.td

@@ -319,7 +319,7 @@ def OMPMarkDeclareTargetPass
 
 def OMPFunctionFiltering : Pass<"omp-function-filtering"> {
   let summary = "Filters out functions intended for the host when compiling "
-                "for the device and vice versa.";
+                "for the target device.";
   let constructor = "::fir::createOMPFunctionFilteringPass()";
   let dependentDialects = [
     "mlir::func::FuncDialect"

flang/lib/Frontend/FrontendActions.cpp

@@ -316,10 +316,6 @@ bool CodeGenAction::beginSourceFileAction() {
     pm.addPass(fir::createOMPMarkDeclareTargetPass());
     if (isDevice) {
       pm.addPass(fir::createOMPEarlyOutliningPass());
-      // FIXME: This should eventually be moved out of the
-      // if, so that it also functions for host, however,
-      // we must fix the filtering to function reasonably
-      // for host first.  
       pm.addPass(fir::createOMPFunctionFilteringPass());
     }
   }

flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp

@@ -24,7 +24,6 @@ namespace fir {
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 } // namespace fir
 
-using namespace fir;
 using namespace mlir;
 
 namespace {
@@ -35,11 +34,10 @@ class OMPFunctionFilteringPass
 
   void runOnOperation() override {
     auto op = dyn_cast<omp::OffloadModuleInterface>(getOperation());
-    if (!op)
+    if (!op || !op.getIsTargetDevice())
       return;
 
-    bool isDeviceCompilation = op.getIsTargetDevice();
-    op->walk<WalkOrder::PostOrder>([&](func::FuncOp funcOp) {
+    op->walk<WalkOrder::PreOrder>([&](func::FuncOp funcOp) {
       // Do not filter functions with target regions inside, because they have
       // to be available for both host and device so that regular and reverse
       // offloading can be supported.
@@ -58,11 +56,21 @@ class OMPFunctionFilteringPass
       if (declareTargetOp && declareTargetOp.isDeclareTarget())
         declareType = declareTargetOp.getDeclareTargetDeviceType();
 
-      if ((isDeviceCompilation &&
-           declareType == omp::DeclareTargetDeviceType::host) ||
-          (!isDeviceCompilation &&
-           declareType == omp::DeclareTargetDeviceType::nohost))
-        funcOp->erase();
+      // Filtering a function here means removing its body and explicitly
+      // setting its omp.declare_target attribute, so that following
+      // translation/lowering/transformation passes will skip processing its
+      // contents, but preventing the calls to undefined symbols that could
+      // result if the function were deleted. The second stage of function
+      // filtering, at the MLIR to LLVM IR translation level, will remove these
+      // from the IR thanks to the mismatch between the omp.declare_target
+      // attribute and the target device.
+      if (declareType == omp::DeclareTargetDeviceType::host) {
+        funcOp.eraseBody();
+        funcOp.setVisibility(SymbolTable::Visibility::Private);
+        if (declareTargetOp)
+          declareTargetOp.setDeclareTarget(declareType,
+                                           omp::DeclareTargetCaptureClause::to);
+      }
     });
   }
 };

flang/test/Lower/OpenMP/function-filtering-2.f90

@@ -1,34 +1,37 @@
-! RUN: %flang_fc1 -fopenmp -emit-llvm %s -o - | FileCheck --check-prefixes=LLVM-HOST %s
-! RUN: %flang_fc1 -fopenmp -emit-mlir %s -o - | FileCheck --check-prefix=MLIR-HOST %s
-! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-llvm %s -o - | FileCheck --check-prefixes=LLVM-DEVICE %s
-! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-mlir %s -o - | FileCheck --check-prefix=MLIR-DEVICE %s
+! RUN: %flang_fc1 -fopenmp -emit-llvm %s -o - | FileCheck --check-prefix=LLVM %s
+! RUN: %flang_fc1 -fopenmp -emit-mlir %s -o - | FileCheck --check-prefix=MLIR %s
+! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-llvm %s -o - | FileCheck --check-prefix=LLVM %s
+! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-mlir %s -o - | FileCheck --check-prefix=MLIR %s
 
-! MLIR-HOST: func.func @{{.*}}implicit_invocation(
-! MLIR-DEVICE: func.func @{{.*}}implicit_invocation() attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>}
-! LLVM-HOST: define {{.*}} @{{.*}}implicit_invocation{{.*}}(
-! LLVM-DEVICE: define {{.*}} @{{.*}}implicit_invocation{{.*}}(
+! MLIR: func.func @{{.*}}implicit_invocation() attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>}
+! MLIR: return
+! LLVM: define {{.*}} @{{.*}}implicit_invocation{{.*}}(
 subroutine implicit_invocation()
 end subroutine implicit_invocation
 
-! MLIR-HOST: func.func @{{.*}}declaretarget(
-! MLIR-DEVICE: func.func @{{.*}}declaretarget() attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>}
-! LLVM-HOST: define {{.*}} @{{.*}}declaretarget{{.*}}(
-! LLVM-DEVICE: define {{.*}} @{{.*}}declaretarget{{.*}}(
+! MLIR: func.func @{{.*}}declaretarget() attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>}
+! MLIR: return
+! LLVM: define {{.*}} @{{.*}}declaretarget{{.*}}(
 subroutine declaretarget()
 !$omp declare target to(declaretarget) device_type(nohost)
     call implicit_invocation()
 end subroutine declaretarget
 
-! MLIR-HOST: func.func @{{.*}}no_declaretarget(
-! MLIR-DEVICE: func.func @{{.*}}no_declaretarget() attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>}
-! LLVM-HOST: define {{.*}} @{{.*}}no_declaretarget{{.*}}(
-! LLVM-DEVICE: define {{.*}} @{{.*}}no_declaretarget{{.*}}(
+! MLIR: func.func @{{.*}}no_declaretarget() attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>}
+! MLIR: return
+! LLVM: define {{.*}} @{{.*}}no_declaretarget{{.*}}(
 subroutine no_declaretarget()
 end subroutine no_declaretarget
 
 ! MLIR-HOST: func.func @{{.*}}main(
-! MLIR-DEVICE: func.func @{{.*}}main_omp_outline{{.*}}() attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>, omp.outline_parent_name = "_QQmain"} 
+! MLIR-HOST-NOT: func.func @{{.*}}main_omp_outline{{.*}}()
+! MLIR-DEVICE-NOT: func.func @{{.*}}main(
+! MLIR-DEVICE: func.func @{{.*}}main_omp_outline{{.*}}() attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>, omp.outline_parent_name = "_QQmain"}
+! MLIR-ALL: return
+
 ! LLVM-HOST: define {{.*}} @{{.*}}main{{.*}}(
+! LLVM-HOST-NOT: {{.*}} @{{.*}}__omp_offloading{{.*}}main_{{.*}}(
+! LLVM-DEVICE-NOT: {{.*}} @{{.*}}main{{.*}}(
 ! LLVM-DEVICE: define {{.*}} @{{.*}}__omp_offloading{{.*}}main_{{.*}}(
 program main
 !$omp target

flang/test/Lower/OpenMP/function-filtering.f90

@@ -1,43 +1,45 @@
 ! RUN: %flang_fc1 -fopenmp -emit-llvm %s -o - | FileCheck --check-prefixes=LLVM-HOST,LLVM-ALL %s
-! RUN: %flang_fc1 -fopenmp -emit-mlir %s -o - | FileCheck --check-prefix=MLIR-HOST %s
+! RUN: %flang_fc1 -fopenmp -emit-mlir %s -o - | FileCheck --check-prefixes=MLIR-HOST,MLIR-ALL %s
 ! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-llvm %s -o - | FileCheck --check-prefixes=LLVM-DEVICE,LLVM-ALL %s
-! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-mlir %s -o - | FileCheck --check-prefix=MLIR-DEVICE %s
+! RUN: %flang_fc1 -fopenmp -fopenmp-is-target-device -emit-mlir %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.
 
-! DISABLED, this portion of the test is disabled via the removal of the colon for the time 
-! being as filtering is enabled for device only for the time being while a fix is in progress. 
-! MLIR-HOST-NOT func.func @{{.*}}device_fn(
-! LLVM-HOST-NOT define {{.*}} @{{.*}}device_fn{{.*}}(
+! MLIR-ALL: func.func @{{.*}}device_fn(
+! MLIR-ALL: return
 
-! MLIR-DEVICE: func.func @{{.*}}device_fn(
-! LLVM-DEVICE: define {{.*}} @{{.*}}device_fn{{.*}}(
+! LLVM-ALL: define {{.*}} @{{.*}}device_fn{{.*}}(
 function device_fn() result(x)
   !$omp declare target to(device_fn) device_type(nohost)
   integer :: x
   x = 10
 end function device_fn
 
 ! MLIR-HOST: func.func @{{.*}}host_fn(
-! MLIR-DEVICE-NOT: func.func @{{.*}}host_fn(
+! MLIR-HOST: return
+! MLIR-DEVICE: func.func private @{{.*}}host_fn(
+! MLIR-DEVICE-NOT: return
+
 ! LLVM-HOST: define {{.*}} @{{.*}}host_fn{{.*}}(
-! LLVM-DEVICE-NOT: define {{.*}} @{{.*}}host_fn{{.*}}(
+! LLVM-DEVICE-NOT: {{.*}} @{{.*}}host_fn{{.*}}(
 function host_fn() result(x)
   !$omp declare target to(host_fn) device_type(host)
   integer :: x
   x = 10
 end function host_fn
 
 ! MLIR-HOST: func.func @{{.*}}target_subr(
+! MLIR-HOST: return
 ! MLIR-HOST-NOT: func.func @{{.*}}target_subr_omp_outline_0(
 ! MLIR-DEVICE-NOT: func.func @{{.*}}target_subr(
 ! MLIR-DEVICE: func.func @{{.*}}target_subr_omp_outline_0(
+! MLIR-DEVICE: return
 
 ! LLVM-ALL-NOT: define {{.*}} @{{.*}}target_subr_omp_outline_0{{.*}}(
 ! LLVM-HOST: define {{.*}} @{{.*}}target_subr{{.*}}(
-! LLVM-DEVICE-NOT: define {{.*}} @{{.*}}target_subr{{.*}}(
+! LLVM-DEVICE-NOT: {{.*}} @{{.*}}target_subr{{.*}}(
 ! LLVM-ALL: define {{.*}} @__omp_offloading_{{.*}}_{{.*}}_target_subr__{{.*}}(
 subroutine target_subr(x)
   integer, intent(out) :: x

flang/test/Transforms/omp-function-filtering.mlir

@@ -1,12 +1,19 @@
 // RUN: fir-opt -split-input-file --omp-function-filtering %s | FileCheck %s
 
-// CHECK:     func.func @any
-// CHECK:     func.func @nohost
-// CHECK-NOT: func.func @host
-// CHECK-NOT: func.func @none
-// CHECK:     func.func @nohost_target
-// CHECK:     func.func @host_target
-// CHECK:     func.func @none_target
+// CHECK: func.func @any
+// CHECK: return
+// CHECK: func.func @nohost
+// CHECK: return
+// CHECK: func.func private @host
+// CHECK-NOT: return
+// CHECK: func.func private @none
+// CHECK-NOT: return
+// CHECK: func.func @nohost_target
+// CHECK: return
+// CHECK: func.func @host_target
+// CHECK: return
+// CHECK: func.func @none_target
+// CHECK: return
 module attributes {omp.is_target_device = true} {
   func.func @any() -> ()
       attributes {
@@ -56,13 +63,20 @@ module attributes {omp.is_target_device = true} {
 
 // -----
 
-// CHECK:     func.func @any
-// CHECK-NOT: func.func @nohost
-// CHECK:     func.func @host
-// CHECK:     func.func @none
-// CHECK:     func.func @nohost_target
-// CHECK:     func.func @host_target
-// CHECK:     func.func @none_target
+// CHECK: func.func @any
+// CHECK: return
+// CHECK: func.func @nohost
+// CHECK: return
+// CHECK: func.func @host
+// CHECK: return
+// CHECK: func.func @none
+// CHECK: return
+// CHECK: func.func @nohost_target
+// CHECK: return
+// CHECK: func.func @host_target
+// CHECK: return
+// CHECK: func.func @none_target
+// CHECK: return
 module attributes {omp.is_target_device = false} {
   func.func @any() -> ()
       attributes {

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -1763,19 +1763,12 @@ convertDeclareTargetAttr(Operation *op,
   if (FunctionOpInterface funcOp = dyn_cast<FunctionOpInterface>(op)) {
     if (auto offloadMod = dyn_cast<omp::OffloadModuleInterface>(
             op->getParentOfType<ModuleOp>().getOperation())) {
-      bool isDeviceCompilation = offloadMod.getIsTargetDevice();
-      // FIXME: Temporarily disabled for host as it causes some issues when
-      // lowering while removing functions at the current time.
-      if (!isDeviceCompilation)
+      if (!offloadMod.getIsTargetDevice())
         return success();
 
       omp::DeclareTargetDeviceType declareType =
           declareTargetAttr.getDeviceType().getValue();
-
-      if ((isDeviceCompilation &&
-           declareType == omp::DeclareTargetDeviceType::host) ||
-          (!isDeviceCompilation &&
-           declareType == omp::DeclareTargetDeviceType::nohost)) {
+      if (declareType == omp::DeclareTargetDeviceType::host) {
         llvm::Function *llvmFunc =
             moduleTranslation.lookupFunction(funcOp.getName());
         llvmFunc->dropAllReferences();

mlir/test/Target/LLVMIR/openmp-llvm.mlir

@@ -2547,19 +2547,17 @@ module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true,
 // -----
 
 module attributes {omp.is_target_device = false} {
-  // DISABLED, this portion of the test is disabled via the removal of the colon for the time 
-  // being as filtering is enabled for device only for the time being while a fix is in progress. 
-  // CHECK-NOT @filter_host_nohost
-  llvm.func @filter_host_nohost() -> ()
+  // CHECK: define void @filter_nohost
+  llvm.func @filter_nohost() -> ()
       attributes {
         omp.declare_target =
           #omp.declaretarget<device_type = (nohost), capture_clause = (to)>
       } {
     llvm.return
   }
 
-  // CHECK: @filter_host_host
-  llvm.func @filter_host_host() -> ()
+  // CHECK: define void @filter_host
+  llvm.func @filter_host() -> ()
       attributes {
         omp.declare_target =
           #omp.declaretarget<device_type = (host), capture_clause = (to)>
@@ -2571,17 +2569,17 @@ module attributes {omp.is_target_device = false} {
 // -----
 
 module attributes {omp.is_target_device = true} {
-  // CHECK: @filter_device_nohost
-  llvm.func @filter_device_nohost() -> ()
+  // CHECK: define void @filter_nohost
+  llvm.func @filter_nohost() -> ()
       attributes {
         omp.declare_target =
           #omp.declaretarget<device_type = (nohost), capture_clause = (to)>
       } {
     llvm.return
   }
 
-  // CHECK-NOT: @filter_device_host
-  llvm.func @filter_device_host() -> ()
+  // CHECK-NOT: define void @filter_host
+  llvm.func @filter_host() -> ()
       attributes {
         omp.declare_target =
           #omp.declaretarget<device_type = (host), capture_clause = (to)>