[Flang][OpenMP][MLIR] An mlir transformation pass for marking FuncOp's implicitly...

This pass will mark functions called from TargetOp's
and declare target functions as implicitly declare
target by adding the MLIR declare target attribute
directly to the function.

This pass executes after the initial lowering of Fortran's PFT
to MLIR (FIR/OMP+Arith etc.) and is one of a series of passes
that aim to clean up the MLIR for offloading (seperate passes
in different patches, one for early outlining, another for declare
target function filtering).

Reviewers: jsjodin, skatrak, kiaranchandramohan

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

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h
index 3272cb3..8aeb3e3 100644 (file)
--- a/flang/include/flang/Optimizer/Transforms/Passes.h
+++ b/flang/include/flang/Optimizer/Transforms/Passes.h
@@ -77,6 +77,8 @@ std::unique_ptr<mlir::Pass> createPolymorphicOpConversionPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 createOMPEarlyOutliningPass();
 std::unique_ptr<mlir::Pass> createOMPFunctionFilteringPass();
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+createOMPMarkDeclareTargetPass();
 
 // declarative passes
 #define GEN_PASS_REGISTRATION

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 40a08c9..66725a1 100644 (file)
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -308,6 +308,12 @@ def OMPEarlyOutliningPass
     the optimizer to perform transforms across target region boundaries.
   }];
   let constructor = "::fir::createOMPEarlyOutliningPass()";
+}
+
+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"];
 }
 

diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index c03c3fd..7c44b46 100644 (file)
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/flang/lib/Frontend/FrontendActions.cpp
@@ -310,6 +310,7 @@ bool CodeGenAction::beginSourceFileAction() {
             mlirModule->getOperation()))
       isDevice = offloadMod.getIsTargetDevice();
 
+    pm.addPass(fir::createOMPMarkDeclareTargetPass());
     if (isDevice)
       pm.addPass(fir::createOMPEarlyOutliningPass());
     pm.addPass(fir::createOMPFunctionFilteringPass());

diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index 1808542..e3943be 100644 (file)
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -18,6 +18,7 @@ add_flang_library(FIRTransforms
   LoopVersioning.cpp
   OMPEarlyOutlining.cpp
   OMPFunctionFiltering.cpp
+  OMPMarkDeclareTarget.cpp
 
   DEPENDS
   FIRDialect

diff --git a/flang/lib/Optimizer/Transforms/OMPMarkDeclareTarget.cpp b/flang/lib/Optimizer/Transforms/OMPMarkDeclareTarget.cpp
new file mode 100644 (file)
index 0000000..bebcb2a
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/OMPMarkDeclareTarget.cpp
@@ -0,0 +1,97 @@
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SmallPtrSet.h"
+
+namespace fir {
+#define GEN_PASS_DEF_OMPMARKDECLARETARGETPASS
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+namespace {
+class OMPMarkDeclareTargetPass
+    : public fir::impl::OMPMarkDeclareTargetPassBase<OMPMarkDeclareTargetPass> {
+
+  void markNestedFuncs(mlir::omp::DeclareTargetDeviceType parentDevTy,
+                       mlir::omp::DeclareTargetCaptureClause parentCapClause,
+                       mlir::Operation *currOp,
+                       llvm::SmallPtrSet<mlir::Operation *, 16> visited) {
+    if (visited.contains(currOp))
+      return;
+    visited.insert(currOp);
+
+    currOp->walk([&, this](mlir::Operation *op) {
+      if (auto callOp = llvm::dyn_cast<mlir::CallOpInterface>(op)) {
+        if (auto symRef = llvm::dyn_cast_if_present<mlir::SymbolRefAttr>(
+                callOp.getCallableForCallee())) {
+          if (auto currFOp =
+                  getOperation().lookupSymbol<mlir::func::FuncOp>(symRef)) {
+            auto current = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
+                currFOp.getOperation());
+
+            if (current.isDeclareTarget()) {
+              auto currentDt = current.getDeclareTargetDeviceType();
+
+              // Found the same function twice, with different device_types,
+              // mark as Any as it belongs to both
+              if (currentDt != parentDevTy &&
+                  currentDt != mlir::omp::DeclareTargetDeviceType::any) {
+                current.setDeclareTarget(
+                    mlir::omp::DeclareTargetDeviceType::any,
+                    current.getDeclareTargetCaptureClause());
+              }
+            } else {
+              current.setDeclareTarget(parentDevTy, parentCapClause);
+            }
+
+            markNestedFuncs(parentDevTy, parentCapClause, currFOp, visited);
+          }
+        }
+      }
+    });
+  }
+
+  // This pass executes on mlir::ModuleOp's marking functions contained within
+  // as implicitly declare target if they are called from within an explicitly
+  // marked declare target function or a target region (TargetOp)
+  void runOnOperation() override {
+    for (auto functionOp : getOperation().getOps<mlir::func::FuncOp>()) {
+      auto declareTargetOp = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
+          functionOp.getOperation());
+      if (declareTargetOp.isDeclareTarget()) {
+        llvm::SmallPtrSet<mlir::Operation *, 16> visited;
+        markNestedFuncs(declareTargetOp.getDeclareTargetDeviceType(),
+                        declareTargetOp.getDeclareTargetCaptureClause(),
+                        functionOp, visited);
+      }
+    }
+
+    // TODO: Extend to work with reverse-offloading, this shouldn't
+    // require too much effort, just need to check the device clause
+    // when it's lowering has been implemented and change the
+    // DeclareTargetDeviceType argument from nohost to host depending on
+    // the contents of the device clause
+    getOperation()->walk([&](mlir::omp::TargetOp tarOp) {
+      llvm::SmallPtrSet<mlir::Operation *, 16> visited;
+      markNestedFuncs(mlir::omp::DeclareTargetDeviceType::nohost,
+                      mlir::omp::DeclareTargetCaptureClause::to, tarOp,
+                      visited);
+    });
+  }
+};
+
+} // namespace
+
+namespace fir {
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+createOMPMarkDeclareTargetPass() {
+  return std::make_unique<OMPMarkDeclareTargetPass>();
+}
+} // namespace fir

