unittests/Pass/PassManagerTest.cpp - llvm-project/mlir - Git at Google

 //===- PassManagerTest.cpp - PassManager unit tests -----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//

 #include "mlir/Pass/PassManager.h"
 #include "mlir/Debug/BreakpointManagers/TagBreakpointManager.h"
 #include "mlir/Debug/ExecutionContext.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/Pass/Pass.h"
 #include "gtest/gtest.h"

 #include <memory>

 using namespace mlir;
 using namespace mlir::detail;

 namespace {
 /// Analysis that operates on any operation.
 struct GenericAnalysis {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GenericAnalysis)

   GenericAnalysis(Operation *op) : isFunc(isa<func::FuncOp>(op)) {}
   const bool isFunc;
 };

 /// Analysis that operates on a specific operation.
 struct OpSpecificAnalysis {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpSpecificAnalysis)

   OpSpecificAnalysis(func::FuncOp op) : isSecret(op.getName() == "secret") {}
   const bool isSecret;
 };

 /// Simple pass to annotate a func::FuncOp with the results of analysis.
 struct AnnotateFunctionPass
     : public PassWrapper<AnnotateFunctionPass, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AnnotateFunctionPass)

   void runOnOperation() override {
     func::FuncOp op = getOperation();
     Builder builder(op->getParentOfType<ModuleOp>());

     auto &ga = getAnalysis<GenericAnalysis>();
     auto &sa = getAnalysis<OpSpecificAnalysis>();

     op->setAttr("isFunc", builder.getBoolAttr(ga.isFunc));
     op->setAttr("isSecret", builder.getBoolAttr(sa.isSecret));
   }
 };

 TEST(PassManagerTest, OpSpecificAnalysis) {
   MLIRContext context;
   context.loadDialect<func::FuncDialect>();
   Builder builder(&context);

   // Create a module with 2 functions.
   OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
   for (StringRef name : {"secret", "not_secret"}) {
     auto func = func::FuncOp::create(
         builder.getUnknownLoc(), name,
         builder.getFunctionType(std::nullopt, std::nullopt));
     func.setPrivate();
     module->push_back(func);
   }

   // Instantiate and run our pass.
   auto pm = PassManager::on<ModuleOp>(&context);
   pm.addNestedPass<func::FuncOp>(std::make_unique<AnnotateFunctionPass>());
   LogicalResult result = pm.run(module.get());
   EXPECT_TRUE(succeeded(result));

   // Verify that each function got annotated with expected attributes.
   for (func::FuncOp func : module->getOps<func::FuncOp>()) {
     ASSERT_TRUE(isa<BoolAttr>(func->getDiscardableAttr("isFunc")));
     EXPECT_TRUE(cast<BoolAttr>(func->getDiscardableAttr("isFunc")).getValue());

     bool isSecret = func.getName() == "secret";
     ASSERT_TRUE(isa<BoolAttr>(func->getDiscardableAttr("isSecret")));
     EXPECT_EQ(cast<BoolAttr>(func->getDiscardableAttr("isSecret")).getValue(),
               isSecret);
   }
 }

 /// Simple pass to annotate a func::FuncOp with a single attribute `didProcess`.
 struct AddAttrFunctionPass
     : public PassWrapper<AddAttrFunctionPass, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddAttrFunctionPass)

   void runOnOperation() override {
     func::FuncOp op = getOperation();
     Builder builder(op->getParentOfType<ModuleOp>());
     if (op->hasAttr("didProcess"))
       op->setAttr("didProcessAgain", builder.getUnitAttr());

     // We always want to set this one.
     op->setAttr("didProcess", builder.getUnitAttr());
   }
 };

 /// Simple pass to annotate a func::FuncOp with a single attribute
 /// `didProcess2`.
 struct AddSecondAttrFunctionPass
     : public PassWrapper<AddSecondAttrFunctionPass,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddSecondAttrFunctionPass)

   void runOnOperation() override {
     func::FuncOp op = getOperation();
     Builder builder(op->getParentOfType<ModuleOp>());
     op->setAttr("didProcess2", builder.getUnitAttr());
   }
 };

 TEST(PassManagerTest, ExecutionAction) {
   MLIRContext context;
   context.loadDialect<func::FuncDialect>();
   Builder builder(&context);

   // Create a module with 2 functions.
   OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
   auto f =
       func::FuncOp::create(builder.getUnknownLoc(), "process_me_once",
                            builder.getFunctionType(std::nullopt, std::nullopt));
   f.setPrivate();
   module->push_back(f);

   // Instantiate our passes.
   auto pm = PassManager::on<ModuleOp>(&context);
   auto pass = std::make_unique<AddAttrFunctionPass>();
   auto *passPtr = pass.get();
   pm.addNestedPass<func::FuncOp>(std::move(pass));
   pm.addNestedPass<func::FuncOp>(std::make_unique<AddSecondAttrFunctionPass>());
   // Duplicate the first pass to ensure that we *only* run the *first* pass, not
   // all instances of this pass kind. Notice that this pass (and the test as a
   // whole) are built to ensure that we can run just a single pass out of a
   // pipeline that may contain duplicates.
   pm.addNestedPass<func::FuncOp>(std::make_unique<AddAttrFunctionPass>());

   // Use the action manager to only hit the first pass, not the second one.
   auto onBreakpoint = [&](const tracing::ActionActiveStack *backtrace)
       -> tracing::ExecutionContext::Control {
     // Not a PassExecutionAction, apply the action.
     auto *passExec = dyn_cast<PassExecutionAction>(&backtrace->getAction());
     if (!passExec)
       return tracing::ExecutionContext::Next;

     // If this isn't a function, apply the action.
     if (!isa<func::FuncOp>(passExec->getOp()))
       return tracing::ExecutionContext::Next;

     // Only apply the first function pass. Not all instances of the first pass,
     // only the first pass.
     if (passExec->getPass().getThreadingSiblingOrThis() == passPtr)
       return tracing::ExecutionContext::Next;

     // Do not apply any other passes in the pass manager.
     return tracing::ExecutionContext::Skip;
   };

   // Set up our breakpoint manager.
   tracing::TagBreakpointManager simpleManager;
   tracing::ExecutionContext executionCtx(onBreakpoint);
   executionCtx.addBreakpointManager(&simpleManager);
   simpleManager.addBreakpoint(PassExecutionAction::tag);

   // Register the execution context in the MLIRContext.
   context.registerActionHandler(executionCtx);

   // Run the pass manager, expecting our handler to be called.
   LogicalResult result = pm.run(module.get());
   EXPECT_TRUE(succeeded(result));

   // Verify that each function got annotated with `didProcess` and *not*
   // `didProcess2`.
   for (func::FuncOp func : module->getOps<func::FuncOp>()) {
     ASSERT_TRUE(func->getDiscardableAttr("didProcess"));
     ASSERT_FALSE(func->getDiscardableAttr("didProcess2"));
     ASSERT_FALSE(func->getDiscardableAttr("didProcessAgain"));
   }
 }

 namespace {
 struct InvalidPass : Pass {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InvalidPass)

   InvalidPass() : Pass(TypeID::get<InvalidPass>(), StringRef("invalid_op")) {}
   StringRef getName() const override { return "Invalid Pass"; }
   void runOnOperation() override {}
   bool canScheduleOn(RegisteredOperationName opName) const override {
     return true;
   }

   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clonePass() const override {
     return std::make_unique<InvalidPass>(
         *static_cast<const InvalidPass *>(this));
   }
 };
 } // namespace

 TEST(PassManagerTest, InvalidPass) {
   MLIRContext context;
   context.allowUnregisteredDialects();

   // Create a module
   OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));

   // Add a single "invalid_op" operation
   OpBuilder builder(&module->getBodyRegion());
   OperationState state(UnknownLoc::get(&context), "invalid_op");
   builder.insert(Operation::create(state));

   // Register a diagnostic handler to capture the diagnostic so that we can
   // check it later.
   std::unique_ptr<Diagnostic> diagnostic;
   context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
     diagnostic = std::make_unique<Diagnostic>(std::move(diag));
   });

   // Instantiate and run our pass.
   auto pm = PassManager::on<ModuleOp>(&context);
   pm.nest("invalid_op").addPass(std::make_unique<InvalidPass>());
   LogicalResult result = pm.run(module.get());
   EXPECT_TRUE(failed(result));
   ASSERT_TRUE(diagnostic.get() != nullptr);
   EXPECT_EQ(
       diagnostic->str(),
       "'invalid_op' op trying to schedule a pass on an unregistered operation");

   // Check that clearing the pass manager effectively removed the pass.
   pm.clear();
   result = pm.run(module.get());
   EXPECT_TRUE(succeeded(result));

   // Check that adding the pass at the top-level triggers a fatal error.
   ASSERT_DEATH(pm.addPass(std::make_unique<InvalidPass>()),
                "Can't add pass 'Invalid Pass' restricted to 'invalid_op' on a "
                "PassManager intended to run on 'builtin.module', did you "
                "intend to nest?");
 }

 /// Simple pass to annotate a func::FuncOp with the results of analysis.
 struct InitializeCheckingPass
     : public PassWrapper<InitializeCheckingPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InitializeCheckingPass)
   LogicalResult initialize(MLIRContext *ctx) final {
     initialized = true;
     return success();
   }
   bool initialized = false;

   void runOnOperation() override {
     if (!initialized) {
       getOperation()->emitError() << "Pass isn't initialized!";
       signalPassFailure();
     }
   }
 };

 TEST(PassManagerTest, PassInitialization) {
   MLIRContext context;
   context.allowUnregisteredDialects();

   // Create a module
   OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));

   // Instantiate and run our pass.
   auto pm = PassManager::on<ModuleOp>(&context);
   pm.addPass(std::make_unique<InitializeCheckingPass>());
   EXPECT_TRUE(succeeded(pm.run(module.get())));

   // Adding a second copy of the pass, we should also initialize it!
   pm.addPass(std::make_unique<InitializeCheckingPass>());
   EXPECT_TRUE(succeeded(pm.run(module.get())));
 }

 } // namespace
	//===- PassManagerTest.cpp - PassManager unit tests -----------------------===//
	//
	// 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
	//
	//===----------------------------------------------------------------------===//

	#include "mlir/Pass/PassManager.h"
	#include "mlir/Debug/BreakpointManagers/TagBreakpointManager.h"
	#include "mlir/Debug/ExecutionContext.h"
	#include "mlir/Dialect/Func/IR/FuncOps.h"
	#include "mlir/IR/Builders.h"
	#include "mlir/IR/BuiltinOps.h"
	#include "mlir/IR/Diagnostics.h"
	#include "mlir/Pass/Pass.h"
	#include "gtest/gtest.h"

	#include <memory>

	using namespace mlir;
	using namespace mlir::detail;

	namespace {
	/// Analysis that operates on any operation.
	struct GenericAnalysis {
	MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GenericAnalysis)

	GenericAnalysis(Operation *op) : isFunc(isa<func::FuncOp>(op)) {}
	const bool isFunc;
	};

	/// Analysis that operates on a specific operation.
	struct OpSpecificAnalysis {
	MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpSpecificAnalysis)

	OpSpecificAnalysis(func::FuncOp op) : isSecret(op.getName() == "secret") {}
	const bool isSecret;
	};

	/// Simple pass to annotate a func::FuncOp with the results of analysis.
	struct AnnotateFunctionPass
	: public PassWrapper<AnnotateFunctionPass, OperationPass<func::FuncOp>> {
	MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AnnotateFunctionPass)

	void runOnOperation() override {
	func::FuncOp op = getOperation();
	Builder builder(op->getParentOfType<ModuleOp>());

	auto &ga = getAnalysis<GenericAnalysis>();
	auto &sa = getAnalysis<OpSpecificAnalysis>();

	op->setAttr("isFunc", builder.getBoolAttr(ga.isFunc));
	op->setAttr("isSecret", builder.getBoolAttr(sa.isSecret));
	}
	};

	TEST(PassManagerTest, OpSpecificAnalysis) {
	MLIRContext context;
	context.loadDialect<func::FuncDialect>();
	Builder builder(&context);

	// Create a module with 2 functions.
	OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
	for (StringRef name : {"secret", "not_secret"}) {
	auto func = func::FuncOp::create(
	builder.getUnknownLoc(), name,
	builder.getFunctionType(std::nullopt, std::nullopt));
	func.setPrivate();
	module->push_back(func);
	}

	// Instantiate and run our pass.
	auto pm = PassManager::on<ModuleOp>(&context);
	pm.addNestedPass<func::FuncOp>(std::make_unique<AnnotateFunctionPass>());
	LogicalResult result = pm.run(module.get());
	EXPECT_TRUE(succeeded(result));

	// Verify that each function got annotated with expected attributes.
	for (func::FuncOp func : module->getOps<func::FuncOp>()) {
	ASSERT_TRUE(isa<BoolAttr>(func->getDiscardableAttr("isFunc")));
	EXPECT_TRUE(cast<BoolAttr>(func->getDiscardableAttr("isFunc")).getValue());

	bool isSecret = func.getName() == "secret";
	ASSERT_TRUE(isa<BoolAttr>(func->getDiscardableAttr("isSecret")));
	EXPECT_EQ(cast<BoolAttr>(func->getDiscardableAttr("isSecret")).getValue(),
	isSecret);
	}
	}

	/// Simple pass to annotate a func::FuncOp with a single attribute `didProcess`.
	struct AddAttrFunctionPass
	: public PassWrapper<AddAttrFunctionPass, OperationPass<func::FuncOp>> {
	MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddAttrFunctionPass)

	void runOnOperation() override {
	func::FuncOp op = getOperation();
	Builder builder(op->getParentOfType<ModuleOp>());
	if (op->hasAttr("didProcess"))
	op->setAttr("didProcessAgain", builder.getUnitAttr());

	// We always want to set this one.
	op->setAttr("didProcess", builder.getUnitAttr());
	}
	};

	/// Simple pass to annotate a func::FuncOp with a single attribute
	/// `didProcess2`.
	struct AddSecondAttrFunctionPass
	: public PassWrapper<AddSecondAttrFunctionPass,
	OperationPass<func::FuncOp>> {
	MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddSecondAttrFunctionPass)

	void runOnOperation() override {
	func::FuncOp op = getOperation();
	Builder builder(op->getParentOfType<ModuleOp>());
	op->setAttr("didProcess2", builder.getUnitAttr());
	}
	};

	TEST(PassManagerTest, ExecutionAction) {
	MLIRContext context;
	context.loadDialect<func::FuncDialect>();
	Builder builder(&context);

	// Create a module with 2 functions.
	OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
	auto f =
	func::FuncOp::create(builder.getUnknownLoc(), "process_me_once",
	builder.getFunctionType(std::nullopt, std::nullopt));
	f.setPrivate();
	module->push_back(f);

	// Instantiate our passes.
	auto pm = PassManager::on<ModuleOp>(&context);
	auto pass = std::make_unique<AddAttrFunctionPass>();
	auto *passPtr = pass.get();
	pm.addNestedPass<func::FuncOp>(std::move(pass));
	pm.addNestedPass<func::FuncOp>(std::make_unique<AddSecondAttrFunctionPass>());
	// Duplicate the first pass to ensure that we only run the first pass, not
	// all instances of this pass kind. Notice that this pass (and the test as a
	// whole) are built to ensure that we can run just a single pass out of a
	// pipeline that may contain duplicates.
	pm.addNestedPass<func::FuncOp>(std::make_unique<AddAttrFunctionPass>());

	// Use the action manager to only hit the first pass, not the second one.
	auto onBreakpoint = [&](const tracing::ActionActiveStack *backtrace)
	-> tracing::ExecutionContext::Control {
	// Not a PassExecutionAction, apply the action.
	auto *passExec = dyn_cast<PassExecutionAction>(&backtrace->getAction());
	if (!passExec)
	return tracing::ExecutionContext::Next;

	// If this isn't a function, apply the action.
	if (!isa<func::FuncOp>(passExec->getOp()))
	return tracing::ExecutionContext::Next;

	// Only apply the first function pass. Not all instances of the first pass,
	// only the first pass.
	if (passExec->getPass().getThreadingSiblingOrThis() == passPtr)
	return tracing::ExecutionContext::Next;

	// Do not apply any other passes in the pass manager.
	return tracing::ExecutionContext::Skip;
	};

	// Set up our breakpoint manager.
	tracing::TagBreakpointManager simpleManager;
	tracing::ExecutionContext executionCtx(onBreakpoint);
	executionCtx.addBreakpointManager(&simpleManager);
	simpleManager.addBreakpoint(PassExecutionAction::tag);

	// Register the execution context in the MLIRContext.
	context.registerActionHandler(executionCtx);

	// Run the pass manager, expecting our handler to be called.
	LogicalResult result = pm.run(module.get());
	EXPECT_TRUE(succeeded(result));

	// Verify that each function got annotated with `didProcess` and not
	// `didProcess2`.
	for (func::FuncOp func : module->getOps<func::FuncOp>()) {
	ASSERT_TRUE(func->getDiscardableAttr("didProcess"));
	ASSERT_FALSE(func->getDiscardableAttr("didProcess2"));
	ASSERT_FALSE(func->getDiscardableAttr("didProcessAgain"));
	}
	}

	namespace {
	struct InvalidPass : Pass {
	MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InvalidPass)

	InvalidPass() : Pass(TypeID::get<InvalidPass>(), StringRef("invalid_op")) {}
	StringRef getName() const override { return "Invalid Pass"; }
	void runOnOperation() override {}
	bool canScheduleOn(RegisteredOperationName opName) const override {
	return true;
	}

	/// A clone method to create a copy of this pass.
	std::unique_ptr<Pass> clonePass() const override {
	return std::make_unique<InvalidPass>(
	static_cast<const InvalidPass >(this));
	}
	};
	} // namespace

	TEST(PassManagerTest, InvalidPass) {
	MLIRContext context;
	context.allowUnregisteredDialects();

	// Create a module
	OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));

	// Add a single "invalid_op" operation
	OpBuilder builder(&module->getBodyRegion());
	OperationState state(UnknownLoc::get(&context), "invalid_op");
	builder.insert(Operation::create(state));

	// Register a diagnostic handler to capture the diagnostic so that we can
	// check it later.
	std::unique_ptr<Diagnostic> diagnostic;
	context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
	diagnostic = std::make_unique<Diagnostic>(std::move(diag));
	});

	// Instantiate and run our pass.
	auto pm = PassManager::on<ModuleOp>(&context);
	pm.nest("invalid_op").addPass(std::make_unique<InvalidPass>());
	LogicalResult result = pm.run(module.get());
	EXPECT_TRUE(failed(result));
	ASSERT_TRUE(diagnostic.get() != nullptr);
	EXPECT_EQ(
	diagnostic->str(),
	"'invalid_op' op trying to schedule a pass on an unregistered operation");

	// Check that clearing the pass manager effectively removed the pass.
	pm.clear();
	result = pm.run(module.get());
	EXPECT_TRUE(succeeded(result));

	// Check that adding the pass at the top-level triggers a fatal error.
	ASSERT_DEATH(pm.addPass(std::make_unique<InvalidPass>()),
	"Can't add pass 'Invalid Pass' restricted to 'invalid_op' on a "
	"PassManager intended to run on 'builtin.module', did you "
	"intend to nest?");
	}

	/// Simple pass to annotate a func::FuncOp with the results of analysis.
	struct InitializeCheckingPass
	: public PassWrapper<InitializeCheckingPass, OperationPass<ModuleOp>> {
	MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InitializeCheckingPass)
	LogicalResult initialize(MLIRContext *ctx) final {
	initialized = true;
	return success();
	}
	bool initialized = false;

	void runOnOperation() override {
	if (!initialized) {
	getOperation()->emitError() << "Pass isn't initialized!";
	signalPassFailure();
	}
	}
	};

	TEST(PassManagerTest, PassInitialization) {
	MLIRContext context;
	context.allowUnregisteredDialects();

	// Create a module
	OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));

	// Instantiate and run our pass.
	auto pm = PassManager::on<ModuleOp>(&context);
	pm.addPass(std::make_unique<InitializeCheckingPass>());
	EXPECT_TRUE(succeeded(pm.run(module.get())));

	// Adding a second copy of the pass, we should also initialize it!
	pm.addPass(std::make_unique<InitializeCheckingPass>());
	EXPECT_TRUE(succeeded(pm.run(module.get())));
	}

	} // namespace