[flang] de-duplicate AbstractResult pass

flang/include/flang/Optimizer/Dialect/FIROpsSupport.h

@@ -173,6 +173,13 @@ inline mlir::NamedAttribute getAdaptToByRefAttr(Builder &builder) {
           builder.getUnitAttr()};
 }
 
+/// Returns true if the operation name is for a container operation expected to
+/// contain (HL)FIR operations which need to be lowered by FIR passes. The
+/// simplest example of this is func.func.
+/// This operates on mlir::RegisteredOperationName so that it can be used to
+/// implement mlir::Pass::canScheduleOn.
+bool isa_toplevel(mlir::RegisteredOperationName opName);
+
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_DIALECT_FIROPSSUPPORT_H

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

@@ -31,8 +31,7 @@ namespace fir {
 // Passes defined in Passes.td
 //===----------------------------------------------------------------------===//
 
-#define GEN_PASS_DECL_ABSTRACTRESULTONFUNCOPT
-#define GEN_PASS_DECL_ABSTRACTRESULTONGLOBALOPT
+#define GEN_PASS_DECL_ABSTRACTRESULTOPT
 #define GEN_PASS_DECL_AFFINEDIALECTPROMOTION
 #define GEN_PASS_DECL_AFFINEDIALECTDEMOTION
 #define GEN_PASS_DECL_ANNOTATECONSTANTOPERANDS
@@ -50,8 +49,7 @@ namespace fir {
 #define GEN_PASS_DECL_OPENACCDATAOPERANDCONVERSION
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 
-std::unique_ptr<mlir::Pass> createAbstractResultOnFuncOptPass();
-std::unique_ptr<mlir::Pass> createAbstractResultOnGlobalOptPass();
+std::unique_ptr<mlir::Pass> createAbstractResultOptPass();
 std::unique_ptr<mlir::Pass> createAffineDemotionPass();
 std::unique_ptr<mlir::Pass>
 createArrayValueCopyPass(fir::ArrayValueCopyOptions options = {});

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

@@ -16,8 +16,8 @@
 
 include "mlir/Pass/PassBase.td"
 
-class AbstractResultOptBase<string optExt, string operation> 
-  : Pass<"abstract-result-on-" # optExt # "-opt", operation> {
+def AbstractResultOpt
+  : Pass<"abstract-result"> {
   let summary = "Convert fir.array, fir.box and fir.rec function result to "
                 "function argument";
   let description = [{
@@ -33,14 +33,7 @@ class AbstractResultOptBase<string optExt, string operation>
            "Pass fir.array<T> result as fir.box<fir.array<T>> argument instead"
            " of fir.ref<fir.array<T>>.">
   ];
-}
-
-def AbstractResultOnFuncOpt : AbstractResultOptBase<"func", "mlir::func::FuncOp"> {
-  let constructor = "::fir::createAbstractResultOnFuncOptPass()";
-}
-
-def AbstractResultOnGlobalOpt : AbstractResultOptBase<"global", "fir::GlobalOp"> {
-  let constructor = "::fir::createAbstractResultOnGlobalOptPass()";
+  let constructor = "::fir::createAbstractResultOptPass()";
 }
 
 def AffineDialectPromotion : Pass<"promote-to-affine", "::mlir::func::FuncOp"> {

flang/include/flang/Tools/CLOptions.inc

@@ -19,6 +19,7 @@
 #include "flang/Optimizer/Transforms/Passes.h"
 #include "llvm/Passes/OptimizationLevel.h"
 #include "llvm/Support/CommandLine.h"
+#include <type_traits>
 
 #define DisableOption(DOName, DOOption, DODescription) \
   static llvm::cl::opt<bool> disable##DOName("disable-" DOOption, \
@@ -86,6 +87,31 @@ DisableOption(BoxedProcedureRewrite, "boxed-procedure-rewrite",
 DisableOption(ExternalNameConversion, "external-name-interop",
     "convert names with external convention");
 
+// TODO: remove once these are used for non-codegen passes
+#if !defined(FLANG_EXCLUDE_CODEGEN)
+using PassConstructor = std::function<std::unique_ptr<mlir::Pass>()>;
+
+template <typename OP>
+void addNestedPassToOps(mlir::PassManager &pm, PassConstructor ctor) {
+  pm.addNestedPass<OP>(ctor());
+}
+
+template <typename OP, typename... OPS,
+    typename = std::enable_if_t<sizeof...(OPS) != 0>>
+void addNestedPassToOps(mlir::PassManager &pm, PassConstructor ctor) {
+  addNestedPassToOps<OP>(pm, ctor);
+  addNestedPassToOps<OPS...>(pm, ctor);
+}
+
+void addNestedPassToAllTopLevelOperations(
+    mlir::PassManager &pm, PassConstructor ctor) {
+  // TODO: add more operations that might need full lowering support
+  // any operations also need to be added to fir::isa_toplevel
+  addNestedPassToOps<mlir::func::FuncOp, mlir::omp::DeclareReductionOp,
+      fir::GlobalOp>(pm, ctor);
+}
+#endif
+
 /// Generic for adding a pass to the pass manager if it is not disabled.
 template <typename F>
 void addPassConditionally(
@@ -304,9 +330,7 @@ inline void createDebugPasses(
 inline void createDefaultFIRCodeGenPassPipeline(
     mlir::PassManager &pm, MLIRToLLVMPassPipelineConfig config) {
   fir::addBoxedProcedurePass(pm);
-  pm.addNestedPass<mlir::func::FuncOp>(
-      fir::createAbstractResultOnFuncOptPass());
-  pm.addNestedPass<fir::GlobalOp>(fir::createAbstractResultOnGlobalOptPass());
+  addNestedPassToAllTopLevelOperations(pm, fir::createAbstractResultOptPass);
   fir::addCodeGenRewritePass(pm);
   fir::addTargetRewritePass(pm);
   fir::addExternalNameConversionPass(pm, config.Underscoring);

flang/lib/Optimizer/Dialect/FIROps.cpp

@@ -3846,6 +3846,19 @@ std::optional<std::int64_t> fir::getIntIfConstant(mlir::Value value) {
   return {};
 }
 
+bool fir::isa_toplevel(mlir::RegisteredOperationName opName) {
+  const std::initializer_list<llvm::StringLiteral> topLevelOps{
+      fir::GlobalOp::getOperationName(),
+      mlir::func::FuncOp::getOperationName(),
+      mlir::omp::DeclareReductionOp::getOperationName(),
+  };
+
+  llvm::StringRef opStr = opName.getStringRef();
+  return llvm::any_of(topLevelOps, [&](const llvm::StringRef &topLevelOp) {
+    return opStr == topLevelOp;
+  });
+}
+
 mlir::Type fir::applyPathToType(mlir::Type eleTy, mlir::ValueRange path) {
   for (auto i = path.begin(), end = path.end(); eleTy && i < end;) {
     eleTy = llvm::TypeSwitch<mlir::Type, mlir::Type>(eleTy)

flang/lib/Optimizer/Transforms/AbstractResult.cpp

@@ -21,8 +21,7 @@
 #include "llvm/ADT/TypeSwitch.h"
 
 namespace fir {
-#define GEN_PASS_DEF_ABSTRACTRESULTONFUNCOPT
-#define GEN_PASS_DEF_ABSTRACTRESULTONGLOBALOPT
+#define GEN_PASS_DEF_ABSTRACTRESULTOPT
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 } // namespace fir
 
@@ -285,58 +284,8 @@ class AddrOfOpConversion : public mlir::OpRewritePattern<fir::AddrOfOp> {
   bool shouldBoxResult;
 };
 
-/// @brief Base CRTP class for AbstractResult pass family.
-/// Contains common logic for abstract result conversion in a reusable fashion.
-/// @tparam Pass target class that implements operation-specific logic.
-/// @tparam PassBase base class template for the pass generated by TableGen.
-/// The `Pass` class must define runOnSpecificOperation(OpTy, bool,
-/// mlir::RewritePatternSet&, mlir::ConversionTarget&) member function.
-/// This function should implement operation-specific functionality.
-template <typename Pass, template <typename> class PassBase>
-class AbstractResultOptTemplate : public PassBase<Pass> {
-public:
-  void runOnOperation() override {
-    auto *context = &this->getContext();
-    auto op = this->getOperation();
-
-    mlir::RewritePatternSet patterns(context);
-    mlir::ConversionTarget target = *context;
-    const bool shouldBoxResult = this->passResultAsBox.getValue();
-
-    auto &self = static_cast<Pass &>(*this);
-    self.runOnSpecificOperation(op, shouldBoxResult, patterns, target);
-
-    // Convert the calls and, if needed,  the ReturnOp in the function body.
-    target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect,
-                           mlir::func::FuncDialect>();
-    target.addIllegalOp<fir::SaveResultOp>();
-    target.addDynamicallyLegalOp<fir::CallOp>([](fir::CallOp call) {
-      return !hasAbstractResult(call.getFunctionType());
-    });
-    target.addDynamicallyLegalOp<fir::AddrOfOp>([](fir::AddrOfOp addrOf) {
-      if (auto funTy = addrOf.getType().dyn_cast<mlir::FunctionType>())
-        return !hasAbstractResult(funTy);
-      return true;
-    });
-    target.addDynamicallyLegalOp<fir::DispatchOp>([](fir::DispatchOp dispatch) {
-      return !hasAbstractResult(dispatch.getFunctionType());
-    });
-
-    patterns.insert<CallConversion<fir::CallOp>>(context, shouldBoxResult);
-    patterns.insert<CallConversion<fir::DispatchOp>>(context, shouldBoxResult);
-    patterns.insert<SaveResultOpConversion>(context);
-    patterns.insert<AddrOfOpConversion>(context, shouldBoxResult);
-    if (mlir::failed(
-            mlir::applyPartialConversion(op, target, std::move(patterns)))) {
-      mlir::emitError(op.getLoc(), "error in converting abstract results\n");
-      this->signalPassFailure();
-    }
-  }
-};
-
-class AbstractResultOnFuncOpt
-    : public AbstractResultOptTemplate<AbstractResultOnFuncOpt,
-                                       fir::impl::AbstractResultOnFuncOptBase> {
+class AbstractResultOpt
+    : public fir::impl::AbstractResultOptBase<AbstractResultOpt> {
 public:
   void runOnSpecificOperation(mlir::func::FuncOp func, bool shouldBoxResult,
                               mlir::RewritePatternSet &patterns,
@@ -386,40 +335,76 @@ class AbstractResultOnFuncOpt
       }
     }
   }
-};
 
-inline static bool containsFunctionTypeWithAbstractResult(mlir::Type type) {
-  return mlir::TypeSwitch<mlir::Type, bool>(type)
-      .Case([](fir::BoxProcType boxProc) {
-        return fir::hasAbstractResult(
-            boxProc.getEleTy().cast<mlir::FunctionType>());
-      })
-      .Case([](fir::PointerType pointer) {
-        return fir::hasAbstractResult(
-            pointer.getEleTy().cast<mlir::FunctionType>());
-      })
-      .Default([](auto &&) { return false; });
-}
+  inline static bool containsFunctionTypeWithAbstractResult(mlir::Type type) {
+    return mlir::TypeSwitch<mlir::Type, bool>(type)
+        .Case([](fir::BoxProcType boxProc) {
+          return fir::hasAbstractResult(
+              boxProc.getEleTy().cast<mlir::FunctionType>());
+        })
+        .Case([](fir::PointerType pointer) {
+          return fir::hasAbstractResult(
+              pointer.getEleTy().cast<mlir::FunctionType>());
+        })
+        .Default([](auto &&) { return false; });
+  }
 
-class AbstractResultOnGlobalOpt
-    : public AbstractResultOptTemplate<
-          AbstractResultOnGlobalOpt, fir::impl::AbstractResultOnGlobalOptBase> {
-public:
   void runOnSpecificOperation(fir::GlobalOp global, bool,
                               mlir::RewritePatternSet &,
                               mlir::ConversionTarget &) {
     if (containsFunctionTypeWithAbstractResult(global.getType())) {
       TODO(global->getLoc(), "support for procedure pointers");
     }
   }
+
+  virtual bool canScheduleOn(RegisteredOperationName opName) const override {
+    return fir::isa_toplevel(opName);
+  }
+
+  void runOnOperation() override {
+    auto *context = &this->getContext();
+    mlir::Operation *op = this->getOperation();
+
+    mlir::RewritePatternSet patterns(context);
+    mlir::ConversionTarget target = *context;
+    const bool shouldBoxResult = this->passResultAsBox.getValue();
+
+    mlir::TypeSwitch<mlir::Operation *, void>(op)
+        .Case<mlir::func::FuncOp, fir::GlobalOp>([&](auto op) {
+          runOnSpecificOperation(op, shouldBoxResult, patterns, target);
+        });
+
+    // Convert the calls and, if needed,  the ReturnOp in the function body.
+    target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect,
+                           mlir::func::FuncDialect>();
+    target.addIllegalOp<fir::SaveResultOp>();
+    target.addDynamicallyLegalOp<fir::CallOp>([](fir::CallOp call) {
+      return !hasAbstractResult(call.getFunctionType());
+    });
+    target.addDynamicallyLegalOp<fir::AddrOfOp>([](fir::AddrOfOp addrOf) {
+      if (auto funTy = addrOf.getType().dyn_cast<mlir::FunctionType>())
+        return !hasAbstractResult(funTy);
+      return true;
+    });
+    target.addDynamicallyLegalOp<fir::DispatchOp>([](fir::DispatchOp dispatch) {
+      return !hasAbstractResult(dispatch.getFunctionType());
+    });
+
+    patterns.insert<CallConversion<fir::CallOp>>(context, shouldBoxResult);
+    patterns.insert<CallConversion<fir::DispatchOp>>(context, shouldBoxResult);
+    patterns.insert<SaveResultOpConversion>(context);
+    patterns.insert<AddrOfOpConversion>(context, shouldBoxResult);
+    if (mlir::failed(
+            mlir::applyPartialConversion(op, target, std::move(patterns)))) {
+      mlir::emitError(op->getLoc(), "error in converting abstract results\n");
+      this->signalPassFailure();
+    }
+  }
 };
+
 } // end anonymous namespace
 } // namespace fir
 
-std::unique_ptr<mlir::Pass> fir::createAbstractResultOnFuncOptPass() {
-  return std::make_unique<AbstractResultOnFuncOpt>();
-}
-
-std::unique_ptr<mlir::Pass> fir::createAbstractResultOnGlobalOptPass() {
-  return std::make_unique<AbstractResultOnGlobalOpt>();
+std::unique_ptr<mlir::Pass> fir::createAbstractResultOptPass() {
+  return std::make_unique<AbstractResultOpt>();
 }

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

@@ -72,11 +72,13 @@
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 ! ALL-NEXT: BoxedProcedurePass
 
-! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func']
+! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction']
 ! ALL-NEXT:   'fir.global' Pipeline
-! ALL-NEXT:   AbstractResultOnGlobalOpt
+! ALL-NEXT:     AbstractResultOpt
 ! ALL-NEXT:   'func.func' Pipeline
-! ALL-NEXT:   AbstractResultOnFuncOpt
+! ALL-NEXT:     AbstractResultOpt
+! ALL-NEXT:   'omp.declare_reduction' Pipeline
+! ALL-NEXT:     AbstractResultOpt
 
 ! ALL-NEXT: CodeGenRewrite
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations eliminated

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

@@ -67,11 +67,13 @@
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 ! ALL-NEXT: BoxedProcedurePass
 
-! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func']
+! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction']
 ! ALL-NEXT:   'fir.global' Pipeline
-! ALL-NEXT:    AbstractResultOnGlobalOpt
+! ALL-NEXT:    AbstractResultOpt
 ! ALL-NEXT:  'func.func' Pipeline
-! ALL-NEXT:    AbstractResultOnFuncOpt
+! ALL-NEXT:    AbstractResultOpt
+! ALL-NEXT:  'omp.declare_reduction' Pipeline
+! ALL-NEXT:    AbstractResultOpt
 
 ! ALL-NEXT: CodeGenRewrite
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations eliminated

flang/test/Fir/abstract-result-2.fir

@@ -1,4 +1,4 @@
-// RUN: fir-opt %s --abstract-result-on-func-opt | FileCheck %s 
+// RUN: fir-opt %s --pass-pipeline="builtin.module(func.func(abstract-result))" | FileCheck %s
 
 // Check that the attributes are shifted along with their corresponding arguments
 

flang/test/Fir/abstract-results.fir

@@ -1,10 +1,10 @@
 // Test rewrite of functions that return fir.array<>, fir.type<>, fir.box<> to
 // functions that take an additional argument for the result.
 
-// RUN: fir-opt %s --abstract-result-on-func-opt | FileCheck %s --check-prefix=FUNC-REF
-// RUN: fir-opt %s --abstract-result-on-func-opt=abstract-result-as-box | FileCheck %s --check-prefix=FUNC-BOX
-// RUN: fir-opt %s --abstract-result-on-global-opt | FileCheck %s --check-prefix=GLOBAL-REF
-// RUN: fir-opt %s --abstract-result-on-global-opt=abstract-result-as-box | FileCheck %s --check-prefix=GLOBAL-BOX
+// RUN: fir-opt %s --pass-pipeline="builtin.module(func.func(abstract-result))" | FileCheck %s --check-prefix=FUNC-REF
+// RUN: fir-opt %s --pass-pipeline="builtin.module(func.func(abstract-result{abstract-result-as-box}))" | FileCheck %s --check-prefix=FUNC-BOX
+// RUN: fir-opt %s --pass-pipeline="builtin.module(fir.global(abstract-result))" | FileCheck %s --check-prefix=GLOBAL-REF
+// RUN: fir-opt %s --pass-pipeline="builtin.module(fir.global(abstract-result{abstract-result-as-box}))" | FileCheck %s --check-prefix=GLOBAL-BOX
 
 // ----------------------- Test declaration rewrite ----------------------------
 

flang/test/Fir/basic-program.fir

@@ -74,11 +74,13 @@ func.func @_QQmain() {
 // PASSES-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 // PASSES-NEXT: BoxedProcedurePass
 
-// PASSES-NEXT: Pipeline Collection : ['fir.global', 'func.func']
+// PASSES-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction']
 // PASSES-NEXT:   'fir.global' Pipeline
-// PASSES-NEXT:    AbstractResultOnGlobalOpt
+// PASSES-NEXT:    AbstractResultOpt
 // PASSES-NEXT:  'func.func' Pipeline
-// PASSES-NEXT:    AbstractResultOnFuncOpt
+// PASSES-NEXT:    AbstractResultOpt
+// PASSES-NEXT:  'omp.declare_reduction' Pipeline
+// PASSES-NEXT:    AbstractResultOpt
 
 // PASSES-NEXT: CodeGenRewrite
 // PASSES-NEXT:   (S) 0 num-dce'd - Number of operations eliminated

flang/test/Fir/non-trivial-procedure-binding-description.f90

@@ -1,5 +1,5 @@
 ! RUN: %flang_fc1 -emit-mlir %s -o - | FileCheck %s --check-prefix=BEFORE
-! RUN: %flang_fc1 -emit-mlir %s -o - | fir-opt --abstract-result-on-global-opt | FileCheck %s --check-prefix=AFTER
+! RUN: %flang_fc1 -emit-mlir %s -o - | fir-opt --pass-pipeline="builtin.module(fir.global(abstract-result))" | FileCheck %s --check-prefix=AFTER
 module a
   type f
   contains