Catalyst

Development

Examples

API

  1. Docs
  2. MLIR Plugins

MLIR Plugins

This page describes how to develop an MLIR plugin that can work with Catalyst. An MLIR plugin is a shared object, defined and built outside of a project’s primary source tree, that implements custom dialects or compilation passes compatible with the MLIR framework. Because Catalyst is built natively on top of MLIR, it seamlessly supports these external plugins. This extensibility allows developers to integrate custom quantum compilation passes and dialects directly into the Catalyst compilation pipeline.

Building the Standalone Plugin

Catalyst provides a pre-configured Makefile rule to build the standalone-plugin provided in the upstream MLIR source repository. To compile the plugin, execute the following command:

make plugin

Upon successful compilation, the shared object file StandalonePlugin.so will be located in the catalyst/mlir/standalone/build/lib/ directory. The StandalonePlugin.so file is a minimal reference plugin that introduces a custom dialect (called Standalone) and a single transformation pass that renames symbols from bar to foo. It is intended as a structural template for building external MLIR plugins, rather than a comprehensive demonstration of advanced MLIR capabilities.

You can use the StandalonePlugin.so plugin to transform a quantum program

  • with either quantum-opt or catalyst, or

  • by loading it from a Python program.

If you are interested in using it from the command-line interface, you can add the following flags to load the Standalone plugin:

  • --load-pass-plugin=/path/to/StandalonePlugin.so

  • --load-dialect-plugin=/path/to/StandalonePlugin.so

Notice that you will still be able to use the rest of the flags and arguments accepted by the command-line interface. For example, passing --help as in

quantum-opt --load-pass-plugin=/path/to/StandalonePlugin.so --help

outputs the documentation available for the Standalone pass:

--standalone-switch-bar-foo    - Switches the name of a FuncOp named `bar` to `foo` and folds.

Taking into account the description of the pass standalone-switch-bar-foo, let’s write the most minimal program that would be transformed by this transformation:

module @module {
  func.func private @bar() -> (tensor<i64>) {
    %c = stablehlo.constant dense<0> : tensor<i64>
    return %c : tensor<i64>
  }
}

You can schedule this pass as any other pass:

quantum-opt \
    --load-pass-plugin=/path/to/StandalonePlugin.so \
    --pass-pipeline='builtin.module(standalone-switch-bar-foo)' \
    example.mlir

Note

The current implementation of MLIR plugins only supports using the --pass-pipeline option for specifying passes.

And you have your transformed program:

module @module {
  func.func private @foo() -> tensor<i64> {
    %c = stablehlo.constant dense<0> : tensor<i64>
    return %c : tensor<i64>
  }
}

Notice that the name of the function bar has been changed to foo.

Pass Plugins vs Dialect Plugins

You may now be asking, “how come we used the option --load-pass-plugin but we didn’t use the option --load-dialect-plugin?” The --load-pass-plugin option is used to load passes, while the --load-dialect-plugin is used to load dialects. As mentioned earlier, the StandalonePlugin.so file also contains a dialect. It is a simple dialect intended only for testing purposes, and it only contains a single operation, standalone.foo. (Please do not confuse this operation with symbols named foo).

We can write a program that contains operations in the standalone dialect:

module @module {
  func.func private @bar() -> (i32) {
    %0 = arith.constant 0 : i32
    %1 = standalone.foo %0 : i32
    return %1 : i32
  }
}

But if we try to run it, using the same command as shown earlier

quantum-opt \
    --load-pass-plugin=/path/to/StandalonePlugin.so \
    --pass-pipeline='builtin.module(standalone-switch-bar-foo)' \
    example.mlir

the compilation will fail with the following message:

