Writing Passes for the MQSS

Table of Contents

• Registering a new pass
• Writing a new pass
• Building your pass
• Using your new pass

The following sections describe how expand this collection of MLIR passes by defining custom passes.

Registering a new pass

Into the MQSS, the passes are classified as transforms, decompositions and code generation passes. To define a new pass you have to register it into the include/Passes/Transforms.hpp, include/Passes/Decompositions.hpp or include/Passes/CodeGen.hpp files, respectively.

Let say, you want to include a custom pass called CustomExamplePass as part of the decompositions collection. The method to create the pass must be registered in the file include/Passes/Decompositions.hpp as follows:

namespace mqss::opt {
  std::unique_ptr<mlir::Pass> createCustomExamplePass();
}

If the pass requires arguments, those have to be also declared into the signature of the method that creates the pass. For instance, we declare a CustomExampleArgumentPass that receives the argument int value.

namespace mqss::opt {
  std::unique_ptr<mlir::Pass> createCustomExampleArgumentPass(int value);
}

Notice that the pass has to be declared into namespace mqss::opt to be integrated as part of the MQSS.

Writing a new pass

This example serves as a very simple template for creating custom MLIR passes using QUAKE MLIR dialect and perform some general transformation. In this example, you can create a rewrite pattern that replaces Hadamard operations with S operations.

#include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
#include "cudaq/Support/Plugin.h"
#include "mlir/Rewrite/FrozenRewritePatternSet.h"
#include "mlir/Transforms/DialectConversion.h"
#include "Passes.hpp"
 
// The pass here is simple, replace Hadamard operations with S operations.
 
using namespace mlir;
 
namespace {
 
struct ReplaceH : public OpRewritePattern<quake::HOp> {
  using OpRewritePattern::OpRewritePattern;
  LogicalResult matchAndRewrite(quake::HOp hOp,
                                PatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<quake::SOp>(
        hOp, hOp.isAdj(), hOp.getParameters(), hOp.getControls(),
        hOp.getTargets());
    return success();
  }
};
 
class CustomExamplePassPlugin
    : public PassWrapper<CustomExamplePassPlugin, OperationPass<func::FuncOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CustomExamplePassPlugin)
 
  llvm::StringRef getArgument() const override { return "cudaq-custom-pass"; }
 
  void runOnOperation() override {
    auto circuit = getOperation();
    auto ctx = circuit.getContext();
 
    RewritePatternSet patterns(ctx);
    patterns.insert<ReplaceH>(ctx);
    ConversionTarget target(*ctx);
    target.addLegalDialect<quake::QuakeDialect>();
    target.addIllegalOp<quake::HOp>();
    if (failed(applyPartialConversion(circuit, target, std::move(patterns)))) {
      circuit.emitOpError("simple pass failed");
      signalPassFailure();
    }
  }
};
} // namespace
 
std::unique_ptr<Pass> mqss::opt::createCustomExamplePass(){
  return std::make_unique<CustomExamplePassPlugin>();
}

In order to be integrated into this project, the file CustomExamplePass.cpp must be in the lib/Passes/Decompositions, lib/Passes/Transforms or lib/Passes/CodeGen directory of this repository.

Building your pass

After including your pass in the project, you can build it as follows:

bash build.sh --mlir-dir "dir-to-mlir" --clang-dir "dir-to-clang"
              --llvm-dir "dir-to-llvm"

Using your new pass

Once your pass is integrated into this project. You can use it to transform any given MLIR/Quake module. Assuming that your MLIR module is into a string named quakeModule. First, you need to get the context and the module itself as follows:

auto [mlirModule, contextPtr] = extractMLIRContext(quakeModule);
mlir::MLIRContext &context = *contextPtr;

Apart of getting context and module, the function extractMLIRContext register the dialects to be used, in our case Quake dialect too. This tells to MLIR to recognize operations and functions belonging to Quake dialect.

Next, you have to declare a pass manager mlir::PassManager and load your custom pass mqss::opt::createCustomExamplePass as follows:

// creating pass manager
mlir::PassManager pm(&context);
// Adding custom pass
pm.addNestedPass<mlir::func::FuncOp>(mqss::opt::createCustomExamplePass());

To apply your custom pass on the MLIR module mlirModule, you have to run the pass manager as follows:

// running the pass
if(mlir::failed(pm.run(mlirModule)))
  std::runtime_error("The pass failed...");

If your pass is successfully applied, you can dump your module to visualize the effects of your pass in the module, as follows:

// Convert the module to a string
std::string moduleAsString;
llvm::raw_string_ostream stringStream(moduleAsString);
// Dump module to string
mlirModule->print(stringStream);
// Printing the transformed module
std::cout << "Module after Pass\n" << moduleAsString << std::endl;