Antes de adentrarme en el fantástico mundo del compilador Cuda-Q, pensé que valdría la pena escribir una breve nota sobre el arte de la reescritura de AST.
Cabe preguntarse: ¿cuál es la diferencia entre la reescritura del AST y la compilación? En resumen, toda compilación incluye la reescritura del AST, pero no toda reescritura del AST se considera compilación.
Además, la compilación es el proceso de convertir el código de un lenguaje a otro, mientras que la reescritura del AST generalmente se mantiene dentro del mismo lenguaje.
API del árbol de sintaxis abstracta (AST) de Python
Link to heading
Python proporciona una biblioteca estándar para trabajar con AST. Hay aproximadamente cuatro categorías de nodos en un AST de Python:
Por ejemplo, 10, "hola", Verdadero.
Constant(value=10)
Constant(value="hello")
Constant(value=True)
Por ejemplo, x, y, z. Se puede acceder a las variables en modo de lectura (carga) o de escritura (almacenamiento).
Name(id='x', ctx=Load())
Name(id='y', ctx=Store())
Por ejemplo, 1 + 2, x + 2, x > y. Las expresiones son nodos AST que producen un valor.
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())])
Por ejemplo, if, for, while, return, break, continue. Las sentencias son nodos AST que realizan acciones y aparecen como construcciones de nivel superior o de bloque en el código Python.
Por ejemplo, el siguiente código:
es equivalente a:
If(
test=Compare(left=Name("x", Load()),ops=[Gt()],comparators=[Constant(value=1)]),
body=[Pass()],
orelse=[]
)
Ejemplo concreto: Condicionales en núcleos cuánticos
Link to heading
Imaginemos que queremos transformar este código:
@kernel
def conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
if iq.i > 0.5:
qubit.main.play("waveform")
en
def kernelized_conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
with cc._if(iq.i > 0.5):
qubit.main.play('waveform')
La idea es que iq es una Readout en tiempo de ejecución que no debe ser evaluada por el intérprete de pyton, sino por el controlador cuántico. Para que esto funcione, es necesario crear un decorador 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)
Las funciones decompile, parse_to_ast y recompile son simplemente envolturas delgadas alrededor de las funciones de la biblioteca estándar (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)
El transformador es una subclase de ast.NodeTransformer que recorre el AST y lo transforma en un generador para el modelo. Este es un patrón de diseño muy genérico que se puede usar para transformar cualquier AST en cualquier otro 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
Introducir el concepto de decorador “kernel” para transformar una instrucción de Python en una ISA QCS puede resultar confuso. ¿Qué ocurre si se olvida aplicar el decorador? El código kernel no funcionaría, ya que la evaluación la realiza el intérprete de Python, lo que podría resultar muy confuso para el desarrollador.
Pero existen muchas soluciones para este problema. Se podría usar una herramienta de análisis estático, o linter, para comprobar si la función está decorada con el decorador kernel. Esta opción puede ser muy útil, sobre todo al empezar a usar LLM para automatizar la reescritura de código.
(Fuente de la imagen: Incrustación de cambios de código sensible al contexto)
En este breve apunte, he demostrado que es posible reescribir el AST de una función de Python de una manera muy sencilla. Mediante este método, resulta mucho más fácil escribir kernels cuánticos que se pueden ejecutar en la pila de control cuántico, sin tener que usar la sintaxis engorrosa with _xxx.
Por supuesto, cabe preguntarse qué sentido tiene manipular el AST en Python, dado que la industria avanza rápidamente hacia una arquitectura unificada basada en LLVM con dialectos MLIR para la computación cuántica. Esto es cierto, y en el próximo informe mostraré cómo usar MLIR para lograr el mismo resultado.
También tendré que echar un vistazo al compilador NAC3 de M-Labs. En comparación con CudaQ, NAC3 utiliza una implementación basada en Rust. Incluso el código Python se convierte a AST mediante un analizador sintáctico de Rust. Sin embargo, la generación de código se delega a LLVM, al menos para su núcleo blando RISC-V.
Mientras tanto, este memorándum ha sido una forma divertida de aprender sobre el AST de Python, ¡y esto puede ser útil para mejorar los compiladores de circuitos cuánticos basados en Python que aún no utilizan LLVM!