Create your own CUDA-Q Compiler Pass

The CUDA-Q IR can be transformed, analyzed, or optimized using standard MLIR patterns and tools. CUDA-Q provides a registration mechanism for the cudaq-opt tool that allows one to create, load, and use custom MLIR passes on Quake code.

CUDA-Q MLIR Passes can only be created within an existing CUDA-Q development environment. Therefore, you must clone the repository and add your Pass code as part of the existing CUDA-Q CMake system.

As an example, clone the repository and add the following directory structure under lib, lib/Plugins/MyCustomPlugin/. Within this directory create a CMakeLists.txt file and a MyCustomPlugin.cpp file. In the CMake file, add the following:

add_llvm_pass_plugin(MyCustomPlugin MyCustomPlugin.cpp)

Creating a CUDA-Q IR pass starts with the implementation of an mlir::OperationPass. A full discussion of the MLIR Pass infrastructure is beyond the scope of this document; please see MLIR Passes. To create such a pass, start with the following template in the MyCustomPlugin.cpp file:

#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"

// Here is an example MLIR Pass that one can write externally and
// use via the cudaq-opt tool, with the --load-cudaq-plugin flag.
// 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 CustomPassPlugin
    : public PassWrapper<CustomPassPlugin, OperationPass<func::FuncOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CustomPassPlugin)

  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

CUDAQ_REGISTER_MLIR_PASS(CustomPassPlugin)

This example serves as a very simple template for creating custom MLIR Passes that analyze the CUDA-Q Quake representation and perform some general transformation. In this example, we create a rewrite pattern that replaces Hadamard operations with S operations.

Ensure that add_subdirectory(Plugins) is in the lib/CMakeLists.txt file, and also that there is a lib/Plugins/CMakeLists.txt file that adds your plugin directory with add_subdirectory.

Then build CUDA-Q and you will have a MyCustomPlugin.so plugin library in the install. You can load the plugin and use it with cudaq-opt as follows:

cudaq-opt --load-cudaq-plugin MyCustomPlugin.so file.qke -cudaq-custom-pass