[flang-commits] [flang] [mlir] [flang][MLIR] Support delayed privatization for `wsloop` (PFT -> MLIR) (PR #118271)

[Kareem Ergawy via flang-commits]

https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/118271

Adds PFT to MLIR lowering for delayed privatization of `omp.wsloop` ops. Lowering to LLVM IR will be added in a later PR.

>From cd0b6e6a7f542adf7cdc1c4cfa136e26b29e4039 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Sun, 1 Dec 2024 22:48:23 -0600
Subject: [PATCH] [flang][MLIR] Support delayed privatization for `wsloop` (PFT
 -> MLIR)

Adds PFT to MLIR lowering for delayed privatization of `omp.wsloop` ops.
Lowering to LLVM IR will be added in a later PR.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp             |  8 ++--
 .../OpenMP/DelayedPrivatization/wsloop.f90    | 42 +++++++++++++++++++
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  |  4 +-
 3 files changed, 48 insertions(+), 6 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/DelayedPrivatization/wsloop.f90

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 3bb43b766bcebf..d93b4ae8e50944 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2042,11 +2042,10 @@ static void genStandaloneDo(lower::AbstractConverter &converter,
   genWsloopClauses(converter, semaCtx, stmtCtx, item->clauses, loc,
                    wsloopClauseOps, wsloopReductionSyms);
 
-  // TODO: Support delayed privatization.
   DataSharingProcessor dsp(converter, semaCtx, item->clauses, eval,
                            /*shouldCollectPreDeterminedSymbols=*/true,
-                           /*useDelayedPrivatization=*/false, &symTable);
-  dsp.processStep1();
+                           enableDelayedPrivatizationStaging, &symTable);
+  dsp.processStep1(&wsloopClauseOps);
 
   mlir::omp::LoopNestOperands loopNestClauseOps;
   llvm::SmallVector<const semantics::Symbol *> iv;
@@ -2054,7 +2053,8 @@ static void genStandaloneDo(lower::AbstractConverter &converter,
                      loopNestClauseOps, iv);
 
   EntryBlockArgs wsloopArgs;
-  // TODO: Add private syms and vars.
+  wsloopArgs.priv.syms = dsp.getDelayedPrivSymbols();
+  wsloopArgs.priv.vars = wsloopClauseOps.privateVars;
   wsloopArgs.reduction.syms = wsloopReductionSyms;
   wsloopArgs.reduction.vars = wsloopClauseOps.reductionVars;
   auto wsloopOp = genWrapperOp<mlir::omp::WsloopOp>(
diff --git a/flang/test/Lower/OpenMP/DelayedPrivatization/wsloop.f90 b/flang/test/Lower/OpenMP/DelayedPrivatization/wsloop.f90
new file mode 100644
index 00000000000000..66fd120085c782
--- /dev/null
+++ b/flang/test/Lower/OpenMP/DelayedPrivatization/wsloop.f90
@@ -0,0 +1,42 @@
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -mmlir --openmp-enable-delayed-privatization-staging \
+! RUN:   -o - %s 2>&1 | FileCheck %s
+! RUN: bbc -emit-hlfir -fopenmp --openmp-enable-delayed-privatization-staging -o - %s 2>&1 \
+! RUN:   | FileCheck %s
+
+subroutine wsloop_private
+    implicit none
+    integer :: x, i
+
+    !$omp parallel do firstprivate(x)
+    do i = 0, 10
+      x = x + i
+    end do
+end subroutine wsloop_private
+
+! CHECK: omp.private {type = private} @[[I_PRIVATIZER:.*i_private_ref_i32]]
+! CHECK: omp.private {type = firstprivate} @[[X_PRIVATIZER:.*x_firstprivate_ref_i32]]
+
+! CHECK: func.func @{{.*}}() {
+! CHECK:   %[[I_DECL:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "{{.*}}i"}
+! CHECK:   %[[X_DECL:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "{{.*}}x"}
+
+! CHECK:   omp.parallel {
+! CHECK:     omp.wsloop private(
+! CHECK-SAME:  @[[X_PRIVATIZER]] %[[X_DECL]]#0 -> %[[X_ARG:[^[:space:]]+]],
+! CHECK-SAME:  @[[I_PRIVATIZER]] %[[I_DECL]]#0 -> %[[I_ARG:.*]] : {{.*}}) {
+
+! CHECK:       omp.loop_nest (%[[IV:.*]]) : i32 = {{.*}} {
+! CHECK:         %[[X_PRIV_DECL:.*]]:2 = hlfir.declare %[[X_ARG]] {uniq_name = "{{.*}}x"}
+! CHECK:         %[[I_PRIV_DECL:.*]]:2 = hlfir.declare %[[I_ARG]] {uniq_name = "{{.*}}i"}
+! CHECK:         fir.store %[[IV]] to %[[I_PRIV_DECL]]#1
+! CHECK:         %[[X_VAL:.*]] = fir.load %[[X_PRIV_DECL]]#0
+! CHECK:         %[[I_VAL:.*]] = fir.load %[[I_PRIV_DECL]]#0
+! CHECK:         %[[ADD_VAL:.*]] = arith.addi %[[X_VAL]], %[[I_VAL]]
+! CHECK:         hlfir.assign %[[ADD_VAL]] to %[[X_PRIV_DECL]]#0
+! CHECK:         omp.yield
+! CHECK:       }
+! CHECK:     }
+
+! CHECK:     omp.terminator
+! CHECK:   }
+! CHECK: }
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 8c5f79a49a334f..6d7dbbf58bbda7 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2066,8 +2066,8 @@ void WsloopOp::build(OpBuilder &builder, OperationState &state,
       builder, state,
       /*allocate_vars=*/{}, /*allocator_vars=*/{}, clauses.linearVars,
       clauses.linearStepVars, clauses.nowait, clauses.order, clauses.orderMod,
-      clauses.ordered, /*private_vars=*/{}, /*private_syms=*/nullptr,
-      clauses.reductionVars,
+      clauses.ordered, clauses.privateVars,
+      makeArrayAttr(ctx, clauses.privateSyms), clauses.reductionVars,
       makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
       makeArrayAttr(ctx, clauses.reductionSyms), clauses.scheduleKind,
       clauses.scheduleChunk, clauses.scheduleMod, clauses.scheduleSimd);