[flang-commits] [flang] f20b67a - [Flang][MLIR][OpenMP] Improve device-only function filtering

Sergio Afonso via flang-commits flang-commits at lists.llvm.org
Thu Aug 10 03:30:07 PDT 2023


Author: Sergio Afonso
Date: 2023-08-10T11:29:45+01:00
New Revision: f20b67a81c66f44ab2f57b5a2fce04dafe758a86

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

LOG: [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

Added: 
    

Modified: 
    flang/include/flang/Optimizer/Transforms/Passes.td
    flang/lib/Frontend/FrontendActions.cpp
    flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
    flang/test/Lower/OpenMP/function-filtering-2.f90
    flang/test/Lower/OpenMP/function-filtering.f90
    flang/test/Transforms/omp-function-filtering.mlir
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
    mlir/test/Target/LLVMIR/openmp-llvm.mlir

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 66725a113028fc..9474edf13ce463 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/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"

diff  --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index 6cf60dcfccc5a0..d2c6e518a0500d 100644
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/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());
     }
   }

diff  --git a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
index 7784c90e930000..43fa5b7c4de241 100644
--- a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
+++ b/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);
+      }
     });
   }
 };

diff  --git a/flang/test/Lower/OpenMP/function-filtering-2.f90 b/flang/test/Lower/OpenMP/function-filtering-2.f90
index dbad29bb7c1d76..a12b88d21f8785 100644
--- a/flang/test/Lower/OpenMP/function-filtering-2.f90
+++ b/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

diff  --git a/flang/test/Lower/OpenMP/function-filtering.f90 b/flang/test/Lower/OpenMP/function-filtering.f90
index 117cef6677d321..83302205943f33 100644
--- a/flang/test/Lower/OpenMP/function-filtering.f90
+++ b/flang/test/Lower/OpenMP/function-filtering.f90
@@ -1,19 +1,16 @@
 ! 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
@@ -21,9 +18,12 @@ function device_fn() result(x)
 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
@@ -31,13 +31,15 @@ function host_fn() result(x)
 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

diff  --git a/flang/test/Transforms/omp-function-filtering.mlir b/flang/test/Transforms/omp-function-filtering.mlir
index ccb11caf7c81f7..44777e2cac30c5 100644
--- a/flang/test/Transforms/omp-function-filtering.mlir
+++ b/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 {

diff  --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 97f1362f470c7f..ac410e60e0bb6e 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/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();

diff  --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
index 2a76164ffdd2e0..73b6b007e5a56b 100644
--- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
@@ -2547,10 +2547,8 @@ 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)>
@@ -2558,8 +2556,8 @@ module attributes {omp.is_target_device = false} {
     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,8 +2569,8 @@ 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)>
@@ -2580,8 +2578,8 @@ module attributes {omp.is_target_device = true} {
     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)>


        


More information about the flang-commits mailing list