[flang-commits] [flang] [flang][OpenMP] Support delayed privatisation for composite distribute simd (PR #151169)

[Kajetan Puchalski via flang-commits]

https://github.com/mrkajetanp updated https://github.com/llvm/llvm-project/pull/151169

>From cc03d947da62a7bfd6fc77590c2f366044d23740 Mon Sep 17 00:00:00 2001
From: Kajetan Puchalski <kajetan.puchalski at arm.com>
Date: Tue, 29 Jul 2025 15:12:07 +0000
Subject: [PATCH] [flang][OpenMP] Support delayed privatisation for composite
 distribute simd

Implement the lowering for delayed privatisation for composite
"distibute simd"constructs. Fixes new crashes previously masked
by simd information on composite constructs being ignored

Signed-off-by: Kajetan Puchalski <kajetan.puchalski at arm.com>
---
 flang/lib/Lower/OpenMP/OpenMP.cpp           | 23 ++++++++++++++-------
 flang/test/Lower/OpenMP/distribute-simd.f90 | 21 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 12089d6caa5fe..6cbb257a1d933 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -3145,11 +3145,16 @@ static mlir::omp::DistributeOp genCompositeDistributeSimd(
   genSimdClauses(converter, semaCtx, simdItem->clauses, loc, simdClauseOps,
                  simdReductionSyms);
 
-  // TODO: Support delayed privatization.
-  DataSharingProcessor dsp(converter, semaCtx, simdItem->clauses, eval,
-                           /*shouldCollectPreDeterminedSymbols=*/true,
-                           /*useDelayedPrivatization=*/false, symTable);
-  dsp.processStep1();
+  DataSharingProcessor distributeItemDSP(
+      converter, semaCtx, distributeItem->clauses, eval,
+      /*shouldCollectPreDeterminedSymbols=*/false,
+      /*useDelayedPrivatization=*/true, symTable);
+  distributeItemDSP.processStep1(&distributeClauseOps);
+
+  DataSharingProcessor simdItemDSP(converter, semaCtx, simdItem->clauses, eval,
+                                   /*shouldCollectPreDeterminedSymbols=*/true,
+                                   /*useDelayedPrivatization=*/true, symTable);
+  simdItemDSP.processStep1(&simdClauseOps);
 
   // Pass the innermost leaf construct's clauses because that's where COLLAPSE
   // is placed by construct decomposition.
@@ -3160,13 +3165,15 @@ static mlir::omp::DistributeOp genCompositeDistributeSimd(
 
   // Operation creation.
   EntryBlockArgs distributeArgs;
-  // TODO: Add private syms and vars.
+  distributeArgs.priv.syms = distributeItemDSP.getDelayedPrivSymbols();
+  distributeArgs.priv.vars = distributeClauseOps.privateVars;
   auto distributeOp = genWrapperOp<mlir::omp::DistributeOp>(
       converter, loc, distributeClauseOps, distributeArgs);
   distributeOp.setComposite(/*val=*/true);
 
   EntryBlockArgs simdArgs;
-  // TODO: Add private syms and vars.
+  simdArgs.priv.syms = simdItemDSP.getDelayedPrivSymbols();
+  simdArgs.priv.vars = simdClauseOps.privateVars;
   simdArgs.reduction.syms = simdReductionSyms;
   simdArgs.reduction.vars = simdClauseOps.reductionVars;
   auto simdOp =
@@ -3176,7 +3183,7 @@ static mlir::omp::DistributeOp genCompositeDistributeSimd(
   genLoopNestOp(converter, symTable, semaCtx, eval, loc, queue, simdItem,
                 loopNestClauseOps, iv,
                 {{distributeOp, distributeArgs}, {simdOp, simdArgs}},
-                llvm::omp::Directive::OMPD_distribute_simd, dsp);
+                llvm::omp::Directive::OMPD_distribute_simd, simdItemDSP);
   return distributeOp;
 }
 
diff --git a/flang/test/Lower/OpenMP/distribute-simd.f90 b/flang/test/Lower/OpenMP/distribute-simd.f90
index a43600143dc49..d0316d1a136ab 100644
--- a/flang/test/Lower/OpenMP/distribute-simd.f90
+++ b/flang/test/Lower/OpenMP/distribute-simd.f90
@@ -7,7 +7,7 @@
 subroutine distribute_simd_aligned(A)
   use iso_c_binding
   type(c_ptr) :: A
-  
+
   !$omp teams
 
   ! CHECK:      omp.distribute
@@ -57,3 +57,22 @@ subroutine distribute_simd_simdlen()
 
   !$omp end teams
 end subroutine distribute_simd_simdlen
+
+! CHECK-LABEL: func.func @_QPdistribute_simd_private(
+subroutine distribute_simd_private()
+  integer, allocatable :: tmp
+  ! CHECK:      omp.teams
+  !$omp teams
+  ! CHECK:      omp.distribute
+  ! CHECK:      omp.simd
+  ! CHECK-SAME: private(@[[PRIV_BOX_SYM:.*]] %{{.*}} -> %[[PRIV_BOX:.*]], @[[PRIV_IVAR_SYM:.*]] %{{.*}} -> %[[PRIV_IVAR:.*]] : !fir.ref<!fir.box<!fir.heap<i32>>>, !fir.ref<i32>)
+  ! CHECK-NEXT: omp.loop_nest (%[[IVAR:.*]]) : i32
+  !$omp distribute simd private(tmp)
+  do index_ = 1, 10
+  ! CHECK:      %[[PRIV_BOX_DECL:.*]]:2 = hlfir.declare %[[PRIV_BOX]]
+  ! CHECK:      %[[PRIV_IVAR_DECL:.*]]:2 = hlfir.declare %[[PRIV_IVAR]]
+  ! CHECK:      hlfir.assign %[[IVAR]] to %[[PRIV_IVAR_DECL]]#0
+  end do
+  !$omp end distribute simd
+  !$omp end teams
+end subroutine distribute_simd_private