[Mlir-commits] [mlir] b003fac - [flang][OpenMP] Add `OutlineableOpenMPOpInterface` to `omp.teams` (#131109)

[llvmlistbot at llvm.org]

Author: Kareem Ergawy
Date: 2025-03-13T16:03:19+01:00
New Revision: b003face11fadc526a6f816243441f486ffc958d

URL: https://github.com/llvm/llvm-project/commit/b003face11fadc526a6f816243441f486ffc958d
DIFF: https://github.com/llvm/llvm-project/commit/b003face11fadc526a6f816243441f486ffc958d.diff

LOG: [flang][OpenMP] Add `OutlineableOpenMPOpInterface` to `omp.teams` (#131109)

Given the following input:
```fortran
program rep_loopbind
  implicit none
  integer :: i
  real :: priv_val

  !$omp teams private(priv_val)
    !$omp distribute
    do i=1,1000
    end do
  !$omp end teams
end program
```
the `AllocaOpConversion` pattern in `FIRToLLVMLowering` would **move**
the private allocations that belong to the `teams` directive (i.e. the
allocations needed for the private copies of `priv_val` and the loop's
iteration variable) from the the `omp.teams` op to the outside scope.

This is not correct since these allocations should be eventually emitted
inside the outlined region for the `teams` directive. Without this fix,
these allocation would be emitted in the parent function (or the parent
scope whatever it is).

Added: 
    flang/test/Fir/omp-teams.fir

Modified: 
    flang/test/Transforms/stack-arrays-hlfir.f90
    mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

Removed: 
    


################################################################################
diff  --git a/flang/test/Fir/omp-teams.fir b/flang/test/Fir/omp-teams.fir
new file mode 100644
index 0000000000000..76599955fbac9
--- /dev/null
+++ b/flang/test/Fir/omp-teams.fir
@@ -0,0 +1,38 @@
+// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s
+
+// Tests that private allocations of an `omp.teams` op remain nested inside the
+// owning `omp.teams` op after FIR to LLVM conversion.
+
+func.func @foo() {
+  omp.teams {
+    %2 = fir.alloca f32 {bindc_name = "priv_val", pinned, uniq_name = "_QFEpriv_val"}
+    %3 = fir.alloca i32 {bindc_name = "i", pinned, uniq_name = "_QFEi"}
+    %c1_i32 = arith.constant 1 : i32
+    %c1000_i32 = arith.constant 1000 : i32
+    %c1_i32_0 = arith.constant 1 : i32
+    omp.distribute {
+      omp.loop_nest (%arg0) : i32 = (%c1_i32) to (%c1000_i32) inclusive step (%c1_i32_0) {
+        fir.store %arg0 to %3 : !fir.ref<i32>
+        omp.yield
+      }
+    }
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL:   llvm.func @foo() {
+// CHECK:           omp.teams {
+// CHECK:             %[[TEAMS_IV_ALLOC:.*]] = llvm.alloca %{{.*}} x i32 {bindc_name = "i", pinned}
+// CHECK:             %[[TEAMS_PRIV_ALLOC:.*]] = llvm.alloca %{{.*}} x f32 {bindc_name = "priv_val", pinned}
+// CHECK:             omp.distribute {
+// CHECK:               omp.loop_nest (%{{.*}}) : i32 = (%{{.*}}) to (%{{.*}}) inclusive step (%{{.*}}) {
+// CHECK:                 llvm.store %{{.*}}, %[[TEAMS_IV_ALLOC]] : i32, !llvm.ptr
+// CHECK:                 omp.yield
+// CHECK:               }
+// CHECK:             }
+// CHECK:             omp.terminator
+// CHECK:           }
+// CHECK:           llvm.return
+// CHECK:         }
+

diff  --git a/flang/test/Transforms/stack-arrays-hlfir.f90 b/flang/test/Transforms/stack-arrays-hlfir.f90
index 296edbaf35417..06749b7ca88af 100644
--- a/flang/test/Transforms/stack-arrays-hlfir.f90
+++ b/flang/test/Transforms/stack-arrays-hlfir.f90
@@ -71,8 +71,8 @@ end subroutine omp_target_wsloop
 ! CHECK:         omp.target {{.*}} {
 ! CHECK-NOT:       fir.allocmem
 ! CHECK-NOT:       fir.freemem
-! CHECK:           fir.alloca !fir.array<2xi64>
 ! CHECK:         omp.teams {
+! CHECK:           fir.alloca !fir.array<2xi64>
 ! CHECK:         omp.distribute {
 ! CHECK:         omp.loop_nest {{.*}} {
 ! CHECK-NOT:       fir.allocmem

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 5752446494280..8d657444a9b9e 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -236,7 +236,7 @@ def TerminatorOp : OpenMP_Op<"terminator", [Terminator, Pure]> {
 // 2.7 teams Construct
 //===----------------------------------------------------------------------===//
 def TeamsOp : OpenMP_Op<"teams", traits = [
-    AttrSizedOperandSegments, RecursiveMemoryEffects
+    AttrSizedOperandSegments, RecursiveMemoryEffects, OutlineableOpenMPOpInterface
   ], clauses = [
     OpenMP_AllocateClause, OpenMP_IfClause, OpenMP_NumTeamsClause,
     OpenMP_PrivateClause, OpenMP_ReductionClause, OpenMP_ThreadLimitClause