Prima di immergermi nel fantastico mondo del compilatore Cuda-Q, ho pensato che valesse la pena scrivere un breve promemoria sull’arte della riscrittura dell’AST. Abstract Syntax Tree

Ci si potrebbe chiedere: qual è la differenza tra riscrittura dell’AST e compilazione? In breve, ogni compilazione include una riscrittura dell’AST, ma non ogni riscrittura dell’AST è da considerarsi una compilazione.

Inoltre, la compilazione è il processo di abbassamento del codice da un linguaggio a un altro, mentre la riscrittura dell’AST di solito rimane all’interno dello stesso linguaggio.

API dell’albero sintattico astratto (AST) di Python Link to heading

Python fornisce una libreria standard per lavorare con gli AST. Esistono all’incirca quattro categorie di nodi in un AST Python:

Letterali o costanti Link to heading

Ad esempio 10, "ciao", vero.

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

Variabili Link to heading

Ad esempio x, y, z. Le variabili possono essere accessibili in modalità di lettura (caricamento) o scrittura (memorizzazione).

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

Espressioni Link to heading

Ad esempio 1 + 2, x + 2, x > y. Le espressioni sono nodi AST che producono un valore.

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())])

Dichiarazioni Link to heading

Ad esempio, if, for, while, return, break, continue. Le istruzioni sono nodi AST che eseguono azioni e appaiono come costrutti di livello superiore o a livello di blocco nel codice Python.

Ad esempio, il seguente codice:

if x > 1:
    pass

equivale a:

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

Esempio concreto: Condizionali nei kernel quantistici Link to heading

Immaginiamo di voler trasformare questo codice:

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

in

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

L’idea è che iq sia un Readout a tempo di esecuzione che non dovrebbe essere valutato dall’interprete Python, bensì dal controller quantistico. Affinché ciò funzioni, è necessario creare un decoratore kernel:

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)

Le funzioni decompile, parse_to_ast e recompile sono semplicemente dei wrapper attorno alle funzioni standard della libreria (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)

Il trasformatore è una sottoclasse di ast.NodeTransformer che attraversa l’AST e lo trasforma in un generatore per il modello. Si tratta di un modello di progettazione molto generico che può essere utilizzato per trasformare qualsiasi AST in qualsiasi altro 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

Introdurre il concetto di decoratore “kernel” per trasformare un’istruzione Python in un’architettura ISA QCS può essere fonte di confusione. Cosa succede se ci si dimentica di applicare il decoratore? Il codice kernel non funzionerebbe, poiché la valutazione verrebbe eseguita dall’interprete Python, il che potrebbe essere molto disorientante per lo sviluppatore.

Esistono però diverse soluzioni a questo problema. Si potrebbe utilizzare uno strumento di analisi statica, o linter, per verificare se la funzione è decorata con il decoratore kernel. Questo approccio può rivelarsi molto efficace, soprattutto quando si inizia a utilizzare i LLM per automatizzare la riscrittura del codice.

Incorporamento del cambiamento di codice sensibile al contesto (fonte dell’immagine: Incorporamento del cambiamento di codice sensibile al contesto)

Conclusion Link to heading

In questa breve nota, ho dimostrato che è possibile riscrivere l’AST di una funzione Python in modo molto semplice. Utilizzando questo metodo, diventa molto più facile scrivere kernel quantistici eseguibili sullo stack di controllo quantistico, senza dover ricorrere alla barbara sintassi with _xxx.

Certo, ci si potrebbe chiedere a cosa serva manipolare l’AST in Python, visto che il settore si sta muovendo rapidamente verso uno stack unificato basato su LLVM con dialetti MLIR per il calcolo quantistico. Questo è effettivamente vero, e nel prossimo promemoria mostrerò come utilizzare MLIR per ottenere lo stesso risultato.

Dovrò anche dare un’occhiata al compilatore NAC3 di M-Labs. Rispetto a CudaQ, NAC3 utilizza un’implementazione basata su Rust. Anche il codice Python viene convertito in AST utilizzando un parser Rust. Tuttavia, la generazione del codice è delegata a LLVM, almeno per il loro soft-core RISC-V.

Nel frattempo, questo promemoria si è rivelato un modo divertente per imparare a conoscere l’AST di Python, e questo può essere utile per migliorare i compilatori di circuiti quantistici basati su Python che non utilizzano ancora LLVM!

References Link to heading