[flang-commits] [flang] [Flang][OpenMP] Make boxed procedure pass aware of OpenMP private ops (PR #118261)

Kiran Chandramohan via flang-commits flang-commits at lists.llvm.org
Mon Dec 9 05:27:35 PST 2024


https://github.com/kiranchandramohan updated https://github.com/llvm/llvm-project/pull/118261

>From 2ff85cceeb0e09ef251b1d59d2ac323f322930f2 Mon Sep 17 00:00:00 2001
From: Kiran Chandramohan <kiran.chandramohan at arm.com>
Date: Mon, 2 Dec 2024 06:26:31 +0000
Subject: [PATCH 1/4] [Flang][OpenMP] Make boxed procedure pass aware of OpenMP
 private ops

Fixes #109727
---
 .../lib/Optimizer/CodeGen/BoxedProcedure.cpp  | 10 +++
 flang/test/Fir/boxproc-openmp.fir             | 87 +++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 flang/test/Fir/boxproc-openmp.fir

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index c536fd19fcc69a..850b5b4367a30d 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -16,6 +16,7 @@
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
 #include "flang/Optimizer/Support/FatalError.h"
 #include "flang/Optimizer/Support/InternalNames.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
@@ -374,6 +375,15 @@ class BoxedProcedurePass
               op->getResult(i.index()).setType(toTy);
             }
           rewriter.finalizeOpModification(op);
