[Flang][OpenMP] Make boxed procedure pass aware of OpenMP private ops #118261

kiranchandramohan · 2024-12-02T06:34:15Z

graphite-app · 2024-12-02T06:34:22Z

Your org has enabled the Graphite merge queue for merging into main

Add the label “FP Bundles” to the PR and Graphite will automatically add it to the merge queue when it’s ready to merge.

You must have a Graphite account and log in to Graphite in order to use the merge queue. Sign up using this link.

llvmbot · 2024-12-02T06:34:48Z

@llvm/pr-subscribers-flang-codegen

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

Author: Kiran Chandramohan (kiranchandramohan)

Changes

Fixes #109727

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

2 Files Affected:

(modified) flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp (+10)
(added) flang/test/Fir/boxproc-openmp.fir (+86)

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..01582a13c0da4d
--- /dev/null
+++ b/flang/test/Fir/boxproc-openmp.fir
@@ -0,0 +1,86 @@
+// RUN: fir-opt --boxed-procedure %s | FileCheck %s
+// Test the boxed procedure pass with OpenMP declaration operations
+
+// 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(%arg0: !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(%1 : !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 %1 -> %arg0 : !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(%arg0: !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(%1 : !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(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.type<_QFsub2TtUnboxProc{p1:() -> ()}>>>>,
+//CHECK-SAME: %arg1: !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(%arg1 : !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(%arg0: !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 %1 -> %arg0 :
+//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}

Fixes llvm#109727

jeanPerier

Thanks Kiran. I wonder if the generic fix is not to remove the "op->getDialect() == firDialect" in the case just above. I do not recall the rational of limiting it to fir ops. Most likely this was an early compiler-perf optimization that was based on the wrong assumption that such FIR type would only be returned by FIR ops, but I see no harm removing that limitation (the needsConversion lookup cost should be pretty limited for most non FIR operation result types).

The right thing would likely to be to have a real dialect conversion pass, but these are expensive, so it may be better to keep the current approach for now (the next bag of issue I see would come from FIR type inside attributes of operations not handled here. There is some work to make dialect conversion cheaper happening in MLIR, maybe we can revisit at that point).

kiranchandramohan · 2024-12-02T12:34:28Z

Thanks @jeanPerier for the quick response.

I wonder if the generic fix is not to remove the "op->getDialect() == firDialect" in the case just above. I do not recall the rational of limiting it to fir ops.

This is a good idea. The fir dialect-specific check currently converts the results. For this particular case, see the next point.

(the next bag of issue I see would come from FIR type inside attributes of operations not handled here.

The type is stored as an attribute in the OpenMP PrivateClauseOp, which is why a separate else statement is required here.

llvm-project/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

Line 101 in 0611fdd

TypeAttrOf<AnyType>:$type,

Some kind of generalization might still be required for the other declaration/recipe-based operations.

jeanPerier

Thanks Kiran for the details, then the only generic solution I see with the current pass would be to also update the TypeAttr attribute in the code that updates the result type (like what is done for SymbolRefAttr in ExternalNameConversion).

If that does not work, you can assign an issue to me and I am OK with your patch.

kiranchandramohan · 2024-12-04T14:19:28Z

Thanks Kiran for the details, then the only generic solution I see with the current pass would be to also update the TypeAttr attribute in the code that updates the result type (like what is done for SymbolRefAttr in ExternalNameConversion).

If that does not work, you can assign an issue to me and I am OK with your patch.

Thanks @jeanPerier for the reference. Just an update here.

I added something like the following, which iterates through the attributes and updates if it is a typeAttr that needs conversion. This works fine for my usecase.

          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));
              }

But I see a failure in the gfortan test. I will have a detailed look later this week.

