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

[skatrak]

Summary

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

[kiranchandramohan]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

1064–1068 ↗

(On Diff #542462)

This part has to come in a separate patch.

[skatrak]

Update patch after splitting some work into D156988.

[kiranchandramohan]

Looks OK to me. @agozillon or @jsjodin can accept the patch if they agree.

[jsjodin]

This looks good to me.

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

322

Nit: I think we should say "target device" instead of "device" if it is visible to the user.

[agozillon]

LGTM, great job. You can wait on @jsjodin to review as well if you wish to have another set of eyes on the patch though!

[skatrak]

Rebase and address review comment.

[skatrak]

Thank you all for giving this a look. I'll land it tomorrow.

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

def OMPMarkDeclareTargetPass
     : Pass<"omp-mark-declare-target", "mlir::ModuleOp"> {
   let summary = "Marks all functions called by an OpenMP declare target function as declare target";
   let constructor = "::fir::createOMPMarkDeclareTargetPass()";
   let dependentDialects = ["mlir::omp::OpenMPDialect"];
 }
 
 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"
   ];
 }
 
 #endif // FLANG_OPTIMIZER_TRANSFORMS_PASSES

flang/lib/Frontend/FrontendActions.cpp

if (ci.getInvocation().getFrontendOpts().features.IsEnabled(
     bool isDevice = false;
     if (auto offloadMod = llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(
             mlirModule->getOperation()))
       isDevice = offloadMod.getIsTargetDevice();
 
     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());
     }
   }
 
   pm.enableVerifier(/*verifyPasses=*/true);
   pm.addPass(std::make_unique<Fortran::lower::VerifierPass>());
 
   if (mlir::failed(pm.run(*mlirModule))) {

flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp

 #include "mlir/IR/BuiltinOps.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace fir {
 #define GEN_PASS_DEF_OMPFUNCTIONFILTERING
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 } // namespace fir
 
-using namespace fir;
 using namespace mlir;
 
 namespace {
 class OMPFunctionFilteringPass
     : public fir::impl::OMPFunctionFilteringBase<OMPFunctionFilteringPass> {
 public:
   OMPFunctionFilteringPass() = default;
 
   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.
       bool hasTargetRegion =
           funcOp
               ->walk<WalkOrder::PreOrder>(
                   [&](omp::TargetOp) { return WalkResult::interrupt(); })
               .wasInterrupted();
       if (hasTargetRegion)
         return;
 
       omp::DeclareTargetDeviceType declareType =
           omp::DeclareTargetDeviceType::host;
       auto declareTargetOp =
           dyn_cast<omp::DeclareTargetInterface>(funcOp.getOperation());
       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);
+      }
     });
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> fir::createOMPFunctionFilteringPass() {
   return std::make_unique<OMPFunctionFilteringPass>();
 }

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

-! 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
     call declaretarget()
     call no_declaretarget()
 !$omp end target
 end program main

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

 ! 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
   !$omp target map(from:x)
     x = 10
   !$omp end target
 end subroutine target_subr

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

 // 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 {
         omp.declare_target =
           #omp.declaretarget<device_type = (any), capture_clause = (to)>
       } {
     func.return
   }
module attributes {omp.is_target_device = true} {
   func.func @none_target() -> () {
     omp.target {}
     func.return
   }
 }
 
 // -----
 
-// 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 {
         omp.declare_target =
           #omp.declaretarget<device_type = (any), capture_clause = (to)>
       } {
     func.return
   }

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

convertDeclareTargetAttr(Operation *op,
   // because the omp.target operation inside must be translated to LLVM, but the
   // wrapper functions themselves must not remain at the end of the process.
   // We know that functions where omp.declare_target does not match
   // omp.is_target_device at this stage can only be wrapper functions because
   // those that aren't are removed earlier as an MLIR transformation pass.
   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();
         llvmFunc->eraseFromParent();
       }
     }
   }
   return success();

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

 // CHECK: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
 // CHECK: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 1
 // CHECK: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 0
 module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true, assume_no_thread_state = 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)>
       } {
     llvm.return
   }
 }
 
 // -----
 
 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)>
       } {
     llvm.return
   }
 }