[flang] Adapt PolymorphicOpConversion to run on all top level ops (#90597)

We might use polymorphic ops in top-level operations other than
functions some time in the future. We need to ensure that these
operations can be lowered.

See RFC:

https://discourse.llvm.org/t/rfc-add-an-interface-for-top-level-container-operations

Some of the changes are from moving declaration and definition of the
constructor function into tablegen (as requested in code review when
altering another pass).

Author: tblah

flang/include/flang/Optimizer/Transforms/Passes.h

@@ -67,7 +67,6 @@ std::unique_ptr<mlir::Pass> createAnnotateConstantOperandsPass();
 std::unique_ptr<mlir::Pass> createAlgebraicSimplificationPass();
 std::unique_ptr<mlir::Pass>
 createAlgebraicSimplificationPass(const mlir::GreedyRewriteConfig &config);
-std::unique_ptr<mlir::Pass> createPolymorphicOpConversionPass();
 
 std::unique_ptr<mlir::Pass> createOMPDescriptorMapInfoGenPass();
 std::unique_ptr<mlir::Pass> createOMPFunctionFilteringPass();

flang/include/flang/Optimizer/Transforms/Passes.td

@@ -298,14 +298,13 @@ def AlgebraicSimplification : Pass<"flang-algebraic-simplification"> {
   let constructor = "::fir::createAlgebraicSimplificationPass()";
 }
 
-def PolymorphicOpConversion : Pass<"fir-polymorphic-op", "::mlir::func::FuncOp"> {
+def PolymorphicOpConversion : Pass<"fir-polymorphic-op"> {
   let summary =
     "Simplify operations on polymorphic types";
   let description = [{
     This pass breaks up the lowering of operations on polymorphic types by 
     introducing an intermediate FIR level that simplifies code geneation. 
   }];
-  let constructor = "::fir::createPolymorphicOpConversionPass()";
   let dependentDialects = [
     "fir::FIROpsDialect", "mlir::func::FuncDialect"
   ];

flang/include/flang/Tools/CLOptions.inc

@@ -271,7 +271,7 @@ inline void createDefaultFIROptimizerPassPipeline(
   pm.addPass(mlir::createCSEPass());
 
   // Polymorphic types
-  pm.addPass(fir::createPolymorphicOpConversionPass());
+  addNestedPassToAllTopLevelOperations(pm, fir::createPolymorphicOpConversion);
 
   if (pc.AliasAnalysis && !disableFirAliasTags && !useOldAliasTags)
     pm.addPass(fir::createAliasTagsPass());

flang/lib/Optimizer/Transforms/PolymorphicOpConversion.cpp

@@ -229,7 +229,9 @@ class PolymorphicOpConversion
 
   void runOnOperation() override {
     auto *context = &getContext();
-    auto mod = getOperation()->getParentOfType<ModuleOp>();
+    auto mod = mlir::dyn_cast_or_null<mlir::ModuleOp>(getOperation());
+    if (!mod)
+      mod = getOperation()->getParentOfType<ModuleOp>();
     mlir::RewritePatternSet patterns(context);
 
     BindingTables bindingTables;
@@ -471,7 +473,3 @@ SelectTypeConv::collectAncestors(fir::TypeInfoOp dt, mlir::ModuleOp mod) const {
   }
   return ancestors;
 }
-
-std::unique_ptr<mlir::Pass> fir::createPolymorphicOpConversionPass() {
-  return std::make_unique<PolymorphicOpConversion>();
-}

flang/test/Driver/bbc-mlir-pass-pipeline.f90

@@ -47,13 +47,16 @@
 
 ! CHECK-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! CHECK-NEXT: 'fir.global' Pipeline
+! CHECK-NEXT:   PolymorphicOpConversion
 ! CHECK-NEXT:   CFGConversion
 ! CHECK-NEXT: 'func.func' Pipeline
 ! CHECK-NEXT:   PolymorphicOpConversion
 ! CHECK-NEXT:   CFGConversion
 ! CHECK-NEXT: 'omp.declare_reduction' Pipeline
+! CHECK-NEXT:   PolymorphicOpConversion
 ! CHECK-NEXT:   CFGConversion
 ! CHECK-NEXT: 'omp.private' Pipeline
+! CHECK-NEXT:   PolymorphicOpConversion
 ! CHECK-NEXT:   CFGConversion
 
 ! CHECK-NEXT: SCFToControlFlow

flang/test/Driver/mlir-debug-pass-pipeline.f90

@@ -67,13 +67,16 @@
 
 ! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! ALL-NEXT:   'fir.global' Pipeline
+! ALL-NEXT:     PolymorphicOpConversion
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT:   'func.func' Pipeline
 ! ALL-NEXT:     PolymorphicOpConversion
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT:   'omp.declare_reduction' Pipeline
+! ALL-NEXT:     PolymorphicOpConversion
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT:   'omp.private' Pipeline
+! ALL-NEXT:     PolymorphicOpConversion
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT: SCFToControlFlow
 ! ALL-NEXT: Canonicalizer

flang/test/Driver/mlir-pass-pipeline.f90

@@ -56,18 +56,28 @@
 ! ALL-NEXT:   (S) 0 num-cse'd - Number of operations CSE'd
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 
+! O2-NEXT:  Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
+! O2-NEXT:    'fir.global' Pipeline
+! O2-NEXT:      PolymorphicOpConversion
 ! O2-NEXT:    'func.func' Pipeline
 ! O2-NEXT:      PolymorphicOpConversion
+! O2-NEXT:    'omp.declare_reduction' Pipeline
+! O2-NEXT:      PolymorphicOpConversion
+! O2-NEXT:    'omp.private' Pipeline
+! O2-NEXT:      PolymorphicOpConversion
 ! O2-NEXT:  AddAliasTags
 ! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! ALL-NEXT:    'fir.global' Pipeline
+! NOTO2-NEXT:      PolymorphicOpConversion
 ! ALL-NEXT:      CFGConversion
 ! ALL-NEXT:    'func.func' Pipeline
 ! NOTO2-NEXT:      PolymorphicOpConversion
 ! ALL-NEXT:      CFGConversion
 ! ALL-NEXT:   'omp.declare_reduction' Pipeline
+! NOTO2-NEXT:      PolymorphicOpConversion
 ! ALL-NEXT:      CFGConversion
 ! ALL-NEXT:   'omp.private' Pipeline
+! NOTO2-NEXT:    PolymorphicOpConversion
 ! ALL-NEXT:      CFGConversion
 
 ! ALL-NEXT: SCFToControlFlow

flang/test/Fir/basic-program.fir

@@ -62,8 +62,15 @@ func.func @_QQmain() {
 // PASSES-NEXT:   (S) 0 num-cse'd - Number of operations CSE'd
 // PASSES-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 
+// PASSES-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
+// PASSES-NEXT: 'fir.global' Pipeline
+// PASSES-NEXT:   PolymorphicOpConversion
 // PASSES-NEXT: 'func.func' Pipeline
 // PASSES-NEXT:   PolymorphicOpConversion
+// PASSES-NEXT: 'omp.declare_reduction' Pipeline
+// PASSES-NEXT:   PolymorphicOpConversion
+// PASSES-NEXT: 'omp.private' Pipeline
+// PASSES-NEXT:   PolymorphicOpConversion
 
 // PASSES-NEXT: AddAliasTags