diff --git a/flang/test/Lower/OpenMP/omp-declare-target-data.f90 b/flang/test/Lower/OpenMP/declare-target-data.f90
similarity index 100%
rename from flang/test/Lower/OpenMP/omp-declare-target-data.f90
rename to flang/test/Lower/OpenMP/declare-target-data.f90

diff --git a/flang/test/Lower/OpenMP/omp-declare-target-func-and-subr.f90 b/flang/test/Lower/OpenMP/declare-target-func-and-subr.f90
similarity index 100%
rename from flang/test/Lower/OpenMP/omp-declare-target-func-and-subr.f90
rename to flang/test/Lower/OpenMP/declare-target-func-and-subr.f90

diff --git a/flang/test/Lower/OpenMP/declare-target-implicit-func-and-subr-cap.f90 b/flang/test/Lower/OpenMP/declare-target-implicit-func-and-subr-cap.f90
new file mode 100644 (file)
index 0000000..a9f0c88
--- /dev/null
+++ b/flang/test/Lower/OpenMP/declare-target-implicit-func-and-subr-cap.f90
@@ -0,0 +1,216 @@
+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
+!RUN: %flang_fc1 -emit-fir -fopenmp -fopenmp-is-device %s -o - | FileCheck %s  --check-prefix=DEVICE
+
+! CHECK-LABEL: func.func @_QPimplicitly_captured
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function implicitly_captured(toggle) result(k)
+   integer :: i, j, k
+   logical :: toggle
+   i = 10
+   j = 5
+   if (toggle) then
+      k = i
+   else
+      k = j
+   end if
+end function implicitly_captured
+
+
+! CHECK-LABEL: func.func @_QPtarget_function
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function target_function(toggle) result(i)
+!$omp declare target
+   integer :: i
+   logical :: toggle
+   i = implicitly_captured(toggle)
+end function target_function
+
+!! -----
+
+! CHECK-LABEL: func.func @_QPimplicitly_captured_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function implicitly_captured_twice() result(k)
+   integer :: i
+   i = 10
+   k = i
+end function implicitly_captured_twice
+
+! CHECK-LABEL: func.func @_QPtarget_function_twice_host
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>{{.*}}}
+function target_function_twice_host() result(i)
+!$omp declare target to(target_function_twice_host) device_type(host)
+   integer :: i
+   i = implicitly_captured_twice()
+end function target_function_twice_host
+
+! DEVICE-LABEL: func.func @_QPtarget_function_twice_device
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function target_function_twice_device() result(i)
+!$omp declare target to(target_function_twice_device) device_type(nohost)
+   integer :: i
+   i = implicitly_captured_twice()
+end function target_function_twice_device
+
+!! -----
+
+! DEVICE-LABEL: func.func @_QPimplicitly_captured_nest
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function implicitly_captured_nest() result(k)
+   integer :: i
+   i = 10
+   k = i
+end function implicitly_captured_nest
+
+! DEVICE-LABEL: func.func @_QPimplicitly_captured_one
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to){{.*}}}
+function implicitly_captured_one() result(k)
+   k = implicitly_captured_nest()
+end function implicitly_captured_one
+
+! DEVICE-LABEL: func.func @_QPimplicitly_captured_two
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function implicitly_captured_two() result(k)
+   integer :: i
+   i = 10
+   k = i
+end function implicitly_captured_two
+
+! DEVICE-LABEL: func.func @_QPtarget_function_test
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function target_function_test() result(j)
+!$omp declare target to(target_function_test) device_type(nohost)
+   integer :: i, j
+   i = implicitly_captured_one()
+   j = implicitly_captured_two() + i
+end function target_function_test
+
+!! -----
+
+! CHECK-LABEL: func.func @_QPimplicitly_captured_nest_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function implicitly_captured_nest_twice() result(k)
+   integer :: i
+   i = 10
+   k = i
+end function implicitly_captured_nest_twice
+
+! CHECK-LABEL: func.func @_QPimplicitly_captured_one_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function implicitly_captured_one_twice() result(k)
+   k = implicitly_captured_nest_twice()
+end function implicitly_captured_one_twice
+
+! CHECK-LABEL: func.func @_QPimplicitly_captured_two_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function implicitly_captured_two_twice() result(k)
+   integer :: i
+   i = 10
+   k = i
+end function implicitly_captured_two_twice
+
+! DEVICE-LABEL: func.func @_QPtarget_function_test_device
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function target_function_test_device() result(j)
+   !$omp declare target to(target_function_test_device) device_type(nohost)
+   integer :: i, j
+   i = implicitly_captured_one_twice()
+   j = implicitly_captured_two_twice() + i
+end function target_function_test_device
+
+! CHECK-LABEL: func.func @_QPtarget_function_test_host
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>{{.*}}}
+function target_function_test_host() result(j)
+   !$omp declare target to(target_function_test_host) device_type(host)
+   integer :: i, j
+   i = implicitly_captured_one_twice()
+   j = implicitly_captured_two_twice() + i
+end function target_function_test_host
+
+!! -----
+
+! DEVICE-LABEL: func.func @_QPimplicitly_captured_with_dev_type_recursive
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+recursive function implicitly_captured_with_dev_type_recursive(increment) result(k)
+!$omp declare target to(implicitly_captured_with_dev_type_recursive) device_type(host)
+   integer :: increment, k
+   if (increment == 10) then
+      k = increment
+   else
+      k = implicitly_captured_with_dev_type_recursive(increment + 1)
+   end if
+end function implicitly_captured_with_dev_type_recursive
+
+! DEVICE-LABEL: func.func @_QPtarget_function_with_dev_type_recurse
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function target_function_with_dev_type_recurse() result(i)
+!$omp declare target to(target_function_with_dev_type_recurse) device_type(nohost)
+   integer :: i
+   i = implicitly_captured_with_dev_type_recursive(0)
+end function target_function_with_dev_type_recurse
+
+!! -----
+
+module test_module
+contains
+! CHECK-LABEL: func.func @_QMtest_modulePimplicitly_captured_nest_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+   function implicitly_captured_nest_twice() result(i)
+      integer :: i
+      i = 10
+   end function implicitly_captured_nest_twice
+
+! CHECK-LABEL: func.func @_QMtest_modulePimplicitly_captured_one_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+   function implicitly_captured_one_twice() result(k)
+      !$omp declare target to(implicitly_captured_one_twice) device_type(host)
+      k = implicitly_captured_nest_twice()
+   end function implicitly_captured_one_twice
+
+! DEVICE-LABEL: func.func @_QMtest_modulePimplicitly_captured_two_twice
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+   function implicitly_captured_two_twice() result(y)
+      integer :: y
+      y = 5
+   end function implicitly_captured_two_twice
+
+! DEVICE-LABEL: func.func @_QMtest_modulePtarget_function_test_device
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+   function target_function_test_device() result(j)
+      !$omp declare target to(target_function_test_device) device_type(nohost)
+      integer :: i, j
+      i = implicitly_captured_one_twice()
+      j = implicitly_captured_two_twice() + i
+   end function target_function_test_device
+end module test_module
+
+!! -----
+
+program mb
+   interface
+      subroutine caller_recursive
+         !$omp declare target to(caller_recursive) device_type(nohost)
+      end subroutine
+
+      recursive subroutine implicitly_captured_recursive(increment)
+         integer :: increment
+      end subroutine
+   end interface
+end program
+
+! DEVICE-LABEL: func.func @_QPimplicitly_captured_recursive
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+recursive subroutine implicitly_captured_recursive(increment)
+   integer :: increment
+   if (increment == 10) then
+      return
+   else
+      call implicitly_captured_recursive(increment + 1)
+   end if
+end subroutine
+
+! DEVICE-LABEL: func.func @_QPcaller_recursive
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+subroutine caller_recursive
+!$omp declare target to(caller_recursive) device_type(nohost)
+   call implicitly_captured_recursive(0)
+end subroutine

