[Flang][OpenMP][MLIR] Add lowering from PFT to MLIR (FIR) for OpenMP declare target directive in Flang

[agozillon]

Summary

This patch adds PFT lowering for the OpenMP declare target directive
in Flang to the omp dialects declare target attribute, which currently
applies to function or global operations.

This builds on subsequent patches which add a semantic analysis pass
for implicit declare target capture and another which adds the declare
target MLIR attribute itself. Further lowering of the MLIR attribute comes
in additional patches on top of this.

[agozillon]

This is the PFT lowering to MLIR segment of the larger declare target patch: https://reviews.llvm.org/D146063 This segment has underwent no additions or changes during the split. Although, it will currently fail the buildbot as this patch isn't built on top of the other patches using arcanist, that's some arcanist wizardry I still need to learn.

[agozillon]

This comment has been deleted.

[agozillon]

Wrong patch! Was intended for the semantic analysis segment... so I've deleted the misinformation

[agozillon]

I believe this patch can go in without the semantic analysis pass patch (but it depends on the attribute patch), I just might need to remove some tests that depend on the pass and move them to the semantic pass patch for later addition, so if it would be possible to get a seperate review on this it would be appreciated.

[agozillon]

As the attribute has currently been accepted it would be great to get a review on this segment which lowers from the PFT to the attribute!

[kiranchandramohan]

Could you fix the CI issues?

[agozillon]

In D150329#4380628, @kiranchandramohan wrote:

Could you fix the CI issues?

No problem at all!

[agozillon]

• [Flang][OpenMP][MLIR] Rebase and remove tests dependent on other components

[agozillon]

latest patch should hopefully appease the CI, it involved a rebase and removal of some tests that will fail without the semantic analysis pass (and are now already tested in that patch).

[kiranchandramohan]

Looks OK.

In general, can you spell the types unless the type is on the RHS?

Variables in the main program have an implicit save and can hence be treated as globals. The following code will currently crash. I believe we have the same issue with threadprivate as well. If you can catch this and issue an appropriate TODO error that would be great.

program main
  integer :: i
  !$omp declare target to(i)
end

flang/lib/Lower/OpenMP.cpp

2518

Can you add a Fortran common block example to test this?

2565

2566

Can you add an assert here?

2572

[agozillon]

In D150329#4385587, @kiranchandramohan wrote:

Looks OK.

In general, can you spell the types unless the type is on the RHS?

Variables in the main program have an implicit save and can hence be treated as globals. The following code will currently crash. I believe we have the same issue with threadprivate as well. If you can catch this and issue an appropriate TODO error that would be great.

program main
  integer :: i
  !$omp declare target to(i)
end

Working on these alterations at the moment, it might take a little longer than I initially expected as it seems common blocks aren't accepted in declare target clauses at the moment, at least the following segment doesn't work at the moment:

PROGRAM main
    REAL :: one = 1
    REAL :: two = 2
    COMMON /numbers/ one, two
    !$omp declare target(/numbers/)
END

It results in a semantic error at the moment. I believe the above segment of code is the correct way to provide a common block to a clause (specifying just "numbers" works, but it results in the incorrect symbol and I do not think that's the correct way to specify the name from examples I've seen and a little test with gfortran), but please do point out if it is incorrect, I'm rather unfamiliar with them.

So I'll have a look at how easy that is to fix, if it is the correct syntax.

[kiranchandramohan]

In D150329#4387497, @agozillon wrote:

In D150329#4385587, @kiranchandramohan wrote:

Looks OK.

In general, can you spell the types unless the type is on the RHS?

Variables in the main program have an implicit save and can hence be treated as globals. The following code will currently crash. I believe we have the same issue with threadprivate as well. If you can catch this and issue an appropriate TODO error that would be great.

program main
  integer :: i
  !$omp declare target to(i)
end

Working on these alterations at the moment, it might take a little longer than I initially expected as it seems common blocks aren't accepted in declare target clauses at the moment, at least the following segment doesn't work at the moment:

PROGRAM main
    REAL :: one = 1
    REAL :: two = 2
    COMMON /numbers/ one, two
    !$omp declare target(/numbers/)
END

It results in a semantic error at the moment. I believe the above segment of code is the correct way to provide a common block to a clause (specifying just "numbers" works, but it results in the incorrect symbol and I do not think that's the correct way to specify the name from examples I've seen and a little test with gfortran), but please do point out if it is incorrect, I'm rather unfamiliar with them.

