From 062fce6f4dd142cf9275abb3358f88b61a2db4e8 Mon Sep 17 00:00:00 2001 From: Andrew Gozillon Date: Mon, 17 Jul 2023 08:20:13 -0500 Subject: [PATCH] [Flang][OpenMP][MLIR] An mlir transformation pass for marking FuncOp's implicitly called from TargetOp's and declare target marked FuncOp's as implicitly declare target 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 --- flang/include/flang/Optimizer/Transforms/Passes.h | 2 + flang/include/flang/Optimizer/Transforms/Passes.td | 6 + flang/lib/Frontend/FrontendActions.cpp | 1 + flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 + .../Optimizer/Transforms/OMPMarkDeclareTarget.cpp | 97 +++++++++ ...are-target-data.f90 => declare-target-data.f90} | 0 ...d-subr.f90 => declare-target-func-and-subr.f90} | 0 .../declare-target-implicit-func-and-subr-cap.f90 | 216 +++++++++++++++++++++ .../OpenMP/declare-target-implicit-tarop-cap.f90 | 69 +++++++ .../{omp-host-ir-flag.f90 => host-ir-flag.f90} | 0 .../OpenMP/{omp-is-device.f90 => is-device.f90} | 0 11 files changed, 392 insertions(+) create mode 100644 flang/lib/Optimizer/Transforms/OMPMarkDeclareTarget.cpp rename flang/test/Lower/OpenMP/{omp-declare-target-data.f90 => declare-target-data.f90} (100%) rename flang/test/Lower/OpenMP/{omp-declare-target-func-and-subr.f90 => declare-target-func-and-subr.f90} (100%) create mode 100644 flang/test/Lower/OpenMP/declare-target-implicit-func-and-subr-cap.f90 create mode 100644 flang/test/Lower/OpenMP/declare-target-implicit-tarop-cap.f90 rename flang/test/Lower/OpenMP/{omp-host-ir-flag.f90 => host-ir-flag.f90} (100%) rename flang/test/Lower/OpenMP/{omp-is-device.f90 => is-device.f90} (100%) diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h index 3272cb3..8aeb3e3 100644 --- a/flang/include/flang/Optimizer/Transforms/Passes.h +++ b/flang/include/flang/Optimizer/Transforms/Passes.h @@ -77,6 +77,8 @@ std::unique_ptr createPolymorphicOpConversionPass(); std::unique_ptr> createOMPEarlyOutliningPass(); std::unique_ptr createOMPFunctionFilteringPass(); +std::unique_ptr> +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 --- 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 --- 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 --- 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 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 { + + void markNestedFuncs(mlir::omp::DeclareTargetDeviceType parentDevTy, + mlir::omp::DeclareTargetCaptureClause parentCapClause, + mlir::Operation *currOp, + llvm::SmallPtrSet visited) { + if (visited.contains(currOp)) + return; + visited.insert(currOp); + + currOp->walk([&, this](mlir::Operation *op) { + if (auto callOp = llvm::dyn_cast(op)) { + if (auto symRef = llvm::dyn_cast_if_present( + callOp.getCallableForCallee())) { + if (auto currFOp = + getOperation().lookupSymbol(symRef)) { + auto current = llvm::dyn_cast( + 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()) { + auto declareTargetOp = llvm::dyn_cast( + functionOp.getOperation()); + if (declareTargetOp.isDeclareTarget()) { + llvm::SmallPtrSet 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 visited; + markNestedFuncs(mlir::omp::DeclareTargetDeviceType::nohost, + mlir::omp::DeclareTargetCaptureClause::to, tarOp, + visited); + }); + } +}; + +} // namespace + +namespace fir { +std::unique_ptr> +createOMPMarkDeclareTargetPass() { + return std::make_unique(); +} +} // 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 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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} + 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{{.*}}} + 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{{.*}}} + 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{{.*}}} + 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{{.*}}} +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{{.*}}} +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 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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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{{.*}}} +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 -- 2.7.4