+        } else if (auto privateOp =
+                       mlir::dyn_cast<mlir::omp::PrivateClauseOp>(op)) {
+          mlir::Type dataTy = privateOp.getType();
+          if (typeConverter.needsConversion(dataTy)) {
+            auto toTy = typeConverter.convertType(dataTy);
+            rewriter.startOpModification(privateOp);
+            privateOp.setType(toTy);
+            rewriter.finalizeOpModification(privateOp);
+          }
         }
         // Ensure block arguments are updated if needed.
         if (opIsValid && op->getNumRegions() != 0) {
diff --git a/flang/test/Fir/boxproc-openmp.fir b/flang/test/Fir/boxproc-openmp.fir
new file mode 100644
index 00000000000000..8b714539b5e85a
--- /dev/null
+++ b/flang/test/Fir/boxproc-openmp.fir
@@ -0,0 +1,87 @@
+// RUN: fir-opt --boxed-procedure %s | FileCheck %s
+// Test the boxed procedure pass with OpenMP declaration operations.
+// Check minimally, only arguments, yields and the private types.
+
+// Test a private declaration with one region (alloc)
+//CHECK: omp.private {type = private}  @_QFsub1Et1_private_ref_rec__QFsub1Tt : !fir.ref<!fir.type<_QFsub1TtUnboxProc{p1:() -> ()}>> alloc {
+omp.private {type = private} @_QFsub1Et1_private_ref_rec__QFsub1Tt : !fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>> alloc {
+//CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.type<_QFsub1TtUnboxProc{p1:() -> ()}>>):
+^bb0(%arg0: !fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>):
+  %c1_i32 = arith.constant 1 : i32
+  %0 = fir.alloca !fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}> {bindc_name = "t1", pinned, uniq_name = "_QFsub1Et1"}
+  %1 = fir.declare %0 {uniq_name = "_QFsub1Et1"} : (!fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>) -> !fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>
+  %2 = fir.embox %1 : (!fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>) -> !fir.box<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>
+  %3 = fir.address_of(@_QQclXea6256ba131ddd9c2210e68030a0edd3) : !fir.ref<!fir.char<1,49>>
+  %4 = fir.convert %2 : (!fir.box<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>) -> !fir.box<none>
+  %5 = fir.convert %3 : (!fir.ref<!fir.char<1,49>>) -> !fir.ref<i8>
+  %6 = fir.call @_FortranAInitialize(%4, %5, %c1_i32) fastmath<contract> : (!fir.box<none>, !fir.ref<i8>, i32) -> none
+//CHECK: omp.yield(%{{.*}} : !fir.ref<!fir.type<_QFsub1TtUnboxProc{p1:() -> ()}>>)
+  omp.yield(%1 : !fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>)
+}
+func.func @_QPsub1() {
+  %0 = fir.alloca !fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}> {bindc_name = "t1", uniq_name = "_QFsub1Et1"}
+  %1 = fir.declare %0 {uniq_name = "_QFsub1Et1"} : (!fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>) -> !fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>
+//CHECK: omp.parallel private(@_QFsub1Et1_private_ref_rec__QFsub1Tt %{{.*}} -> %{{.*}} : !fir.ref<!fir.type<_QFsub1TtUnboxProc{p1:() -> ()}>>) {
+  omp.parallel private(@_QFsub1Et1_private_ref_rec__QFsub1Tt %1 -> %arg0 : !fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>) {
+    %2 = fir.declare %arg0 {uniq_name = "_QFsub1Et1"} : (!fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>) -> !fir.ref<!fir.type<_QFsub1Tt{p1:!fir.boxproc<() -> ()>}>>
+    omp.terminator
+  }
+  return
+}
+
+
+// Test a private declaration with all regions (alloc, copy, dealloc)
+//CHECK: omp.private {type = firstprivate} @_QFsub2Et1_firstprivate_ref_box_heap_rec__QFsub2Tt : 
+//CHECK-SAME: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>> alloc {
+omp.private {type = firstprivate} @_QFsub2Et1_firstprivate_ref_box_heap_rec__QFsub2Tt : !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>> alloc {
+//CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>):
+^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>):
+  %0 = fir.alloca !fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>> {bindc_name = "t1", pinned, uniq_name = "_QFsub2Et1"}
+  %1 = fir.declare %0 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub2Et1"} : (!fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>
+//CHECK: omp.yield(%{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>)
+  omp.yield(%1 : !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>)
+} copy {
+//CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>,
+//CHECK-SAME: %{{.*}}: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>):
+^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>, %arg1: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>):
+  %c5_i32 = arith.constant 5 : i32
+  %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>
+  %1 = fir.box_addr %0 : (!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>) -> !fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>
+  %2 = fir.embox %1 : (!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>) -> !fir.box<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>
+  %3 = fir.address_of(@_QQclXea) : !fir.ref<!fir.char<1,8>>
+  %4 = fir.convert %arg1 : (!fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>) -> !fir.ref<!fir.box<none>>
+  %5 = fir.convert %2 : (!fir.box<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>) -> !fir.box<none>
+  %6 = fir.convert %3 : (!fir.ref<!fir.char<1,8>>) -> !fir.ref<i8>
+  %7 = fir.call @_FortranAAssign(%4, %5, %6, %c5_i32) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.ref<i8>, i32) -> none
+//CHECK: omp.yield(%{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>)
+  omp.yield(%arg1 : !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>)
+} dealloc {
+//CHECK: ^bb0(%{{.*}}: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>):
+^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>):
+  %c5_i32 = arith.constant 5 : i32
+  %false = arith.constant false
+  %0 = fir.absent !fir.box<none>
+  %1 = fir.address_of(@_QQclXea) : !fir.ref<!fir.char<1,8>>
+  %2 = fir.convert %arg0 : (!fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>) -> !fir.ref<!fir.box<none>>
+  %3 = fir.convert %1 : (!fir.ref<!fir.char<1,8>>) -> !fir.ref<i8>
+  %4 = fir.call @_FortranAAllocatableDeallocate(%2, %false, %0, %3, %c5_i32) fastmath<contract> : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
+  omp.yield
+}
+func.func @_QPsub2() {
+  %0 = fir.alloca !fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>> {bindc_name = "t1", uniq_name = "_QFsub2Et1"}
+  %1 = fir.declare %0 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub2Et1"} : (!fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>
+//CHECK: omp.parallel private(@_QFsub2Et1_firstprivate_ref_box_heap_rec__QFsub2Tt %{{.*}} -> %{{.*}} :
+//CHECK-SAME: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>) {
+  omp.parallel private(@_QFsub2Et1_firstprivate_ref_box_heap_rec__QFsub2Tt %1 -> %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>) {
+    %2 = fir.declare %arg0 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub2Et1"} : (!fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2Tt{p1:!fir.boxproc<() -> ()>}>>>>
+    omp.terminator
+  }
+  return
+}
+func.func private @_FortranAInitialize(!fir.box<none>, !fir.ref<i8>, i32) -> none attributes {fir.runtime}
+fir.global linkonce @_QQclXea constant : !fir.char<1,8> {
+  %0 = fir.string_lit "pp.f90\00"(8) : !fir.char<1,8>
+  fir.has_value %0 : !fir.char<1,8>
+}
+func.func private @_FortranAAllocatableDeallocate(!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32 attributes {fir.runtime}
+func.func private @_FortranAAssign(!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.ref<i8>, i32) -> none attributes {fir.runtime}

>From c8a07b2b8eb44a3c43be11be48323ab35d94bb77 Mon Sep 17 00:00:00 2001
From: Kiran Chandramohan <kiran.chandramohan at arm.com>
Date: Wed, 4 Dec 2024 12:53:40 +0000
Subject: [PATCH 2/4] Make code generic by updating typeAttrs if needed

---
 .../lib/Optimizer/CodeGen/BoxedProcedure.cpp  | 21 +++++++++----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index 850b5b4367a30d..ade28fdb688e2f 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -221,7 +221,6 @@ class BoxedProcedurePass
       auto *context = &getContext();
       mlir::IRRewriter rewriter(context);
       BoxprocTypeRewriter typeConverter(mlir::UnknownLoc::get(context));
-      mlir::Dialect *firDialect = context->getLoadedDialect("fir");
       getModule().walk([&](mlir::Operation *op) {
         bool opIsValid = true;
         typeConverter.setLocation(op->getLoc());
@@ -367,23 +366,23 @@ class BoxedProcedurePass
                 index, toTy, index.getFieldId(), toOnTy, index.getTypeparams());
             opIsValid = false;
           }
-        } else if (op->getDialect() == firDialect) {
+        } else {
           rewriter.startOpModification(op);
+          // Convert the operands if needed
           for (auto i : llvm::enumerate(op->getResultTypes()))
             if (typeConverter.needsConversion(i.value())) {
               auto toTy = typeConverter.convertType(i.value());
               op->getResult(i.index()).setType(toTy);
             }
+
+          // Convert the type attributes if needed
+          for (const mlir::NamedAttribute &attr : op->getAttrDictionary())
+            if (auto tyAttr = llvm::dyn_cast<mlir::TypeAttr>(attr.getValue()))
+              if (typeConverter.needsConversion(tyAttr.getValue())) {
+                auto toTy = typeConverter.convertType(tyAttr.getValue());
+                op->setAttr(attr.getName(), mlir::TypeAttr::get(toTy));
+              }
           rewriter.finalizeOpModification(op);
-        } else if (auto privateOp =
-                       mlir::dyn_cast<mlir::omp::PrivateClauseOp>(op)) {
-          mlir::Type dataTy = privateOp.getType();
-          if (typeConverter.needsConversion(dataTy)) {
-            auto toTy = typeConverter.convertType(dataTy);
-            rewriter.startOpModification(privateOp);
-            privateOp.setType(toTy);
-            rewriter.finalizeOpModification(privateOp);
-          }
         }
         // Ensure block arguments are updated if needed.
         if (opIsValid && op->getNumRegions() != 0) {

>From 9f1bbde4e193a27cb7e9895e2f37f155173af39c Mon Sep 17 00:00:00 2001
From: Kiran Chandramohan <kiran.chandramohan at arm.com>
Date: Wed, 4 Dec 2024 17:47:21 +0000
Subject: [PATCH 3/4] Add conversion for TypeDescType

---
 flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp | 5 +++++
 flang/test/Fir/boxproc-2.fir                   | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index ade28fdb688e2f..b9dc49bf6ea69c 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -104,6 +104,8 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
       return needsConversion(unwrapRefType(ty));
     if (auto t = mlir::dyn_cast<SequenceType>(ty))
       return needsConversion(unwrapSequenceType(ty));
+    if (auto t = mlir::dyn_cast<TypeDescType>(ty))
+      return needsConversion(t.getOfTy());
     return false;
   }
 
