[flang][OpenMP] Extend locality spec to OMP claues (init and dealloc regions)

[ergawy]

Extends support for locality specifier to OpenMP translation by adding supprot for transling localizers that have init and dealloc regions.

[llvmbot]

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

Author: Kareem Ergawy (ergawy)

Changes

Extends support for locality specifier to OpenMP translation by adding supprot for transling localizers that have init and dealloc regions.

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/OpenMP/DoConcurrentConversion.cpp (+25-4)
• (added) flang/test/Transforms/DoConcurrent/locality_specifiers_init_dealloc.mlir (+51)

diff --git a/flang/lib/Optimizer/OpenMP/DoConcurrentConversion.cpp b/flang/lib/Optimizer/OpenMP/DoConcurrentConversion.cpp
index 283c3052c166c..28f6c8bf02813 100644
--- a/flang/lib/Optimizer/OpenMP/DoConcurrentConversion.cpp
+++ b/flang/lib/Optimizer/OpenMP/DoConcurrentConversion.cpp
@@ -326,16 +326,37 @@ class DoConcurrentConversion
           TODO(localizer.getLoc(),
                "local_init conversion is not supported yet");
 
-        if (!localizer.getInitRegion().empty())
-          TODO(localizer.getLoc(),
-               "non-empty `init` regions are not supported yet");
-
         auto oldIP = rewriter.saveInsertionPoint();
         rewriter.setInsertionPointAfter(localizer);
         auto privatizer = rewriter.create<mlir::omp::PrivateClauseOp>(
             localizer.getLoc(), sym.getLeafReference().str() + ".omp",
             localizer.getTypeAttr().getValue(),
             mlir::omp::DataSharingClauseType::Private);
+
+        if (!localizer.getInitRegion().empty()) {
+          rewriter.cloneRegionBefore(localizer.getInitRegion(),
+                                     privatizer.getInitRegion(),
+                                     privatizer.getInitRegion().begin());
+          auto firYield = mlir::cast<fir::YieldOp>(
+              privatizer.getInitRegion().back().getTerminator());
+          rewriter.setInsertionPoint(firYield);
+          rewriter.create<mlir::omp::YieldOp>(firYield.getLoc(),
+                                              firYield.getOperands());
+          rewriter.eraseOp(firYield);
+        }
+
+        if (!localizer.getDeallocRegion().empty()) {
+          rewriter.cloneRegionBefore(localizer.getDeallocRegion(),
+                                     privatizer.getDeallocRegion(),
+                                     privatizer.getDeallocRegion().begin());
+          auto firYield = mlir::cast<fir::YieldOp>(
+              privatizer.getDeallocRegion().back().getTerminator());
+          rewriter.setInsertionPoint(firYield);
+          rewriter.create<mlir::omp::YieldOp>(firYield.getLoc(),
+                                              firYield.getOperands());
+          rewriter.eraseOp(firYield);
+        }
+
         rewriter.restoreInsertionPoint(oldIP);
 
         wsloopClauseOps.privateVars.push_back(op);
diff --git a/flang/test/Transforms/DoConcurrent/locality_specifiers_init_dealloc.mlir b/flang/test/Transforms/DoConcurrent/locality_specifiers_init_dealloc.mlir
new file mode 100644
index 0000000000000..a82d8d1715f56
--- /dev/null
+++ b/flang/test/Transforms/DoConcurrent/locality_specifiers_init_dealloc.mlir
@@ -0,0 +1,51 @@
+// Tests mapping `local` locality specifier to `private` clauses for non-empty
+// `init` and `dealloc` regions.
+
+// RUN: fir-opt --omp-do-concurrent-conversion="map-to=host" %s | FileCheck %s
+
+func.func @my_allocator() {
+  return
+}
+
+func.func @my_deallocator() {
+  return
+}
+
+fir.local {type = local} @_QFlocal_assocEaa_private_box_10xf32 : !fir.box<!fir.array<10xf32>> init {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<10xf32>>>, %arg1: !fir.ref<!fir.box<!fir.array<10xf32>>>):
+  fir.call @my_allocator() : () -> ()
+  fir.yield(%arg1 : !fir.ref<!fir.box<!fir.array<10xf32>>>)
+} dealloc {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<10xf32>>>):
+  fir.call @my_deallocator() : () -> ()
+  fir.yield
+}
+
+func.func @_QPlocal_assoc() {
+  %0 = fir.alloca !fir.box<!fir.array<10xf32>>
+  %c1 = arith.constant 1 : index
+
+  fir.do_concurrent {
+    %9 = fir.alloca i32 {bindc_name = "i"}
+    %10:2 = hlfir.declare %9 {uniq_name = "_QFlocal_assocEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+    fir.do_concurrent.loop (%arg0) = (%c1) to (%c1) step (%c1) local(@_QFlocal_assocEaa_private_box_10xf32 %0 -> %arg1 : !fir.ref<!fir.box<!fir.array<10xf32>>>) {
+      %11 = fir.convert %arg0 : (index) -> i32
+      fir.store %11 to %10#0 : !fir.ref<i32>
+    }
+  }
+
+  return
+}
+
+// CHECK:      omp.private {type = private} @[[PRIVATIZER:.*]] : !fir.box<!fir.array<10xf32>> init {
+// CHECK-NEXT: ^bb0(%{{.*}}: !{{.*}}, %{{.*}}: !{{.*}}):
+// CHECK-NEXT:   fir.call @my_allocator() : () -> ()
+// CHECK-NEXT:   omp.yield(%{{.*}})
+// CHECK-NEXT: } dealloc {
+// CHECK-NEXT: ^bb0(%{{.*}}: !{{.*}}):
+// CHECK-NEXT:   fir.call @my_deallocator() : () -> ()
+// CHECK-NEXT:   omp.yield
+// CHECK-NEXT: }
+
+// CHECK: %[[LOCAL_ALLOC:.*]] = fir.alloca !fir.box<!fir.array<10xf32>>
+// CHECK: omp.wsloop private(@[[PRIVATIZER]] %[[LOCAL_ALLOC]] -> %{{.*}} : !{{.*}})

[tblah]

Please could you update the test to use the block arguments so we can see they are mapped correctly