[flang-commits] [flang] [Flang][OpenMP] Support teams reductions lowering (PR #122683)

[Sergio Afonso via flang-commits]

https://github.com/skatrak created https://github.com/llvm/llvm-project/pull/122683

This patch adds PFT to MLIR lowering of teams reductions. Since there is still no MLIR to LLVM IR translation implemented, compilation of programs including these constructs will still trigger not-yet-implemented errors.

>From 5c34eaefc8bb4c6fc210abef138a2715838a4a6f Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Mon, 13 Jan 2025 10:12:11 +0000
Subject: [PATCH] [Flang][OpenMP] Support teams reductions lowering

This patch adds PFT to MLIR lowering of teams reductions. Since there is still
no MLIR to LLVM IR translation implemented, compilation of programs including
these constructs will still trigger not-yet-implemented errors.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp             | 33 ++++++++++++++-----
 .../Lower/OpenMP/Todo/reduction-teams.f90     | 12 -------
 flang/test/Lower/OpenMP/reduction-teams.f90   | 18 ++++++++++
 3 files changed, 42 insertions(+), 21 deletions(-)
 delete mode 100644 flang/test/Lower/OpenMP/Todo/reduction-teams.f90
 create mode 100644 flang/test/Lower/OpenMP/reduction-teams.f90

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index cd4b25a17722c1..ddbca3c7291ae6 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -1336,19 +1336,18 @@ static void genWorkshareClauses(lower::AbstractConverter &converter,
   cp.processNowait(clauseOps);
 }
 
-static void genTeamsClauses(lower::AbstractConverter &converter,
-                            semantics::SemanticsContext &semaCtx,
-                            lower::StatementContext &stmtCtx,
-                            const List<Clause> &clauses, mlir::Location loc,
-                            mlir::omp::TeamsOperands &clauseOps) {
+static void genTeamsClauses(
+    lower::AbstractConverter &converter, semantics::SemanticsContext &semaCtx,
+    lower::StatementContext &stmtCtx, const List<Clause> &clauses,
+    mlir::Location loc, mlir::omp::TeamsOperands &clauseOps,
+    llvm::SmallVectorImpl<const semantics::Symbol *> &reductionSyms) {
   ClauseProcessor cp(converter, semaCtx, clauses);
   cp.processAllocate(clauseOps);
   cp.processIf(llvm::omp::Directive::OMPD_teams, clauseOps);
   cp.processNumTeams(stmtCtx, clauseOps);
   cp.processThreadLimit(stmtCtx, clauseOps);
+  cp.processReduction(loc, clauseOps, reductionSyms);
   // TODO Support delayed privatization.
-
-  cp.processTODO<clause::Reduction>(loc, llvm::omp::Directive::OMPD_teams);
 }
 
 static void genWsloopClauses(
@@ -2015,13 +2014,29 @@ genTeamsOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
            mlir::Location loc, const ConstructQueue &queue,
            ConstructQueue::const_iterator item) {
   lower::StatementContext stmtCtx;
+
   mlir::omp::TeamsOperands clauseOps;
-  genTeamsClauses(converter, semaCtx, stmtCtx, item->clauses, loc, clauseOps);
+  llvm::SmallVector<const semantics::Symbol *> reductionSyms;
+  genTeamsClauses(converter, semaCtx, stmtCtx, item->clauses, loc, clauseOps,
+                  reductionSyms);
+
+  EntryBlockArgs args;
+  // TODO: Add private syms and vars.
+  args.reduction.syms = reductionSyms;
+  args.reduction.vars = clauseOps.reductionVars;
+
+  auto genRegionEntryCB = [&](mlir::Operation *op) {
+    genEntryBlock(converter.getFirOpBuilder(), args, op->getRegion(0));
+    bindEntryBlockArgs(
+        converter, llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(op), args);
+    return llvm::to_vector(args.getSyms());
+  };
 
   return genOpWithBody<mlir::omp::TeamsOp>(
       OpWithBodyGenInfo(converter, symTable, semaCtx, loc, eval,
                         llvm::omp::Directive::OMPD_teams)
-          .setClauses(&item->clauses),
+          .setClauses(&item->clauses)
+          .setGenRegionEntryCb(genRegionEntryCB),
       queue, item, clauseOps);
 }
 
diff --git a/flang/test/Lower/OpenMP/Todo/reduction-teams.f90 b/flang/test/Lower/OpenMP/Todo/reduction-teams.f90
deleted file mode 100644
index db4839593c7e7f..00000000000000
--- a/flang/test/Lower/OpenMP/Todo/reduction-teams.f90
+++ /dev/null
@@ -1,12 +0,0 @@
-! RUN: %not_todo_cmd bbc -emit-fir -fopenmp -o - %s 2>&1 | FileCheck %s
-! RUN: %not_todo_cmd %flang_fc1 -emit-fir -fopenmp -o - %s 2>&1 | FileCheck %s
-
-! CHECK: not yet implemented: Unhandled clause REDUCTION in TEAMS construct
-subroutine reduction_teams()
-  integer :: i
-  i = 0
-
-  !$omp teams reduction(+:i)
-  i = i + 1
-  !$omp end teams
-end subroutine reduction_teams
diff --git a/flang/test/Lower/OpenMP/reduction-teams.f90 b/flang/test/Lower/OpenMP/reduction-teams.f90
new file mode 100644
index 00000000000000..6997e774c2d425
--- /dev/null
+++ b/flang/test/Lower/OpenMP/reduction-teams.f90
@@ -0,0 +1,18 @@
+! RUN: bbc -emit-hlfir -fopenmp -o - %s 2>&1 | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -o - %s 2>&1 | FileCheck %s
+
+! CHECK: omp.declare_reduction @[[RED:.*]] : i32 init {
+
+! CHECK: func.func @_QPreduction_teams() {
+subroutine reduction_teams()
+  integer :: i
+  i = 0
+
+  ! CHECK: omp.teams reduction(@[[RED]] %{{.*}}#0 -> %[[PRIV_I:.*]] : !fir.ref<i32>) {
+  !$omp teams reduction(+:i)
+    ! CHECK: %[[DECL_I:.*]]:2 = hlfir.declare %[[PRIV_I]]
+    ! CHECK: %{{.*}} = fir.load %[[DECL_I]]#0 : !fir.ref<i32>
+    ! CHECK: hlfir.assign %{{.*}} to %[[DECL_I]]#0 : i32, !fir.ref<i32>
+    i = i + 1
+  !$omp end teams
+end subroutine reduction_teams