Bevor ich mich eingehend mit der faszinierenden Welt des Cuda-Q-Compilers (https://github.com/NVIDIA/cuda-quantum/blob/main/tools/nvqpp/nvq%2B%2B.in) befasse, möchte ich kurz die Kunst des AST-Rewritings erläutern.
Man könnte sich fragen: Worin besteht der Unterschied zwischen AST-Rewriting und Kompilierung? Kurz gesagt: Jede Kompilierung beinhaltet AST-Rewriting, aber nicht jedes AST-Rewriting zählt als Kompilierung.
Die Kompilierung ist der Prozess, Code von einer Sprache in eine andere zu übersetzen, während die AST-Neuschreibung in der Regel innerhalb derselben Sprache erfolgt.
Python bietet eine Standardbibliothek für die Arbeit mit ASTs. Es gibt im Wesentlichen vier Knotenkategorien in einem Python-AST:
Zum Beispiel 10, "hello", True.
Constant(value=10)
Constant(value="hello")
Constant(value=True)
Zum Beispiel x, y, z. Variablen können entweder im Lesemodus (Laden) oder im Schreibmodus (Speichern) aufgerufen werden.
Name(id='x', ctx=Load())
Name(id='y', ctx=Store())
Zum Beispiel 1 + 2, x + 2, x > y. Ausdrücke sind AST-Knoten, die einen Wert erzeugen.
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())])
Beispiele hierfür sind if, for, while, return, break und continue. Anweisungen sind AST-Knoten, die Aktionen ausführen und in Python-Code als Konstrukte auf oberster Ebene oder Blockebene auftreten.
Zum Beispiel folgender Code:
ist gleichbedeutend mit:
If(
test=Compare(left=Name("x", Load()),ops=[Gt()],comparators=[Constant(value=1)]),
body=[Pass()],
orelse=[]
)
Konkretes Beispiel: Konditionale Ausdrücke in Quantenkernen
Link zu Überschrift
Stellen wir uns vor, wir möchten diesen Code umwandeln:
@kernel
def conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
if iq.i > 0.5:
qubit.main.play("waveform")
hinein
def kernelized_conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
with cc._if(iq.i > 0.5):
qubit.main.play('waveform')
Die Idee ist, dass iq ein Laufzeit-Readout ist, das nicht vom Python-Interpreter, sondern vom Quantencontroller ausgewertet werden soll. Damit dies funktioniert, muss ein kernel-Dekorator erstellt werden:
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)
Die Funktionen decompile, parse_to_ast und recompile sind lediglich schlanke Wrapper um die Standardbibliotheksfunktionen (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)
Der Transformer ist eine Unterklasse von ast.NodeTransformer, die den AST durchläuft und ihn in einen Generator für das Modell umwandelt. Es handelt sich dabei um ein sehr generisches Entwurfsmuster, mit dem sich jeder beliebige AST in einen anderen AST transformieren lässt.
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
Die Einführung des Konzepts des „Kernel“-Dekorators zur Transformation von Python-Anweisungen in QCS-ISA kann verwirrend sein. Was passiert, wenn man vergisst, den Dekorator anzuwenden? Der Kernel-Code würde nicht funktionieren, da die Auswertung vom Python-Interpreter durchgeführt wird, was für den Entwickler sehr verwirrend sein kann.
Es gibt jedoch viele Lösungsansätze für dieses Problem. Man könnte beispielsweise ein statisches Analysetool, einen sogenannten Linter, verwenden, um zu prüfen, ob die Funktion mit dem kernel-Dekorator versehen ist. Dieser Ansatz kann sehr leistungsfähig sein, insbesondere beim Einsatz von LLMs zur Automatisierung von Code-Rewriting.
(Bildquelle: Context-aware code change embedding)
In diesem kurzen Memo habe ich gezeigt, dass sich der AST einer Python-Funktion auf sehr einfache Weise umschreiben lässt. Mit dieser Methode wird es deutlich einfacher, Quantenkerne zu schreiben, die auf dem Quantenkontrollstapel ausgeführt werden können, ohne die umständliche with _xxx-Syntax verwenden zu müssen.
Man könnte sich natürlich fragen, welchen Sinn die Manipulation des AST in Python hat, da die Branche sich rasant in Richtung eines einheitlichen LLVM-basierten Stacks mit MLIR-Dialekten für Quantencomputing entwickelt. Das stimmt in der Tat, und im nächsten Dokument zeige ich, wie man mit MLIR dasselbe Ergebnis erzielt.
Ich werde mir auch den NAC3-Compiler von M-Labs ansehen müssen. Im Vergleich zu CudaQ verwendet NAC3 eine Rust-basierte Implementierung. Selbst der Python-Code wird mithilfe eines Rust-Parsers in einen AST konvertiert. Die Codegenerierung wird jedoch an LLVM delegiert, zumindest für deren RISC-V-Soft-Core.
In der Zwischenzeit war dieses Memo eine unterhaltsame Möglichkeit, mehr über den Python-AST zu erfahren, und dies kann nützlich sein, um bestehende Python-basierte Quantenschaltungs-Compiler zu verbessern, die noch nicht LLVM verwenden!