[flang-commits] [flang] [flang] Add HLFIR-to-FIR pass pipeline extension points (PR #212194)

[Valentin Churavy via flang-commits]

================
@@ -0,0 +1,223 @@
+//===- PassPipelineTest.cpp -----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for the HLFIR extension points of the HLFIR-to-FIR pass pipeline.
+//
+// The callbacks run when the pipeline is built, so no IR is needed: building
+// the pipeline is enough to observe them.
+//
+//===----------------------------------------------------------------------===//
+
+#include "gtest/gtest.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
+#include "flang/Optimizer/Passes/Pipelines.h"
+#include "flang/Tools/CrossToolHelpers.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+#include <string>
+#include <vector>
+
+namespace {
+
+/// A no-op pass, identifiable by name in a textual pipeline, used to locate an
+/// extension point.
+struct MarkerPass : public mlir::PassWrapper<MarkerPass,
+                        mlir::OperationPass<mlir::ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MarkerPass)
+
+  llvm::StringRef getArgument() const override { return "ep-marker"; }
+  llvm::StringRef getDescription() const override {
+    return "No-op pass used to locate an extension point in a pipeline";
+  }
+  void runOnOperation() override {}
+};
+
+/// Render \p pm as a textual pipeline.
+std::string pipelineAsString(mlir::PassManager &pm) {
+  std::string out;
+  llvm::raw_string_ostream os(out);
+  pm.printAsTextualPipeline(os);
+  return out;
+}
+
+/// Build a HLFIR-to-FIR pipeline at \p level with a MarkerPass injected at each
+/// requested extension point, and return it as a textual pipeline.
+std::string pipelineWithMarkers(
+    llvm::OptimizationLevel level, bool early, bool last) {
+  mlir::MLIRContext context;
+  mlir::PassManager pm(&context, mlir::ModuleOp::getOperationName());
+  MLIRToLLVMPassPipelineConfig config(level);
+  if (early) {
+    config.registerHLFIROptEarlyEPCallbacks(
+        [](mlir::PassManager &nestedPm, llvm::OptimizationLevel) {
+          nestedPm.addPass(std::make_unique<MarkerPass>());
+        });
+  }
+  if (last) {
+    config.registerHLFIROptLastEPCallbacks(
+        [](mlir::PassManager &nestedPm, llvm::OptimizationLevel) {
+          nestedPm.addPass(std::make_unique<MarkerPass>());
+        });
+  }
+  fir::createHLFIRToFIRPassPipeline(pm, fir::EnableOpenMP::None, config);
+  return pipelineAsString(pm);
+}
+
+// Both callbacks are invoked, Early before Last.
+TEST(HLFIRExtensionPoint, CallbacksAreInvokedInOrder) {
+  mlir::MLIRContext context;
+  mlir::PassManager pm(&context, mlir::ModuleOp::getOperationName());
+  MLIRToLLVMPassPipelineConfig config(llvm::OptimizationLevel::O2);
+
+  std::vector<std::string> order;
+  size_t earlySizeAtCall = ~size_t{0}; // sentinel
+
+  config.registerHLFIROptEarlyEPCallbacks(
+      [&](mlir::PassManager &nestedPm, llvm::OptimizationLevel) {
+        order.push_back("early");
+        earlySizeAtCall = nestedPm.size();
+      });
+  config.registerHLFIROptLastEPCallbacks(
+      [&](mlir::PassManager &, llvm::OptimizationLevel) {
+        order.push_back("last");
+      });
+
+  fir::createHLFIRToFIRPassPipeline(pm, fir::EnableOpenMP::None, config);
+
+  ASSERT_EQ(order.size(), 2u);
+  EXPECT_EQ(order[0], "early");
+  EXPECT_EQ(order[1], "last");
+  // Early runs before anything has been added to the pipeline.
+  EXPECT_EQ(earlySizeAtCall, 0u);
+  EXPECT_GT(pm.size(), 0u);
+}
+
+// The callbacks fire at every level, including O0 where the simplification
+// passes are skipped.
+TEST(HLFIRExtensionPoint, CallbacksAreInvokedAtEveryOptLevel) {
+  for (llvm::OptimizationLevel level :
+      {llvm::OptimizationLevel::O0, llvm::OptimizationLevel::O1,
+          llvm::OptimizationLevel::O2, llvm::OptimizationLevel::O3}) {
+    mlir::MLIRContext context;
+    mlir::PassManager pm(&context, mlir::ModuleOp::getOperationName());
+    MLIRToLLVMPassPipelineConfig config(level);
+
+    int earlyCount = 0;
+    int lastCount = 0;
+    llvm::OptimizationLevel seenLevel = llvm::OptimizationLevel::O0;
+    config.registerHLFIROptEarlyEPCallbacks(
+        [&](mlir::PassManager &, llvm::OptimizationLevel cbLevel) {
+          ++earlyCount;
+          seenLevel = cbLevel;
+        });
+    config.registerHLFIROptLastEPCallbacks(
+        [&](mlir::PassManager &, llvm::OptimizationLevel) { ++lastCount; });
+
+    fir::createHLFIRToFIRPassPipeline(pm, fir::EnableOpenMP::None, config);
+
+    EXPECT_EQ(earlyCount, 1);
+    EXPECT_EQ(lastCount, 1);
+    // The callback is handed the level the pipeline was configured with.
+    EXPECT_EQ(seenLevel, level);
+  }
+}
+
+// Early must precede simplify-hlfir-intrinsics and Last must sit between it and
+// lower-hlfir-intrinsics. Asserting on the textual pipeline means a reordering
+// of createHLFIRToFIRPassPipeline breaks this test.
+TEST(HLFIRExtensionPoint, MarkersAreAtTheDocumentedPositions) {
+  std::string pipeline = pipelineWithMarkers(llvm::OptimizationLevel::O2,
+      /*early=*/true, /*last=*/true);
+
+  size_t earlyMarker = pipeline.find("ep-marker");
+  ASSERT_NE(earlyMarker, std::string::npos) << pipeline;
----------------
vchuravy wrote:

This style is used in a couple of places, one example is https://github.com/llvm/llvm-project/blob/d2fcb0c9e11f8d8bedc45db0aeba9e0a53141fe5/flang/unittests/Optimizer/InternalNamesTest.cpp#L44-L45

https://github.com/llvm/llvm-project/pull/212194