[flang-commits] [flang] [flang][OpenMP] Enable lastprivate on simd (PR #93786)

[via flang-commits]

https://github.com/NimishMishra created https://github.com/llvm/llvm-project/pull/93786

This PR enables the lowering of lastprivate when defined on simd construct

>From ebf3eb90e51cd9db2553af8b780583030eab33d0 Mon Sep 17 00:00:00 2001
From: Nimish Mishra <neelam.nimish at gmail.com>
Date: Thu, 30 May 2024 14:04:46 +0530
Subject: [PATCH] [flang][OpenMP] Enable lastprivate on simd

---
 .../lib/Lower/OpenMP/DataSharingProcessor.cpp |  6 ++---
 flang/test/Lower/OpenMP/simd.f90              | 27 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
index 557a9685024c5..ab5f2b8ac2429 100644
--- a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
@@ -199,12 +199,12 @@ void DataSharingProcessor::insertLastPrivateCompare(mlir::Operation *op) {
   for (const omp::Clause &clause : clauses) {
     if (clause.id != llvm::omp::OMPC_lastprivate)
       continue;
-    // TODO: Add lastprivate support for simd construct
-    if (mlir::isa<mlir::omp::WsloopOp>(op)) {
+    if (mlir::isa<mlir::omp::WsloopOp>(op) or
+        mlir::isa<mlir::omp::SimdOp>(op)) {
       // Update the original variable just before exiting the worksharing
       // loop. Conversion as follows:
       //
-      // omp.wsloop {             omp.wsloop {
+      // omp.wsloop / omp.simd {    omp.wsloop / omp.simd {
       //   omp.loop_nest {          omp.loop_nest {
       //     ...                      ...
       //     store          ===>      store
diff --git a/flang/test/Lower/OpenMP/simd.f90 b/flang/test/Lower/OpenMP/simd.f90
index 223b248b79348..2f9720d776490 100644
--- a/flang/test/Lower/OpenMP/simd.f90
+++ b/flang/test/Lower/OpenMP/simd.f90
@@ -182,3 +182,30 @@ subroutine simd_with_collapse_clause(n)
   end do
   !$OMP END SIMD
 end subroutine
+
+!CHECK: func.func @_QPlastprivate_with_simd() {
+subroutine lastprivate_with_simd
+
+!CHECK: %[[VAR_SUM:.*]] = fir.alloca f32 {bindc_name = "sum", uniq_name = "_QFlastprivate_with_simdEsum"}
+!CHECK: %[[VAR_SUM_DECLARE:.*]]:2 = hlfir.declare %[[VAR_SUM]] {{.*}}
+!CHECK: %[[VAR_SUM_PINNED:.*]] = fir.alloca f32 {bindc_name = "sum", pinned, uniq_name = "_QFlastprivate_with_simdEsum"}
+!CHECK: %[[VAR_SUM_PINNED_DECLARE:.*]]:2 = hlfir.declare %[[VAR_SUM_PINNED]] {{.*}}
+  implicit none
+  integer :: i
+  real :: sum
+
+!CHECK: omp.simd {
+!CHECK: omp.loop_nest (%[[ARG:.*]]) : i32 = ({{.*}} to ({{.*}}) inclusive step ({{.*}}) {
+!CHECK: %[[ADD_RESULT:.*]] = arith.addi {{.*}}
+!CHECK: %[[ADD_RESULT_CONVERT:.*]] = fir.convert %[[ADD_RESULT]] : (i32) -> f32
+!CHECK: hlfir.assign %[[ADD_RESULT_CONVERT]] to %[[VAR_SUM_PINNED_DECLARE]]#0 : f32, !fir.ref<f32> 
+  !$omp simd lastprivate(sum)
+  do i = 1, 100
+     sum = i + 1
+  end do
+!CHECK: %[[SELECT_RESULT:.*]] = arith.select {{.*}}, {{.*}}, {{.*}} : i1
+!CHECK: fir.if %[[SELECT_RESULT]] {
+!CHECK: %[[LOADED_SUM:.*]] = fir.load %[[VAR_SUM_PINNED_DECLARE]]#0 : !fir.ref<f32>
+!CHECK: hlfir.assign %[[LOADED_SUM]] to %[[VAR_SUM_DECLARE]]#0 temporary_lhs : f32, !fir.ref<f32>
+!CHECK: }
+end subroutine