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

[Kareem Ergawy]

https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/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).

>From 599da9d50b4102c7756aee674b9dafab3fc0031a Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 13 Mar 2025 04:06:22 -0500
Subject: [PATCH] [flang][OpenMP] Add `OutlineableOpenMPOpInterface` to
 `omp.teams`

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).
---
 flang/test/Fir/omp-teams.fir                  | 38 +++++++++++++++++++
 flang/test/Transforms/stack-arrays-hlfir.f90  |  2 +-
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td |  2 +-
 3 files changed, 40 insertions(+), 2 deletions(-)
 create mode 100644 flang/test/Fir/omp-teams.fir

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 401c4c11d8986..f0e197a87e1c5 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