So I'll have a look at how easy that is to fix, if it is the correct syntax.

That seems to be right. We have similar tests for threadprivate (https://github.com/llvm/llvm-project/blob/main/flang/test/Lower/OpenMP/threadprivate-commonblock.f90). You can check whether there is any special handling for threadprivate.

[agozillon]

Added declare target to the resolution pass in the following patch: https://reviews.llvm.org/D151993 this resolves the symbol for common block (and any other symbols that may need similarly resolved) when captured by the directive, I'll add common block tests in this patch as suggested and it'll test the resolution works appropriately by side effect.

[agozillon]

• [Flang][OpenMP][MLIR] Rebase and remove tests dependent on other components
• [Flang][OpenMP][MLIR] Apply reviewer feedback where possible

[agozillon]

• [Flang][OpenMP][Test] Add missing space

[agozillon]

Added the common block tests and applied the feedback you mentioned! Please have a look incase you spot something I missed from your feedback.

I haven't added code to catch the case you mentioned, e.g.

program main
  integer :: i
  !$omp declare target to(i)
end

As it's currently caught by the assert that's been added and I'm not too sure how to catch this case from inside of the lowering (so I've added some inline comments above the assert for the known cases it'll trigger).

However, threadprivate currently handles this case, it seems to do so by generating an internal fir.GlobalOp which has the symbol information, allowing it to be searched (unlike the fir.AllocaOp) amongst other things. So, it could be possible to do something similar for declare target, upgrade the fir.AllocaOp to include a fir.GlobalOp, alternatively, a search method for the bind or unique name could be made that functions similarly to findSymbol. In either case, that seems like it extends into a future follow up patch.

[kiranchandramohan]

LG. It will be great if someone in your team. can have a look too.

flang/lib/Lower/OpenMP.cpp

2580–2581

Convert this to a TODO. After fixing the Case 3, we can convert it to an assert. This is to avoid adding crashes from gfortran testsuite.

[agozillon]

• [Flang][OpenMP] Replace assert with TODO

[agozillon]

In D150329#4391821, @kiranchandramohan wrote:

LG. It will be great if someone in your team. can have a look too.

No problem, I'll see who's available on my team and have them give a final review the patch, thank you very much for your help and time @kiranchandramohan

[agozillon]

Opened this patch related to the TODO discussed within this patch: https://reviews.llvm.org/D152037 feel free to review and discuss it there for those interested.

No more spam from me tonight, sorry everyone that's subjected to my patches!

[skatrak]

Small nit, otherwise LGTM.

flang/lib/Lower/OpenMP.cpp

2529

Nit: I think in this file it's preferred object initialization using "=" rather than braces. So, to keep the same coding style, I'd recommend changing the initializations of spec, objectList, clauseList, toClause, linkClause and deviceClause below.

[agozillon]

Thank you both for the reviews!

flang/lib/Lower/OpenMP.cpp

2529

it seems fairly normal for unwrap to be used this way, elsewhere in the file as well, the get/get_if is a little different though. So I'll change those instances on commit but not the unwraps to keep the status quo

flang/lib/Lower/OpenMP.cpp

void Fortran::lower::genThreadprivateOp(
   }
 
   fir::ExtendedValue sexv = converter.getSymbolExtendedValue(sym);
   fir::ExtendedValue symThreadprivateExv =
       getExtendedValue(sexv, symThreadprivateValue);
   converter.bindSymbol(sym, symThreadprivateExv);
 }
 
