[llvm-branch-commits] [flang] Revert "[Flang][OpenMP] Disable lowering of omp.simd reductions in co… (PR #113683)

[Sergio Afonso via llvm-branch-commits]

https://github.com/skatrak updated https://github.com/llvm/llvm-project/pull/113683

>From cf57ecd0984bcb2335e8ecfaffdca600a5c7cf1b Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Fri, 25 Oct 2024 12:07:22 +0100
Subject: [PATCH] Revert "[Flang][OpenMP] Disable lowering of omp.simd
 reductions in composites (#112686)"

Lowering of reductions in composite operations can now be re-enabled, since previous commits in this PR stack fix the MLIR representation produced and it no longer triggers a compiler crash during translation to LLVM IR.

This reverts commit c44860c8d2582abd88794267b4fa0fa953bbef80.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp       | 20 ++++++--------------
 flang/test/Lower/OpenMP/wsloop-simd.f90 | 21 +++++++++++++++++++++
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 329cbf3d7539f5..4f9e2347308aa1 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2246,12 +2246,6 @@ static void genCompositeDistributeParallelDoSimd(
   genSimdClauses(converter, semaCtx, simdItem->clauses, loc, simdClauseOps,
                  simdReductionSyms);
 
-  // TODO: Remove this after omp.simd reductions on composite constructs are
-  // supported.
-  simdClauseOps.reductionVars.clear();
-  simdClauseOps.reductionByref.clear();
-  simdClauseOps.reductionSyms.clear();
-
   mlir::omp::LoopNestOperands loopNestClauseOps;
   llvm::SmallVector<const semantics::Symbol *> iv;
   genLoopNestClauses(converter, semaCtx, eval, simdItem->clauses, loc,
@@ -2273,7 +2267,9 @@ static void genCompositeDistributeParallelDoSimd(
   wsloopOp.setComposite(/*val=*/true);
 
   EntryBlockArgs simdArgs;
-  // TODO: Add private and reduction syms and vars.
+  // TODO: Add private syms and vars.
+  simdArgs.reduction.syms = simdReductionSyms;
+  simdArgs.reduction.vars = simdClauseOps.reductionVars;
   auto simdOp =
       genWrapperOp<mlir::omp::SimdOp>(converter, loc, simdClauseOps, simdArgs);
   simdOp.setComposite(/*val=*/true);
@@ -2366,12 +2362,6 @@ static void genCompositeDoSimd(lower::AbstractConverter &converter,
   genSimdClauses(converter, semaCtx, simdItem->clauses, loc, simdClauseOps,
                  simdReductionSyms);
 
-  // TODO: Remove this after omp.simd reductions on composite constructs are
-  // supported.
-  simdClauseOps.reductionVars.clear();
-  simdClauseOps.reductionByref.clear();
-  simdClauseOps.reductionSyms.clear();
-
   // TODO: Support delayed privatization.
   DataSharingProcessor dsp(converter, semaCtx, simdItem->clauses, eval,
                            /*shouldCollectPreDeterminedSymbols=*/true,
@@ -2395,7 +2385,9 @@ static void genCompositeDoSimd(lower::AbstractConverter &converter,
   wsloopOp.setComposite(/*val=*/true);
 
   EntryBlockArgs simdArgs;
-  // TODO: Add private and reduction syms and vars.
+  // TODO: Add private syms and vars.
+  simdArgs.reduction.syms = simdReductionSyms;
+  simdArgs.reduction.vars = simdClauseOps.reductionVars;
   auto simdOp =
       genWrapperOp<mlir::omp::SimdOp>(converter, loc, simdClauseOps, simdArgs);
   simdOp.setComposite(/*val=*/true);
diff --git a/flang/test/Lower/OpenMP/wsloop-simd.f90 b/flang/test/Lower/OpenMP/wsloop-simd.f90
index 899ab59714f144..49a9a523e11fe7 100644
--- a/flang/test/Lower/OpenMP/wsloop-simd.f90
+++ b/flang/test/Lower/OpenMP/wsloop-simd.f90
@@ -45,3 +45,24 @@ subroutine do_simd_simdlen()
     end do
   !$omp end do simd
 end subroutine do_simd_simdlen
+
+! CHECK-LABEL: func.func @_QPdo_simd_reduction(
+subroutine do_simd_reduction()
+  integer :: sum
+  sum = 0
+  ! CHECK:      omp.wsloop
+  ! CHECK-SAME: reduction(@[[RED_SYM:.*]] %{{.*}} -> %[[RED_OUTER:.*]] : !fir.ref<i32>)
+  ! CHECK-NEXT: omp.simd
+  ! CHECK-SAME: reduction(@[[RED_SYM]] %[[RED_OUTER]] -> %[[RED_INNER:.*]] : !fir.ref<i32>)
+  ! CHECK-NEXT: omp.loop_nest
+  ! CHECK:      %[[RED_DECL:.*]]:2 = hlfir.declare %[[RED_INNER]]
+  ! CHECK:      %[[RED:.*]] = fir.load %[[RED_DECL]]#0 : !fir.ref<i32>
+  ! CHECK:      %[[RESULT:.*]] = arith.addi %[[RED]], %{{.*}} : i32
+  ! CHECK:      hlfir.assign %[[RESULT]] to %[[RED_DECL]]#0 : i32, !fir.ref<i32>
+  ! CHECK-NEXT: omp.yield
+  !$omp do simd reduction(+:sum)
+    do index_ = 1, 10
+      sum = sum + 1
+    end do
+  !$omp end do simd
+end subroutine do_simd_reduction