example.mlir:4:10: error: Dialect `standalone' not found for custom op 'standalone.foo'
%1 = standalone.foo %0 : i32
     ^
a.mlir:4:10: note: Registered dialects: acc, affine, amdgpu, amx, arith, arm_neon, arm_sme, arm_sve, async, bufferization, builtin, catalyst, cf, chlo, complex, dlti, emitc, func, gpu, gradient, index, irdl, linalg, llvm, math, memref, mesh, mhlo, mitigation, ml_program, mpi, nvgpu, nvvm, omp, pdl, pdl_interp, polynomial, quant, quantum, rocdl, scf, shape, sparse_tensor, spirv, stablehlo, tensor, test, tosa, transform, ub, vector, vhlo, x86vector, xegpu ; for more info on dialect registration see https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management

In order to parse operations from this dialect, we need to load the dialect from the plugin shared object file:

quantum-opt \
    --load-pass-plugin=/path/to/StandalonePlugin.so \
    --load-dialect-plugin=/path/to/StandalonePlugin.so \
    --pass-pipeline='builtin.module(standalone-switch-bar-foo)' \
    example.mlir

Now, you can parse the program without the error and run the standalone-switch-bar-foo pass:

module @module {
  func.func private @foo() -> i32 {
    %c0_i32 = arith.constant 0 : i32
    %0 = standalone.foo %c0_i32 : i32
    return %0 : i32
  }
}

Creating your own Pass Plugin

The Catalyst repository includes the LLVM project as a Git submodule, which contains the upstream standalone plugin example. Executing make standalone-plugin copies this example directory and applies the necessary patches to ensure compatibility with the Catalyst environment.

Because the standalone plugin is a minimal reference implementation, developing a production-ready plugin capable of manipulating quantum programs requires modifications to its build system. The recommended workflow for adapting the build scripts is detailed below:

1. Add the standalone plugin directory as a subdirectory of Catalyst

diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
index c0b8dfd6c..1b5c2e528 100644
--- a/mlir/CMakeLists.txt
+++ b/mlir/CMakeLists.txt
@@ -73,6 +73,7 @@ add_subdirectory(include)
add_subdirectory(lib)
 add_subdirectory(tools)
 add_subdirectory(test)
+add_subdirectory(standalone)

 if(QUANTUM_ENABLE_BINDINGS_PYTHON)
   message(STATUS "Enabling Python API")

You will also need to make the following change:

diff --git a/mlir/standalone/CMakeLists.txt b/mlir/standalone/CMakeLists.txt
index e999ae34d..fd6ee8f10 100644
--- a/mlir/standalone/CMakeLists.txt
+++ b/mlir/standalone/CMakeLists.txt
@@ -1,6 +1,3 @@
-cmake_minimum_required(VERSION 3.20.0)
-project(standalone-dialect LANGUAGES CXX C)
-
 set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)

 set(CMAKE_CXX_STANDARD 17 CACHE STRING "C++ standard to conform to")

diff --git a/mlir/standalone/CMakeLists.txt b/mlir/standalone/CMakeLists.txt
index 280cd80e1..fd6ee8f10 100644
--- a/mlir/standalone/CMakeLists.txt
+++ b/mlir/standalone/CMakeLists.txt
@@ -32,8 +32,8 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   mlir_configure_python_dev_packages()
 endif()

-set(STANDALONE_SOURCE_DIR ${PROJECT_SOURCE_DIR})
-set(STANDALONE_BINARY_DIR ${PROJECT_BINARY_DIR})
+set(STANDALONE_SOURCE_DIR ${PROJECT_SOURCE_DIR}/standalone)
+set(STANDALONE_BINARY_DIR ${PROJECT_BINARY_DIR}/standalone)
 include_directories(${LLVM_INCLUDE_DIRS})
 include_directories(${MLIR_INCLUDE_DIRS})
 include_directories(${STANDALONE_SOURCE_DIR}/include)

With these changes, you should now be able to use make all and build the standalone plugin. Please note that the location of the StandalonePlugin.so shared object has changed. It will now be stored in the mlir/build/lib/ directory.

2. Include the header files in the standalone plugin pass

diff --git a/mlir/standalone/lib/Standalone/StandalonePasses.cpp b/mlir/standalone/lib/Standalone/StandalonePasses.cpp
index a23d0420f..83e2ce255 100644
--- a/mlir/standalone/lib/Standalone/StandalonePasses.cpp
+++ b/mlir/standalone/lib/Standalone/StandalonePasses.cpp
@@ -12,6 +12,7 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"

 #include "Standalone/StandalonePasses.h"
+#include "Quantum/IR/QuantumOps.h"

 namespace mlir::standalone {
 #define GEN_PASS_DEF_STANDALONESWITCHBARFOO

You can execute make all and see the compilation succeed. Please note that Catalyst has a number of custom dialects, such as the Quantum, Catalyst and Gradient dialects. Depending on which dialect you are interested in, you can include the definition of the operations in a similar way by including the appropriate dialect header file.

3. Marking dialects as dependent in the pass TableGen file

LLVM and MLIR use Tablegen, an embedded domain-specific language (DSL), to declare compiler passes. TableGen automates the generation of required C++ boilerplate code, streamline the development process. When defining transformations via TableGen, you must explicitly register any dependent dialects required by the pass. For instance, to manipulate quantum operations using the custom plugin, the Quantum dialect must be added to the list of dependent dialects in the pass definition file:

diff --git a/mlir/standalone/include/Standalone/StandalonePasses.td b/mlir/standalone/include/Standalone/StandalonePasses.td
index dc8fb43d2..29510d74d 100644
--- a/mlir/standalone/include/Standalone/StandalonePasses.td
+++ b/mlir/standalone/include/Standalone/StandalonePasses.td
@@ -26,6 +26,10 @@ def StandaloneSwitchBarFoo: Pass<"standalone-switch-bar-foo", "::mlir::ModuleOp"
     ```
   }];

