Users can perform runtime validation within QJIT functions

[dime10]

Context

JITing Python functions can provide order(s) of magnitude performance boost during execution, by extracting a representation of the computation and turning it into optimized binary code. However, it also makes it more difficult to understand and debug the program when things go wrong, for instance by not having access to the Python debugger or being able to easily reason about the program's execution.

To counter-act this, we can introduce new utility functions for users that help increase the confidence in their programs and detect error states at runtime.

Goal

We would like to add a debugging function catalyst.assert(bool, str) that asserts at runtime that the given condition holds, and if not raises the provided error message.

@qjit
def f(x: int):
    catalyst.assert(x < 5, "expected x to be less than 5")
    return 2*x

>>> f(3)
6
>>> f(5)
RuntimeError: expected x to be less than 5

Requirements:

• The goal is to implement this functionality using the Catalyst runtime, without relying on callbacks into Python.
• The assertion should raise a C++ exception, which can be turned into a Python error by Pybind11, rather than using C assertions which would terminate the Python kernel.
• An MLIR pass is available that traverses a compiled program and removes all assertions from the IR.
• The pass is conditionally run based on an option to the compiler @qjit(disable_assertions=True) which is off by default

Technical Details

Generally speaking adding a new operator to Catalyst involves providing a representation across its different representations:

• a public API function in Python: This can be provided in a new module within the api_extensions. Public API functions are responsible to (eventually) generate the JAXPR primitive in the next step.
• a JAXPR primitive: A primitive includes a name, how input types map to output types, and how to generate an MLIR operation using the Python bindings for MLIR. The python bindings are a build artifact typically located at catalyst/mlir/build/python_packages/quantum/mlir_quantum/dialects/_catalyst_ops_gen.py.
• an MLIR operation: The operation can be added to the existing Catalyst utility dialect, where we define operations using the ODS tablegen format from MLIR.
  • The operation will also need a lowering rule to the LLVM dialect, which should target the runtime function in the next step. The catalyst->llvm dialect pass runs patterns to perform this lowering, and a new one should be provided for the assert operation.
• a runtime function in our C-API: Once a compiled function has been turned into object code, we link it against the runtime library which provides definitions for all QIR-style and runtime functions. For quantum instructions the runtime provides an interface to actual backend devices, while most classical functions are implemented in the runtime itself. The implementation should live in the corresponding .cpp file.

Once the Operation can be created and lowered all the way to the runtime, we want to make sure we are able to disable it using the aforementioned compiler pass. Write a small transformation pass using the pattern rewriter in MLIR that eliminates all instances of the assert operations from the IR. It can then be included as an optional pass in the Quantum Compilation Pass Pipeline.

Installation Help

Refer to the Catalyst installation guide for how to install a source build of the project.