+void handleDeclareTarget(Fortran::lower::AbstractConverter &converter,
+                         Fortran::lower::pft::Evaluation &eval,
+                         const Fortran::parser::OpenMPDeclareTargetConstruct
+                             &declareTargetConstruct) {
+  llvm::SmallVector<std::pair<mlir::omp::DeclareTargetCaptureClause,
+                              Fortran::semantics::Symbol>,
+                    0>
+      symbolAndClause;
+  auto findFuncAndVarSyms = [&](const Fortran::parser::OmpObjectList &objList,
+                                mlir::omp::DeclareTargetCaptureClause clause) {
+    for (const auto &ompObject : objList.v) {
+      Fortran::common::visit(
+          Fortran::common::visitors{
+              [&](const Fortran::parser::Designator &designator) {
+                if (const Fortran::parser::Name *name =
+                        getDesignatorNameIfDataRef(designator)) {
+                  symbolAndClause.push_back(
+                      std::make_pair(clause, *name->symbol));
+                }
+              },
+              [&](const Fortran::parser::Name &name) {
+                symbolAndClause.push_back(std::make_pair(clause, *name.symbol));
+              }},
+          ompObject.u);
+    }
+  };
+
+  const auto &spec{std::get<Fortran::parser::OmpDeclareTargetSpecifier>(
+      declareTargetConstruct.t)};
+  auto mod = converter.getFirOpBuilder().getModule();
+
+  // The default capture type
+  auto deviceType = Fortran::parser::OmpDeviceTypeClause::Type::Any;
+  if (const auto *objectList{
+          Fortran::parser::Unwrap<Fortran::parser::OmpObjectList>(spec.u)}) {
+    // Case: declare target(func, var1, var2)
+    findFuncAndVarSyms(*objectList, mlir::omp::DeclareTargetCaptureClause::to);
+  } else if (const auto *clauseList{
+                 Fortran::parser::Unwrap<Fortran::parser::OmpClauseList>(
+                     spec.u)}) {
+    if (clauseList->v.empty()) {
+      // Case: declare target, implicit capture of function
+      symbolAndClause.push_back(
+          std::make_pair(mlir::omp::DeclareTargetCaptureClause::to,
+                         eval.getOwningProcedure()->getSubprogramSymbol()));
+    }
+
+    for (const auto &clause : clauseList->v) {
+      if (const auto *toClause{
+              std::get_if<Fortran::parser::OmpClause::To>(&clause.u)}) {
+        // Case: declare target to(func, var1, var2)...
+        findFuncAndVarSyms(toClause->v,
+                           mlir::omp::DeclareTargetCaptureClause::to);
+      } else if (const auto *linkClause{
+                     std::get_if<Fortran::parser::OmpClause::Link>(
+                         &clause.u)}) {
+        // Case: declare target link(var1, var2)...
+        findFuncAndVarSyms(linkClause->v,
+                           mlir::omp::DeclareTargetCaptureClause::link);
+      } else if (const auto *deviceClause{
+                     std::get_if<Fortran::parser::OmpClause::DeviceType>(
+                         &clause.u)}) {
+        // Case: declare target ... device_type(any | host | nohost)
+        deviceType = deviceClause->v.v;
+      }
+    }
+  }
+
+  for (auto sym : symbolAndClause) {
+    auto *op = mod.lookupSymbol(converter.mangleName(std::get<1>(sym)));
+
+    auto declareTargetOp = dyn_cast<mlir::omp::DeclareTargetInterface>(op);
+    if (!declareTargetOp)
+      fir::emitFatalError(
+          converter.getCurrentLocation(),
+          "Attempt to apply declare target on unsupproted operation");
+
+    mlir::omp::DeclareTargetDeviceType newDeviceType;
+    switch (deviceType) {
+    case Fortran::parser::OmpDeviceTypeClause::Type::Nohost:
+      newDeviceType = mlir::omp::DeclareTargetDeviceType::nohost;
+      break;
+    case Fortran::parser::OmpDeviceTypeClause::Type::Host:
+      newDeviceType = mlir::omp::DeclareTargetDeviceType::host;
+      break;
+    case Fortran::parser::OmpDeviceTypeClause::Type::Any:
+      newDeviceType = mlir::omp::DeclareTargetDeviceType::any;
+      break;
+    }
+
+    // The function or global already has a declare target applied to it,
+    // very likely through implicit capture (usage in another declare
+    // target function/subroutine). It should be marked as any if it has
+    // been assigned both host and nohost, else we skip, as there is no
+    // change
+    if (declareTargetOp.isDeclareTarget()) {
+      if (declareTargetOp.getDeclareTargetDeviceType() != newDeviceType)
+        declareTargetOp.setDeclareTarget(
+            mlir::omp::DeclareTargetDeviceType::any, std::get<0>(sym));
+      continue;
+    }
+
+    declareTargetOp.setDeclareTarget(newDeviceType, std::get<0>(sym));
+  }
+}
+
 void Fortran::lower::genOpenMPDeclarativeConstruct(
     Fortran::lower::AbstractConverter &converter,
     Fortran::lower::pft::Evaluation &eval,
     const Fortran::parser::OpenMPDeclarativeConstruct &ompDeclConstruct) {
 
   std::visit(
       common::visitors{
           [&](const Fortran::parser::OpenMPDeclarativeAllocate
                   &declarativeAllocate) {
             TODO(converter.getCurrentLocation(), "OpenMPDeclarativeAllocate");
           },
           [&](const Fortran::parser::OpenMPDeclareReductionConstruct
                   &declareReductionConstruct) {
             TODO(converter.getCurrentLocation(),
                  "OpenMPDeclareReductionConstruct");
           },
           [&](const Fortran::parser::OpenMPDeclareSimdConstruct
                   &declareSimdConstruct) {
             TODO(converter.getCurrentLocation(), "OpenMPDeclareSimdConstruct");
           },
           [&](const Fortran::parser::OpenMPDeclareTargetConstruct
                   &declareTargetConstruct) {
-            TODO(converter.getCurrentLocation(),
-                 "OpenMPDeclareTargetConstruct");
+            handleDeclareTarget(converter, eval, declareTargetConstruct);
           },
           [&](const Fortran::parser::OpenMPRequiresConstruct
                   &requiresConstruct) {
             TODO(converter.getCurrentLocation(), "OpenMPRequiresConstruct");
           },
           [&](const Fortran::parser::OpenMPThreadprivate &threadprivate) {
             // The directive is lowered when instantiating the variable to
             // support the case of threadprivate variable declared in module.

flang/test/Lower/OpenMP/Todo/omp-declare-target.f90

-! This test checks lowering of OpenMP declare target Directive.
-
-// RUN: not flang-new -fc1 -emit-fir -fopenmp %s 2>&1 | FileCheck %s
-
-module mod1
-contains
-  subroutine sub()
-    integer :: x, y
-    // CHECK: not yet implemented: OpenMPDeclareTargetConstruct
-    !$omp declare target
-  end
-end module

flang/test/Lower/OpenMP/omp-declare-target-data.f90

+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s 
+!RUN: %flang_fc1 -emit-fir -fopenmp -fopenmp-is-device %s -o - | FileCheck %s
+
+module test_0
+    implicit none
+
+!CHECK-DAG: fir.global @_QMtest_0Edata_int {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : i32
+INTEGER :: data_int = 10
+!$omp declare target link(data_int)
+
+!CHECK-DAG: fir.global @_QMtest_0Earray_1d({{.*}}) {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : !fir.array<3xi32>
+INTEGER :: array_1d(3) = (/1,2,3/)
+!$omp declare target link(array_1d)
+
+!CHECK-DAG: fir.global @_QMtest_0Earray_2d({{.*}}) {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : !fir.array<2x2xi32> 
+INTEGER :: array_2d(2,2) = reshape((/1,2,3,4/), (/2,2/))
+!$omp declare target link(array_2d)
+
+!CHECK-DAG: fir.global @_QMtest_0Ept1 {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : !fir.box<!fir.ptr<i32>>
+INTEGER, POINTER :: pt1
+!$omp declare target link(pt1)
+
+!CHECK-DAG: fir.global @_QMtest_0Ept2_tar {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} target : i32
+INTEGER, TARGET :: pt2_tar = 5 
+!$omp declare target link(pt2_tar)
+
+!CHECK-DAG: fir.global @_QMtest_0Ept2 {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : !fir.box<!fir.ptr<i32>>
+INTEGER, POINTER :: pt2 => pt2_tar
+!$omp declare target link(pt2)
+
+!CHECK-DAG: fir.global @_QMtest_0Edata_int_to {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} : i32
+INTEGER :: data_int_to = 5
+!$omp declare target to(data_int_to)
+
+!CHECK-DAG: fir.global @_QMtest_0Edata_int_clauseless {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} : i32
+INTEGER :: data_int_clauseless = 1
+!$omp declare target(data_int_clauseless)
+
+!CHECK-DAG: fir.global @_QMtest_0Edata_extended_to_1 {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} : f32
+!CHECK-DAG: fir.global @_QMtest_0Edata_extended_to_2 {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>} : f32
+REAL :: data_extended_to_1 = 2
+REAL :: data_extended_to_2 = 3
+!$omp declare target to(data_extended_to_1, data_extended_to_2)
+
+!CHECK-DAG: fir.global @_QMtest_0Edata_extended_link_1 {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : f32
+!CHECK-DAG: fir.global @_QMtest_0Edata_extended_link_2 {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (link)>} : f32
+REAL :: data_extended_link_1 = 2
+REAL :: data_extended_link_2 = 3
+!$omp declare target link(data_extended_link_1, data_extended_link_2)
+
+contains
+end module test_0

flang/test/Lower/OpenMP/omp-declare-target-func-and-subr.f90

+!RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
+
+! Check specification valid forms of declare target with functions 
+! utilising device_type and to clauses as well as the default 
+! zero clause declare target
+
+! CHECK-LABEL: func.func @_QPfunc_t_device()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}
+FUNCTION FUNC_T_DEVICE() RESULT(I)
+!$omp declare target to(FUNC_T_DEVICE) device_type(nohost)
+    INTEGER :: I
+    I = 1
+END FUNCTION FUNC_T_DEVICE
+
+! CHECK-LABEL: func.func @_QPfunc_t_host()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>{{.*}}
+FUNCTION FUNC_T_HOST() RESULT(I)
+!$omp declare target to(FUNC_T_HOST) device_type(host)
+    INTEGER :: I
+    I = 1
+END FUNCTION FUNC_T_HOST
+
+! CHECK-LABEL: func.func @_QPfunc_t_any()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+FUNCTION FUNC_T_ANY() RESULT(I)
+!$omp declare target to(FUNC_T_ANY) device_type(any)
+    INTEGER :: I
+    I = 1
+END FUNCTION FUNC_T_ANY
+
+! CHECK-LABEL: func.func @_QPfunc_default_t_any()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+FUNCTION FUNC_DEFAULT_T_ANY() RESULT(I)
+!$omp declare target to(FUNC_DEFAULT_T_ANY)
+    INTEGER :: I
+    I = 1
+END FUNCTION FUNC_DEFAULT_T_ANY
+
+! CHECK-LABEL: func.func @_QPfunc_default_any()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+FUNCTION FUNC_DEFAULT_ANY() RESULT(I)
+!$omp declare target
+    INTEGER :: I
+    I = 1
+END FUNCTION FUNC_DEFAULT_ANY
+
+! CHECK-LABEL: func.func @_QPfunc_default_extendedlist()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+FUNCTION FUNC_DEFAULT_EXTENDEDLIST() RESULT(I)
+!$omp declare target(FUNC_DEFAULT_EXTENDEDLIST)
+    INTEGER :: I
+    I = 1
+END FUNCTION FUNC_DEFAULT_EXTENDEDLIST
+
+!! -----
+
+! Check specification valid forms of declare target with subroutines 
+! utilising device_type and to clauses as well as the default 
+! zero clause declare target
+
+! CHECK-LABEL: func.func @_QPsubr_t_device()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}
+SUBROUTINE SUBR_T_DEVICE()
+!$omp declare target to(SUBR_T_DEVICE) device_type(nohost)
+END
+
+! CHECK-LABEL: func.func @_QPsubr_t_host()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>{{.*}}
+SUBROUTINE SUBR_T_HOST()
+!$omp declare target to(SUBR_T_HOST) device_type(host)
+END
+
+! CHECK-LABEL: func.func @_QPsubr_t_any()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+SUBROUTINE SUBR_T_ANY()
+!$omp declare target to(SUBR_T_ANY) device_type(any)
+END
+
+! CHECK-LABEL: func.func @_QPsubr_default_t_any()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+SUBROUTINE SUBR_DEFAULT_T_ANY()
+!$omp declare target to(SUBR_DEFAULT_T_ANY)
+END
+
+! CHECK-LABEL: func.func @_QPsubr_default_any()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+SUBROUTINE SUBR_DEFAULT_ANY()
+!$omp declare target
+END
+
+! CHECK-LABEL: func.func @_QPsubr_default_extendedlist()
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>{{.*}}
+SUBROUTINE SUBR_DEFAULT_EXTENDEDLIST()
+!$omp declare target(SUBR_DEFAULT_EXTENDEDLIST)
+END
+
+!! -----
+
+! CHECK-LABEL: func.func @_QPrecursive_declare_target
+! CHECK-SAME: {{.*}}attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>{{.*}}
+RECURSIVE FUNCTION RECURSIVE_DECLARE_TARGET(INCREMENT) RESULT(K)
+!$omp declare target to(RECURSIVE_DECLARE_TARGET) device_type(nohost)
+    INTEGER :: INCREMENT, K
+    IF (INCREMENT == 10) THEN
+        K = INCREMENT
+    ELSE
+        K = RECURSIVE_DECLARE_TARGET(INCREMENT + 1)
+    END IF
+END FUNCTION RECURSIVE_DECLARE_TARGET