+   let dependentDialects = [
+       "catalyst::quantum::QuantumDialect"
+   ];
+
 }

 #endif // STANDALONE_PASS

Finally, the appropriate library must be linked in (MLIRQuantum for the Quantum dialect), and the plugin tool changed to catalyst-cli:

diff --git a/mlir/standalone/lib/Standalone/CMakeLists.txt b/mlir/standalone/lib/Standalone/CMakeLists.txt
index 0f1705a25..8874e410d 100644
--- a/mlir/standalone/lib/Standalone/CMakeLists.txt
+++ b/mlir/standalone/lib/Standalone/CMakeLists.txt
@@ -10,9 +10,11 @@ add_mlir_dialect_library(MLIRStandalone
         DEPENDS
         MLIRStandaloneOpsIncGen
         MLIRStandalonePassesIncGen
+        MLIRQuantum

         LINK_LIBS PUBLIC
         MLIRIR
         MLIRInferTypeOpInterface
         MLIRFuncDialect
+        MLIRQuantum
         )

diff --git a/mlir/standalone/standalone-plugin/CMakeLists.txt b/mlir/standalone/standalone-plugin/CMakeLists.txt
index 3e3383608..2dbeea9d5 100644
--- a/mlir/standalone/standalone-plugin/CMakeLists.txt
+++ b/mlir/standalone/standalone-plugin/CMakeLists.txt
@@ -5,7 +5,7 @@ add_llvm_library(StandalonePlugin
         DEPENDS
         MLIRStandalone
         PLUGIN_TOOL
-        mlir-opt
+        catalyst-cli

         LINK_LIBS
         MLIRStandalone

If the plugin depends on other dialects defined in Catalyst, they can be added as dependencies in a similar way. For instance, if the plugin depends on the Gradient dialect, add "catalyst::gradient::GradientDialect" to the dependentDialects field and link the MLIRGradient library.

4. Modify the standalone plugin to rewrite quantum operations

Here we will create a very simple pass that will change the quantum qubit allocation from 1 to 42 (for illustration purposes). We recommend reading the MLIR tutorials on how to write MLIR passes, reading the Catalyst source to understand the Catalyst IR, and submitting issues if you are having troubles building your own plugin.

The first thing we need to do is change the OpRewritePattern to match against our quantum::AllocOp, which denotes the allocation of a quantum register containing the given number of qubits.

diff --git a/mlir/standalone/lib/Standalone/StandalonePasses.cpp b/mlir/standalone/lib/Standalone/StandalonePasses.cpp
index 83e2ce255..504cf2d20 100644
--- a/mlir/standalone/lib/Standalone/StandalonePasses.cpp
+++ b/mlir/standalone/lib/Standalone/StandalonePasses.cpp
@@ -19,10 +19,10 @@ namespace mlir::standalone {
 #include "Standalone/StandalonePasses.h.inc"

 namespace {
-class StandaloneSwitchBarFooRewriter : public OpRewritePattern<func::FuncOp> {
+class StandaloneSwitchBarFooRewriter : public OpRewritePattern<catalyst::quantum::AllocOp> {
 public:
-  using OpRewritePattern<func::FuncOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(func::FuncOp op,
+  using OpRewritePattern<catalyst::quantum::AllocOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(catalyst::quantum::AllocOp op,
                                 PatternRewriter &rewriter) const final {
     if (op.getSymName() == "bar") {
       rewriter.modifyOpInPlace(op, [&op]() { op.setSymName("foo"); });

The next step is changing the contents of the function itself:

diff --git a/mlir/standalone/lib/Standalone/StandalonePasses.cpp b/mlir/standalone/lib/Standalone/StandalonePasses.cpp
index 83e2ce255..e8a7f805e 100644
--- a/mlir/standalone/lib/Standalone/StandalonePasses.cpp
+++ b/mlir/standalone/lib/Standalone/StandalonePasses.cpp
@@ -19,15 +19,21 @@ namespace mlir::standalone {
 #include "Standalone/StandalonePasses.h.inc"

 namespace {
-class StandaloneSwitchBarFooRewriter : public OpRewritePattern<func::FuncOp> {
+class StandaloneSwitchBarFooRewriter : public OpRewritePattern<catalyst::quantum::AllocOp> {
 public:
-  using OpRewritePattern<func::FuncOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(func::FuncOp op,
+  using OpRewritePattern<catalyst::quantum::AllocOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(catalyst::quantum::AllocOp op,
                                 PatternRewriter &rewriter) const final {
-    if (op.getSymName() == "bar") {
-      rewriter.modifyOpInPlace(op, [&op]() { op.setSymName("foo"); });
+    // get the number of qubits allocated
+    if (op.getNqubitsAttr().value_or(0) == 1) {
+      Type i64 = rewriter.getI64Type();
+      auto fortytwo = rewriter.getIntegerAttr(i64, 42);
+
+      // modify the allocation to change the number of qubits to 42.
+      rewriter.modifyOpInPlace(op, [&]() { op.setNqubitsAttrAttr(fortytwo); });
       return success();
     }
+    // failure indicates that nothing was modified.
     return failure();
   }
 };

And then we can run make all again. The shared object of the standalone plugin should be available in mlir/build/lib/StandalonePlugin.so. This shared object can be used with both the catalyst and quantum-opt tools. From here, you can change the name of the pass, change the name of the shared object, and implement more complex transformations.

5. Build your own Python wheel and ship your plugin

Now that you have your StandalonePlugin.so, you can ship it in a Python wheel. To allow users to run your pass, we have provided a class called Pass and PassPlugin. You can extend these classes and allow the user to import your derived classes and run passes as a decorator. We provide the apply_pass_plugin() decorator to allow pass plugins to be loaded and executed. For example:

from standalone import getStandalonePluginAbsolutePath

@apply_pass_plugin(getStandalonePluginAbsolutePath(), "standalone-switch-bar-foo")
@qp.qnode(qp.device("lightning.qubit", wires=1))
def qnode():
    return qp.state()

@qp.qjit(target="mlir")
def module():
    return qnode()

print(module.mlir)

If you have followed all the steps in this tutorial and inspect the MLIR sources, you’ll find that the number of qubits allocated will be 42. Take a look into the standalone_plugin_wheel Makefile rule to see how we test shipping a plugin. For more information, please consult our dialect guide, our compiler passes guide, and the MLIR documentation.

You can also register your pass with Catalyst via Python’s entry_points (for reference, we have an example in the Catalyst Github repository that implements the standalone plugin as a Python package). To do this, you only need to define a function named name2pass—it must be named name2pass—that takes a string with the name of the pass (from the user perspective) and returns the absolute path to the plugin stored in your package and the name of the MLIR pass. For the standalone plugin python package we defined:

def name2pass(_name):
    """Example entry point for standalone plugin"""

    return getStandalonePluginAbsolutePath(), "standalone-switch-bar-foo"

You will also need to modify your setup to include the entry_points. See our setup.py file in the standalone plugin python package.

entry_points = {
    "catalyst.passes_resolution": [
        "standalone.passes = standalone_plugin",
    ],
}

setup(
    name="standalone_plugin",
    version="0.1.0",
    # ...
    entry_points=entry_points,
    # ...
)

After this, the user will be able to use your pass with the apply_pass() function.

@apply_pass("standalone.standalone-switch-bar-foo")
@qp.qnode(qp.device("lightning.qubit", wires=1))
def qnode():
    return qp.state()

@qp.qjit(target="mlir")
def module():
    return qnode()

print(module.mlir)

Of course, you can also define your own decorators similar to apply_pass() to check parameters, do some other validation or perhaps just to improve the user interface. For example:

from standalone import SwitchBarToFoo

@SwitchBarToFoo
@qp.qnode(qp.device("lightning.qubit", wires=1))
def qnode():
    return qp.state()

@qp.qjit(target="mlir")
def module():
    return qnode()

print(module.mlir)
Previous
Next