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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Kareem Ergawy (ergawy)

<details>
<summary>Changes</summary>

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).

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


3 Files Affected:

- (added) flang/test/Fir/omp-teams.fir (+38) 
- (modified) flang/test/Transforms/stack-arrays-hlfir.f90 (+1-1) 
- (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+1-1) 


``````````diff
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

``````````

</details>


https://github.com/llvm/llvm-project/pull/131109