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 @@ -20,6 +20,7 @@ class Operation; class Pass; class Region; +class ModuleOp; } // namespace mlir namespace fir { @@ -72,6 +73,8 @@ std::unique_ptr createAlgebraicSimplificationPass(const mlir::GreedyRewriteConfig &config); std::unique_ptr createPolymorphicOpConversionPass(); +std::unique_ptr> +createCaptureImplicitlyDeclareTargetPass(); // declarative passes #define GEN_PASS_REGISTRATION 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 @@ -298,4 +298,11 @@ let dependentDialects = [ "fir::FIROpsDialect" ]; } +def CaptureImplicitlyDeclareTargetPass + : Pass<"capture-implicitly-declare-target", "mlir::ModuleOp"> { + let summary = "Marks all functions referenced by a declare target function as declare target"; + let constructor = "::fir::createCaptureImplicitlyDeclareTargetPass()"; + let dependentDialects = ["mlir::omp::OpenMPDialect"]; +} + #endif // FLANG_OPTIMIZER_TRANSFORMS_PASSES 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 @@ -23,6 +23,7 @@ #include "flang/Optimizer/Dialect/Support/KindMapping.h" #include "flang/Optimizer/Support/InitFIR.h" #include "flang/Optimizer/Support/Utils.h" +#include "flang/Optimizer/Transforms/Passes.h" #include "flang/Parser/dump-parse-tree.h" #include "flang/Parser/parsing.h" #include "flang/Parser/provenance.h" @@ -300,6 +301,13 @@ mlir::PassManager pm((*mlirModule)->getName(), mlir::OpPassManager::Nesting::Implicit); pm.enableVerifier(/*verifyPasses=*/true); + + // Add OpenMP-related passes + if (ci.getInvocation().getFrontendOpts().features.IsEnabled( + Fortran::common::LanguageFeature::OpenMP)) { + pm.addPass(fir::createCaptureImplicitlyDeclareTargetPass()); + } + 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 @@ -16,7 +16,8 @@ AddDebugFoundation.cpp PolymorphicOpConversion.cpp LoopVersioning.cpp - + OMPCaptureImplicitDeclTar.cpp + DEPENDS FIRDialect FIROptTransformsPassIncGen diff --git a/flang/lib/Optimizer/Transforms/OMPCaptureImplicitDeclTar.cpp b/flang/lib/Optimizer/Transforms/OMPCaptureImplicitDeclTar.cpp new file mode 100644 --- /dev/null +++ b/flang/lib/Optimizer/Transforms/OMPCaptureImplicitDeclTar.cpp @@ -0,0 +1,101 @@ +#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/Operation.h" +#include "mlir/Pass/Pass.h" +#include +#include +#include + +namespace fir { +#define GEN_PASS_DEF_CAPTUREIMPLICITLYDECLARETARGETPASS +#include "flang/Optimizer/Transforms/Passes.h.inc" +} // namespace fir + +namespace { +class CaptureImplicitlyDeclareTargetPass + : public fir::impl::CaptureImplicitlyDeclareTargetPassBase< + CaptureImplicitlyDeclareTargetPass> { + + void markNestedFuncs(mlir::omp::DeclareTargetDeviceType parentDevTy, + mlir::omp::DeclareTargetCaptureClause parentCapClause, + mlir::Operation *currOp, + std::vector visited, + mlir::ModuleOp moduleOp) { + if (auto currFuncOp = llvm::dyn_cast(currOp)) { + for (auto v : visited) { + if (currFuncOp.getSymName() == v) + return; + } + visited.push_back(currFuncOp.getSymName()); + } + + 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 = + moduleOp.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, + moduleOp); + } + } + } + }); + } + + void runOnOperation() override { + mlir::ModuleOp moduleOp = getOperation(); + for (auto functionOp : moduleOp.getOps()) { + auto declareTargetOp = llvm::dyn_cast( + functionOp.getOperation()); + if (declareTargetOp.isDeclareTarget()) { + std::vector visited; + markNestedFuncs(declareTargetOp.getDeclareTargetDeviceType(), + declareTargetOp.getDeclareTargetCaptureClause(), + functionOp, visited, moduleOp); + } + } + + // 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 + moduleOp->walk([&](mlir::omp::TargetOp tarOp) { + std::vector visited; + markNestedFuncs(mlir::omp::DeclareTargetDeviceType::nohost, + mlir::omp::DeclareTargetCaptureClause::to, tarOp, visited, + moduleOp); + }); + } +}; + +} // namespace + +namespace fir { +std::unique_ptr> +createCaptureImplicitlyDeclareTargetPass() { + return std::make_unique(); +} +} // namespace fir \ No newline at end of file diff --git a/flang/test/Lower/OpenMP/omp-declare-target-implicit-func-and-subr-cap.f90 b/flang/test/Lower/OpenMP/omp-declare-target-implicit-func-and-subr-cap.f90 new file mode 100644 --- /dev/null +++ b/flang/test/Lower/OpenMP/omp-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/omp-declare-target-implicit-tarop-cap.f90 b/flang/test/Lower/OpenMP/omp-declare-target-implicit-tarop-cap.f90 new file mode 100644 --- /dev/null +++ b/flang/test/Lower/OpenMP/omp-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 \ No newline at end of file