error: loc("/home/kiran/llvm-test-suite/Fortran/gfortran/regression/typebound_operator_9.f03":280:5): 'fir.type_desc' op wrapped type mismatched
error: loc("/home/kiran/llvm-test-suite/Fortran/gfortran/regression/typebound_operator_9.f03":315:5): 'fir.type_desc' op wrapped type mismatched
error: loc("/home/kiran/llvm-test-suite/Fortran/gfortran/regression/typebound_operator_9.f03":299:5): 'fir.type_desc' op wrapped type mismatched
error: loc("/home/kiran/llvm-test-suite/Fortran/gfortran/regression/typebound_operator_9.f03":228:5): 'fir.type_desc' op wrapped type mismatched
error: loc("/home/kiran/llvm-test-suite/Fortran/gfortran/regression/typebound_operator_9.f03":401:3): 'fir.type_desc' op wrapped type mismatched
            ```

jeanPerier · 2024-12-04T14:42:13Z

But I see a failure in the gfortan test. I will have a detailed look later this week.

Thanks for testing @kiranchandramohan, I think the issue is the BoxprocTypeRewriter type converter of this pass is not handling fir::TypeDescType, sorry for having you go through these issues. Do you want me to have a look?

Side note: the TypeDescOp design of duplicating the type (inside the result, and as attribute) is probably debatable/useless.

kiranchandramohan · 2024-12-04T15:21:33Z

But I see a failure in the gfortan test. I will have a detailed look later this week.

Thanks for testing @kiranchandramohan, I think the issue is the BoxprocTypeRewriter type converter of this pass is not handling fir::TypeDescType, sorry for having you go through these issues. Do you want me to have a look?

Side note: the TypeDescOp design of duplicating the type (inside the result, and as attribute) is probably debatable/useless.

Thanks for the hint. I will continue with this.

kiranchandramohan · 2024-12-05T12:42:36Z

Thanks for testing @kiranchandramohan, I think the issue is the BoxprocTypeRewriter type converter of this pass is not handling fir::TypeDescType, sorry for having you go through these issues. Do you want me to have a look?
Side note: the TypeDescOp design of duplicating the type (inside the result, and as attribute) is probably debatable/useless.

Thanks for the hint. I will continue with this.

Done

jeanPerier

The fix looks good (small nit on the include) and I worked fine for me. Thanks for making it generic @kiranchandramohan!

jeanPerier · 2024-12-09T13:22:01Z

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"


nit: not needed anymore

llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir flang:codegen labels Dec 2, 2024

[Flang][OpenMP] Make boxed procedure pass aware of OpenMP private ops

2ff85cc

Fixes llvm#109727

kiranchandramohan force-pushed the pp_openmp branch from ccca3af to 2ff85cc Compare December 2, 2024 06:45

kiranchandramohan requested review from ergawy, clementval, tblah and jeanPerier December 2, 2024 06:47

jeanPerier reviewed Dec 2, 2024

View reviewed changes

jeanPerier approved these changes Dec 2, 2024

View reviewed changes

kiranchandramohan added 2 commits December 5, 2024 10:54

Make code generic by updating typeAttrs if needed

c8a07b2

Add conversion for TypeDescType

9f1bbde

jeanPerier approved these changes Dec 9, 2024

View reviewed changes

Remove OMP file inclusion in boxed procedure pass

cf200db

kiranchandramohan merged commit 4e59721 into llvm:main Dec 9, 2024
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Flang][OpenMP] Make boxed procedure pass aware of OpenMP private ops #118261

[Flang][OpenMP] Make boxed procedure pass aware of OpenMP private ops #118261

Uh oh!

kiranchandramohan commented Dec 2, 2024

Uh oh!

graphite-app bot commented Dec 2, 2024

Uh oh!

llvmbot commented Dec 2, 2024 •

edited

Loading

Uh oh!

jeanPerier left a comment

Uh oh!

kiranchandramohan commented Dec 2, 2024 •

edited

Loading

Uh oh!

jeanPerier left a comment •

edited

Loading

Uh oh!

kiranchandramohan commented Dec 4, 2024

Uh oh!

jeanPerier commented Dec 4, 2024

Uh oh!

kiranchandramohan commented Dec 4, 2024

Uh oh!

kiranchandramohan commented Dec 5, 2024

Uh oh!

jeanPerier left a comment

Uh oh!

jeanPerier Dec 9, 2024

Uh oh!

Uh oh!

Uh oh!

[Flang][OpenMP] Make boxed procedure pass aware of OpenMP private ops #118261

[Flang][OpenMP] Make boxed procedure pass aware of OpenMP private ops #118261

Uh oh!

Conversation

kiranchandramohan commented Dec 2, 2024

Uh oh!

graphite-app bot commented Dec 2, 2024

Your org has enabled the Graphite merge queue for merging into main

Uh oh!

llvmbot commented Dec 2, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jeanPerier left a comment

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan commented Dec 2, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jeanPerier left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan commented Dec 4, 2024

Uh oh!

jeanPerier commented Dec 4, 2024

Uh oh!

kiranchandramohan commented Dec 4, 2024

Uh oh!

kiranchandramohan commented Dec 5, 2024

Uh oh!

jeanPerier left a comment

Choose a reason for hiding this comment

Uh oh!

jeanPerier Dec 9, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

llvmbot commented Dec 2, 2024 •

edited

Loading

kiranchandramohan commented Dec 2, 2024 •

edited

Loading

jeanPerier left a comment •

edited

Loading