diff --git a/flang/test/Lower/OpenMP/declare-target-implicit-tarop-cap.f90 b/flang/test/Lower/OpenMP/declare-target-implicit-tarop-cap.f90
new file mode 100644 (file)
index 0000000..b597513
--- /dev/null
+++ b/flang/test/Lower/OpenMP/declare-target-implicit-tarop-cap.f90
@@ -0,0 +1,69 @@
+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
+!RUN: %flang_fc1 -emit-fir -fopenmp -fopenmp-is-device %s -o - | FileCheck %s  --check-prefix=DEVICE
+
+! DEVICE-LABEL: func.func @_QPimplicit_capture
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function implicit_capture() result(i)
+   implicit none
+   integer :: i
+   i = 1
+end function implicit_capture
+
+subroutine subr_target()
+   integer :: n
+!$omp target map(tofrom:n)
+   n = implicit_capture()
+!$omp end target
+end subroutine
+
+!! -----
+
+! CHECK-LABEL: func.func @_QPimplicitly_captured_nest_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function implicitly_captured_nest_twice() result(i)
+   integer :: i
+   i = 10
+end function implicitly_captured_nest_twice
+
+! CHECK-LABEL: func.func @_QPimplicitly_captured_one_twice
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}}
+function implicitly_captured_one_twice() result(k)
+!$omp declare target to(implicitly_captured_one_twice) device_type(host)
+   k = implicitly_captured_nest_twice()
+end function implicitly_captured_one_twice
+
+! DEVICE-LABEL: func.func @_QPimplicitly_captured_two_twice
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+function implicitly_captured_two_twice() result(y)
+   integer :: y
+   y = 5
+end function implicitly_captured_two_twice
+
+
+function target_function_test_device() result(j)
+   integer :: i, j
+   !$omp target map(tofrom: i, j)
+   i = implicitly_captured_one_twice()
+   j = implicitly_captured_two_twice() + i
+   !$omp end target
+end function target_function_test_device
+
+!! -----
+
+! DEVICE-LABEL: func.func @_QPimplicitly_captured_recursive
+! DEVICE-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}}
+recursive function implicitly_captured_recursive(increment) result(k)
+   integer :: increment, k
+   if (increment == 10) then
+      k = increment
+   else
+      k = implicitly_captured_recursive(increment + 1)
+   end if
+end function implicitly_captured_recursive
+
+function target_function_recurse() result(i)
+   integer :: i
+   !$omp target map(tofrom: i)
+   i = implicitly_captured_recursive(0)
+   !$omp end target
+end function target_function_recurse

diff --git a/flang/test/Lower/OpenMP/omp-host-ir-flag.f90 b/flang/test/Lower/OpenMP/host-ir-flag.f90
similarity index 100%
rename from flang/test/Lower/OpenMP/omp-host-ir-flag.f90
rename to flang/test/Lower/OpenMP/host-ir-flag.f90

diff --git a/flang/test/Lower/OpenMP/omp-is-device.f90 b/flang/test/Lower/OpenMP/is-device.f90
similarity index 100%
rename from flang/test/Lower/OpenMP/omp-is-device.f90
rename to flang/test/Lower/OpenMP/is-device.f90