@@ -168,6 +170,9 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
       rec.finalize(ps, cs);
       return rec;
     });
+    addConversion([&](TypeDescType ty) {
+      return TypeDescType::get(convertType(ty.getOfTy()));
+    });
     addArgumentMaterialization(materializeProcedure);
     addSourceMaterialization(materializeProcedure);
     addTargetMaterialization(materializeProcedure);
diff --git a/flang/test/Fir/boxproc-2.fir b/flang/test/Fir/boxproc-2.fir
index 84132f89afebb2..963dfa03663cb0 100644
--- a/flang/test/Fir/boxproc-2.fir
+++ b/flang/test/Fir/boxproc-2.fir
@@ -13,6 +13,9 @@ func.func private @test3(!fir.boxproc<() -> (!fir.type<a{x:i32, y:f64}>)>) -> no
 //CHECK-LABEL:  func.func private @test5(((i32) -> f32) -> ())
 func.func private @test5(!fir.boxproc<(!fir.boxproc<(i32) -> (f32)>) -> ()>)
 
+//CHECK-LABEL:  func.func private @test6(!fir.tdesc<!fir.type<_QMcartesian_2d_objectsTcartesian_2d_objectUnboxProc{base_pde_object:!fir.type<_QMbase_pde_objectsTbase_pde_objectUnboxProc{process_p:(!fir.class<!fir.type<_QMbase_pde_objectsTbase_pde_objectUnboxProc>>) -> !fir.class<!fir.heap<!fir.type<_QMbase_pde_objectsTbase_pde_objectUnboxProc>>>,source_p:(!fir.class<!fir.type<_QMbase_pde_objectsTbase_pde_objectUnboxProc>>, !fir.ref<f32>) -> !fir.class<!fir.heap<!fir.type<_QMbase_pde_objectsTbase_pde_objectUnboxProc>>>}>,c:!fir.box<!fir.heap<!fir.array<?x?xf32>>>,dx:f32,dy:f32}>>) -> !fir.ref<none>
+func.func private @test6(!fir.tdesc<!fir.type<_QMcartesian_2d_objectsTcartesian_2d_object{base_pde_object:!fir.type<_QMbase_pde_objectsTbase_pde_object{process_p:!fir.boxproc<(!fir.class<!fir.type<_QMbase_pde_objectsTbase_pde_object>>) -> !fir.class<!fir.heap<!fir.type<_QMbase_pde_objectsTbase_pde_object>>>>,source_p:!fir.boxproc<(!fir.class<!fir.type<_QMbase_pde_objectsTbase_pde_object>>, !fir.ref<f32>) -> !fir.class<!fir.heap<!fir.type<_QMbase_pde_objectsTbase_pde_object>>>>}>,c:!fir.box<!fir.heap<!fir.array<?x?xf32>>>,dx:f32,dy:f32}>>) -> !fir.ref<none>
+
 // CHECK-LABEL:   func.func @proc_pointer_component(
 // CHECK-SAME:                                      %[[VAL_0:.*]]: (!fir.ref<i32>) -> f32,
 // CHECK-SAME:                                      %[[VAL_1:.*]]: !fir.ref<i32>) {

>From cf200dbc023f274490d1db820b847933a042604a Mon Sep 17 00:00:00 2001
From: Kiran Chandramohan <kiran.chandramohan at arm.com>
Date: Mon, 9 Dec 2024 13:26:59 +0000
Subject: [PATCH 4/4] Remove OMP file inclusion in boxed procedure pass

---
 flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index b9dc49bf6ea69c..1bb91d252529f0 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -16,7 +16,6 @@
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
 #include "flang/Optimizer/Support/FatalError.h"
 #include "flang/Optimizer/Support/InternalNames.h"
-#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"



More information about the flang-commits mailing list