在深入探索 CUDA-Q 编译器的奇妙世界之前，我想先写一篇关于抽象语法树 (AST) 重写的简短备忘录。
有人可能会问：AST重写和编译有什么区别？简而言之，每次编译都包含AST重写，但并非每次AST重写都算作编译。
此外，编译是将代码从一种语言转换为另一种语言的过程，而抽象语法树（AST）重写通常会在同一种语言内进行。
Python抽象语法树（AST）API
链接到标题
Python 提供了一个用于处理抽象语法树 (AST) 的标准库。Python AST 中的节点大致分为四类：
例如 10、"hello"、True。
Constant(value=10)
Constant(value="hello")
Constant(value=True)
例如 x、y、z。变量既可以以读取（加载）模式访问，也可以以写入（存储）模式访问。
Name(id='x', ctx=Load())
Name(id='y', ctx=Store())
例如 1 + 2、x + 2、x > y。表达式是 AST 节点，用于生成值。
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())])
例如 if、for、while、return、break、continue。语句是抽象语法树 (AST) 节点，用于执行操作，并在 Python 代码中作为顶层或代码块级别的结构出现。
例如，以下代码：
相当于：
If(
test=Compare(left=Name("x", Load()),ops=[Gt()],comparators=[Constant(value=1)]),
body=[Pass()],
orelse=[]
)
具体示例：量子核中的条件语句
链接到标题
假设我们要将这段代码转换成：
@kernel
def conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
if iq.i > 0.5:
qubit.main.play("waveform")
进入
def kernelized_conditional_play(qubit: Qubit):
iq = qubit.ancilla.readout()
with cc._if(iq.i > 0.5):
qubit.main.play('waveform')
其理念是，iq 是一个运行时 Readout，它不应该由 Python 解释器进行评估，而应该由量子控制器进行评估。为了实现这一点，需要创建一个 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)
decompile、parse_to_ast 和 recompile 只是对标准库函数（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)
转换器是 ast.NodeTransformer 的子类，它遍历抽象语法树 (AST) 并将其转换为模型的生成器。这是一种非常通用的设计模式，可用于将任何 AST 转换为任何其他 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
引入“内核”装饰器的概念来将 Python 语句转换为 QCS ISA 可能会令人困惑。如果忘记应用装饰器会怎样？由于 Python 解释器会执行求值操作，内核代码将无法运行，这可能会让开发人员感到非常困惑。
但这个问题有很多解决方案。可以使用静态分析工具（或代码检查器）来检查函数是否使用了 kernel 装饰器。这种方法非常有效，尤其是在开始使用 LLM（生命周期管理）自动化代码重写时。
(图像来源：上下文感知代码更改嵌入)
＃ 结论
在这份简短的备忘录中，我展示了如何以非常简单的方式重写 Python 函数的抽象语法树 (AST)。使用这种方法，编写可在量子控制栈上执行的量子内核就变得容易得多，而无需使用繁琐的 with _xxx 语法。
当然，有人可能会问，既然业界正在快速转向基于LLVM的统一技术栈，并采用MLIR方言进行量子计算，那么在Python中操作抽象语法树（AST）的意义何在呢？的确如此，在接下来的备忘录中，我将展示如何使用MLIR来实现同样的目标。
我还需要研究一下 M-Labs 的 NAC3 编译器。与 CUDAQ 相比，NAC3 使用的是基于 Rust 的实现。甚至连 Python 代码也是通过 Rust 解析器转换为抽象语法树 (AST) 的。不过，至少在他们的 RISC-V 软核中，代码生成是委托给 LLVM 的。
与此同时，这份备忘录提供了一种有趣的学习 Python AST 的方式，这对于改进尚未使用 LLVM 的现有基于 Python 的量子电路编译器非常有用！
＃ 参考