Avant de plonger dans le monde fascinant du compilateur Cuda-Q, j’ai pensé qu’il serait utile de rédiger une petite note sur l’art de la réécriture d’AST.
On peut se demander : quelle est la différence entre la réécriture d’AST et la compilation ? En bref, toute compilation inclut une réécriture d’AST, mais toute réécriture d’AST n’est pas considérée comme une compilation.
De plus, la compilation est le processus de conversion de code d’un langage à un autre, tandis que la réécriture AST reste généralement au sein du même langage.
API d’arbre de syntaxe abstraite (AST) Python
Link to heading
Python fournit une bibliothèque standard pour la manipulation des AST. Un AST Python comporte quatre grandes catégories de nœuds :
Par exemple 10, "bonjour", Vrai.
Constant(value=10)
Constant(value="hello")
Constant(value=True)
Par exemple x, y, z. Les variables peuvent être accessibles en mode lecture (chargement) ou en mode écriture (stockage).
Name(id='x', ctx=Load())
Name(id='y', ctx=Store())
Par exemple, 1 + 2, x + 2, x > y. Les expressions sont des nœuds AST qui produisent une valeur.
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())])
Par exemple, if, for, while, return, break, continue. Les instructions sont des nœuds AST qui effectuent des actions et apparaissent comme des constructions de niveau supérieur ou de niveau bloc dans le code Python.
Par exemple, le code suivant :
est équivalent à :
If(
test=Compare(left=Name("x", Load()),ops=[Gt()],comparators=[Constant(value=1)]),
body=[Pass()],
orelse=[]
)
Exemple concret : Les conditionnelles dans les noyaux quantiques
Link to heading
Imaginons que nous voulions transformer ce code :
@kernel
def conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
if iq.i > 0.5:
qubit.main.play("waveform")
dans
def kernelized_conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
with cc._if(iq.i > 0.5):
qubit.main.play('waveform')
L’idée est que iq est une Readout d’exécution qui ne doit pas être évaluée par l’interpréteur Python, mais plutôt par le contrôleur quantique. Pour que cela fonctionne, il faut créer un décorateur 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)
Les fonctions decompile, parse_to_ast et recompile ne sont que de simples wrappers autour des fonctions de la bibliothèque standard (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)
Le transformateur est une sous-classe de ast.NodeTransformer qui parcourt l’AST et le transforme en un générateur pour le modèle. Il s’agit d’un modèle de conception très générique qui peut être utilisé pour transformer n’importe quel AST en n’importe quel autre 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
L’introduction du concept de décorateur « noyau » pour transformer une instruction Python en ISA QCS peut prêter à confusion. Que se passe-t-il si l’on oublie d’appliquer ce décorateur ? Le code du noyau ne fonctionnera pas, car l’évaluation sera effectuée par l’interpréteur Python, ce qui peut s’avérer très déroutant pour le développeur.
Il existe cependant de nombreuses solutions à ce problème. On peut utiliser un outil d’analyse statique, ou linter, pour vérifier si la fonction est décorée avec le décorateur kernel. Cette approche peut s’avérer très efficace, notamment lorsqu’on commence à utiliser des LLM pour automatiser la réécriture de code.
(source de l’image : Intégration des modifications de code contextuelles)
Dans cette brève note, j’ai démontré qu’il est possible de réécrire l’AST d’une fonction Python de manière très simple. Grâce à cette méthode, il devient beaucoup plus facile d’écrire des noyaux quantiques exécutables sur la pile de contrôle quantique, sans avoir recours à la syntaxe complexe with _xxx.
On pourrait se demander quel est l’intérêt de manipuler l’AST en Python, étant donné que l’industrie évolue rapidement vers une pile unifiée basée sur LLVM avec des dialectes MLIR pour l’informatique quantique. C’est tout à fait exact, et dans la prochaine note, je montrerai comment utiliser MLIR pour obtenir le même résultat.
Je devrai également examiner le compilateur NAC3 de M-Labs. Contrairement à CUDAQ, NAC3 utilise une implémentation basée sur Rust. Même le code Python est converti en AST à l’aide d’un analyseur syntaxique Rust. Cependant, la génération de code est déléguée à LLVM, du moins pour son cœur logiciel RISC-V.
En attendant, cette note a été une manière amusante d’en apprendre davantage sur l’AST de Python, et cela peut être utile pour améliorer les compilateurs de circuits quantiques existants basés sur Python qui n’utilisent pas encore LLVM !