Cuda-Qコンパイラの素晴らしい世界に深く踏み込む前に、AST書き換えの技術について簡単なメモを書いておく価値があると思いました。
AST書き換えとコンパイルの違いは何だろうか？と疑問に思う人もいるかもしれない。簡単に言えば、すべてのコンパイルにはAST書き換えが含まれるが、すべてのAST書き換えがコンパイルとみなされるわけではない。
また、コンパイルはコードをある言語から別の言語に「変換」するプロセスであるのに対し、AST書き換えは通常、同じ言語内で行われます。
Python 抽象構文木 (AST) API
見出しへのリンク
Pythonには、ASTを扱うための標準ライブラリが用意されています。PythonのASTには、大きく分けて4つのカテゴリのノードがあります。
例えば、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を走査してモデルの生成元に変換するast.NodeTransformerのサブクラスです。これは非常に汎用的な設計パターンであり、任意の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構文を使用することなく、量子制御スタック上で実行可能な量子カーネルをはるかに簡単に記述できるようになります。
もちろん、量子コンピューティング向けにMLIR方言を備えた統一LLVMベースのスタックへと業界が急速に移行している中で、PythonでASTを操作することに何の意味があるのかと疑問に思う人もいるでしょう。これは確かにその通りで、次のメモでは、同じ結果を得るためにMLIRを使用する方法を紹介します。
M-LabsのNAC3コンパイラも確認する必要があります。CUDAQと比較すると、NAC3はRustベースの実装を使用しています。Pythonコードでさえ、Rustパーサーを使用してASTに変換されます。ただし、コード生成は、少なくともRISC-Vソフトコアに関してはLLVMに委任されています。
その間、このメモはPythonのASTについて楽しく学ぶ良い機会となり、LLVMをまだ使用していない既存のPythonベースの量子回路コンパイラを改善するのに役立つ可能性があります。