Before deep-diving into the fantastic world of the Cuda-Q compiler, I thought it would be worth writing a small memo about the art of AST rewriting. Abstract Syntax Tree

One may wonder: what is the difference between AST rewriting and compilation? In short, every compilation includes AST rewriting, but not every AST rewrite counts as compilation.

Also, compilation is the process of lowering code from one language into another one, while AST rewrite will usually stay within the same language

Python abstract syntax tree (AST) API Link to heading

Python provides a standard library for working with ASTs. There are roughly four categories of nodes in a Python AST:

Literals or Constants Link to heading

For example 10, "hello", True.

Constant(value=10)
Constant(value="hello")
Constant(value=True)

Variables Link to heading

For example x, y, z. Variables can be either access in read (load) or write (store) mode.

Name(id='x', ctx=Load())
Name(id='y', ctx=Store())

Expressions Link to heading

For example 1 + 2, x + 2, x > y. Expressions are AST nodes that produce a value.

BinOp(left=Constant(value=1), op=Add(), right=Constant(value=2))
BinOp(left=Name("x", Load()),op=Add(),right=Constant(value=2))
Compare(left=Name("x", Load()),ops=[Gt()],comparators=[Name(id='y', ctx=Load())])

Statements Link to heading

For example if, for, while, return, break, continue. Statements are AST nodes that perform actions and appear as top-level or block-level constructs in Python code.

For example, the following code:

if x > 1:
    pass

is equivalent to:

If(
    test=Compare(left=Name("x", Load()),ops=[Gt()],comparators=[Constant(value=1)]),
    body=[Pass()],
    orelse=[]
)

Concrete Example: Conditionals in quantum kernels Link to heading

Let’s imagine we want to turn this code:

@kernel
def conditional_play(qubit: Qubit):
    iq = qubit.ancilla.readout()
    if iq.i > 0.5:
        qubit.main.play("waveform")

into

def kernelized_conditional_play(qubit: Qubit):
    iq = qubit.ancilla.readout()
    with cc._if(iq.i > 0.5):
        qubit.main.play('waveform')

The idea is that iq is a run-time Readout that should not be evaluated by the phyton interpreter, but rather by the quantum controller. For this to work, one needs to create a kernel decorator:

class kernel:

    def __init__(self, func):
        self.func = func

        # Uncompile wrapped function (convert it into a string)
        source = self.decompile(func)

        # Parse the string into an AST
        tree = self.parse_to_ast(source)

        # Transform the AST, converting the "if" into "with if_()"
        tree = Transformer().visit(tree)

        # Recompile the AST into a binary
        binary = self.recompile(tree)

        # Make a namespace for execution
        namespace = func.__globals__.copy()
        # This does not really executes the function, but rather creates 
        # a new function based on the modified AST
        exec(binary, None, namespace)

        # The new function is now available in the namespace
        self.kernel  = namespace["kernelized_" + func.__name__]

    def __call__(self, *args, **kwargs):
        return self.kernel(*args, **kwargs)

The decompile, parse_to_ast and recompile are just thin wrappers around the standard library functions (inspect, compile).

def decompile(self, func):
    return inspect.getsource(func.__code__)

def parse_to_ast(self, source: str):
    return compile(source, filename="<generated>", mode="exec", flags=ast.PyCF_ONLY_AST, dont_inherit=True)

def recompile(self, tree):
    return compile(tree, filename="<generated>", mode="exec", dont_inherit=True)

The transformer is a subclass of ast.NodeTransformer that traverses the AST and transforms it into a generator for the model. That is a very generic design pattern that can be used to transform any AST into any other AST.


class Transformer(ast.NodeTransformer):
    """
    This subclass traverses the AST of the user-written, decorated,
    model specification and transforms it into a generator for the
    model. Subclassing in this way is the idiomatic way to transform
    an AST.

    Specifically:

    1. rewrite all `if` statements into `with cc._if()` blocks
    2. rename the function to `kernelized_` + original function name
    3. Remove the @kernel decorator to prevent from recusion
    """

    def visit_If(self, node):
        self.generic_visit(node)
        modified_node = ast.With(
            items=[
                ast.withitem(
                    context_expr=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id="cc", ctx=ast.Load()),
                            attr="_if",
                            ctx=ast.Load(),
                        ),
                        args=[node.test],
                        keywords=[],
                    ),
                    optional_vars=None,
                )
            ],
            body=node.body,
        )

        ast.copy_location(modified_node, node)
        ast.fix_missing_locations(modified_node)
        return modified_node

    def visit_FunctionDef(self, node):
        modified_node = node
        # Rename the function to `kernelized_` + original function name
        modified_node.name = "kernelized_" + node.name
        # Remove the @kernel decorator to prevent from recusion
        modified_node.decorator_list = []

        # Copy the source location of the original node
        ast.copy_location(modified_node, node)
        ast.fix_missing_locations(modified_node)

        # Do not forget to visit the children of the node
        self.generic_visit(node)

        return modified_node

Digression Link to heading

Introducing the concept “kernel” decorator to transform python statement into QCS ISA can be confusing. What about if one forgets to apply the decorator? The kernel code would not work, due to the evaluation being done by the Python interpreter, which could be very confusing for the developper.

But there are many solutions to this problem. One could use a static analysis tool, or linter, to check if the function is decorated with the kernel decorator. Going in this direction can be quite powerful, especially when starting to use LLMs to automate code rewrite.

Context-aware code change embedding (image source: Context-aware code change embedding)

Conclusion Link to heading

In this quick memo, I have shown that it is possible to rewrite the AST of a Python function in a very simple way. Using this method, it becomes much easier to write quantum kernels that can be executed on the quantum control stack, without having to use barbarian with _xxx syntax.

Of course, one may ask what the point of bothering to manipulate the AST in Python is, since the industry is moving fast towards a unified LLVM-based stack with MLIR dialects for quantum computing. This is indeed true, and in the coming memo, I’ll show how to use MLIR to achieve the same result.

I will also need to have a look at the NAC3 compiler from M-Labs. Compared to CudaQ, NAC3 is using a Rust-based implementation. Even the Python code is converted to AST using a Rust parser. However, the code-generation is delegated to LLVM, at least for their RISC-V soft-core.

In the meantime, this memo has been a fun way to learn about the Python AST, and this can be useful for improving existing Python-based quantum circuit compilers that do not yet use LLVM!

References Link to heading