diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h --- a/flang/include/flang/Optimizer/Transforms/Passes.h +++ b/flang/include/flang/Optimizer/Transforms/Passes.h @@ -75,6 +75,9 @@ std::unique_ptr createPolymorphicOpConversionPass(); std::unique_ptr> createOMPEarlyOutliningPass(); +std::unique_ptr> +createOMPMarkDeclareTargetPass(); + // declarative passes #define GEN_PASS_REGISTRATION #include "flang/Optimizer/Transforms/Passes.h.inc" diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td --- a/flang/include/flang/Optimizer/Transforms/Passes.td +++ b/flang/include/flang/Optimizer/Transforms/Passes.td @@ -308,6 +308,12 @@ 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 --- a/flang/lib/Frontend/FrontendActions.cpp +++ b/flang/lib/Frontend/FrontendActions.cpp @@ -310,11 +310,14 @@ mlirModule->getOperation())) isDevice = offloadMod.getIsTargetDevice(); + pm.addPass(fir::createOMPMarkDeclareTargetPass()); + if (isDevice) pm.addPass(fir::createOMPEarlyOutliningPass()); } pm.enableVerifier(/*verifyPasses=*/true); + pm.addPass(std::make_unique()); if (mlir::failed(pm.run(*mlirModule))) { diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt --- a/flang/lib/Optimizer/Transforms/CMakeLists.txt +++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt @@ -17,7 +17,8 @@ PolymorphicOpConversion.cpp LoopVersioning.cpp OMPEarlyOutlining.cpp - + OMPMarkDeclareTarget.cpp + DEPENDS FIRDialect FIROptTransformsPassIncGen diff --git a/flang/lib/Optimizer/Transforms/OMPMarkDeclareTarget.cpp b/flang/lib/Optimizer/Transforms/OMPMarkDeclareTarget.cpp new file mode 100644 --- /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 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 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 --- /dev/null +++ b/flang/test/Lower/OpenMP/declare-target-implicit-func-and-subr-cap.f90 @@ -0,0 +1,215 @@ +!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s + +! 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 + +! CHECK-LABEL: func.func @_QPtarget_function_twice_device +! CHECK-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 + +!! ----- + +! CHECK-LABEL: func.func @_QPimplicitly_captured_nest +! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget{{.*}}} +function implicitly_captured_nest() result(k) + integer :: i + i = 10 + k = i +end function implicitly_captured_nest + +! CHECK-LABEL: func.func @_QPimplicitly_captured_one +! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget{{.*}}} +function implicitly_captured_two() result(k) + integer :: i + i = 10 + k = i +end function implicitly_captured_two + +! CHECK-LABEL: func.func @_QPtarget_function_test +! CHECK-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 + +! CHECK-LABEL: func.func @_QPtarget_function_test_device +! CHECK-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 + +!! ----- + +! CHECK-LABEL: func.func @_QPimplicitly_captured_with_dev_type_recursive +! CHECK-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 + +! CHECK-LABEL: func.func @_QPtarget_function_with_dev_type_recurse +! CHECK-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 + +! CHECK-LABEL: func.func @_QMtest_modulePimplicitly_captured_two_twice +! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget{{.*}}} + function implicitly_captured_two_twice() result(y) + integer :: y + y = 5 + end function implicitly_captured_two_twice + +! CHECK-LABEL: func.func @_QMtest_modulePtarget_function_test_device +! CHECK-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 + +! CHECK-LABEL: func.func @_QPimplicitly_captured_recursive +! CHECK-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 + +! CHECK-LABEL: func.func @_QPcaller_recursive +! CHECK-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 --- /dev/null +++ b/flang/test/Lower/OpenMP/declare-target-implicit-tarop-cap.f90 @@ -0,0 +1,68 @@ +!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s + +! CHECK-LABEL: func.func @_QPimplicit_capture +! CHECK-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 + +! CHECK-LABEL: func.func @_QPimplicitly_captured_two_twice +! CHECK-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 + +!! ----- + +! CHECK-LABEL: func.func @_QPimplicitly_captured_recursive +! CHECK-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 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 rename from flang/test/Lower/OpenMP/omp-is-device.f90 rename to flang/test/Lower/OpenMP/is-device.f90