[flang][OpenMP] Support delayed privatisation for composite distribute simd

[mrkajetanp]

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

Fixes #151159

[llvmbot]

@llvm/pr-subscribers-flang-openmp

@llvm/pr-subscribers-flang-fir-hlfir

Author: Kajetan Puchalski (mrkajetanp)

Changes

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

Fixes #151159

---

Full diff: https://github.com/llvm/llvm-project/pull/151169.diff

2 Files Affected:

• (modified) flang/lib/Lower/OpenMP/OpenMP.cpp (+15-8)
• (modified) flang/test/Lower/OpenMP/distribute-simd.f90 (+20-1)

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

[tblah]

LGTM. We already did delayed privatization for standalone SIMD. Did you test this against the public test suites?

[mrkajetanp]

LGTM. We already did delayed privatization for standalone SIMD. Did you test this against the public test suites?

I did. No changes in gfortran, fewer failures in Fujitsu :)