コンパイラ/インタープリタ設計完全攻略: パーサからLLVM、AIコンパイラ(TVM/XLA)まで

ソースコード
    ↓  字句解析 (Lexing)
トークン列
    ↓  構文解析 (Parsing)
AST (抽象構文木)
    ↓  意味解析 (Semantic Analysis)
注釈付きAST
    ↓  IR生成
中間表現 (IR)
    ↓  最適化 (Optimization)
最適化されたIR
    ↓  コード生成 (Code Generation)
機械語 / アセンブリ

import re
from enum import Enum, auto
from dataclasses import dataclass
from typing import List, Optional

class TokenType(Enum):
    NUMBER   = auto()
    IDENT    = auto()
    PLUS     = auto()
    MINUS    = auto()
    STAR     = auto()
    SLASH    = auto()
    LPAREN   = auto()
    RPAREN   = auto()
    ASSIGN   = auto()
    SEMICOL  = auto()
    LET      = auto()
    EOF      = auto()

@dataclass
class Token:
    type: TokenType
    value: str
    line: int

KEYWORDS = {'let': TokenType.LET}

TOKEN_PATTERNS = [
    (r'\d+(\.\d+)?',  TokenType.NUMBER),
    (r'[a-zA-Z_]\w*', TokenType.IDENT),
    (r'\+',           TokenType.PLUS),
    (r'-',            TokenType.MINUS),
    (r'\*',           TokenType.STAR),
    (r'/',            TokenType.SLASH),
    (r'\(',           TokenType.LPAREN),
    (r'\)',           TokenType.RPAREN),
    (r'=',            TokenType.ASSIGN),
    (r';',            TokenType.SEMICOL),
]

def tokenize(source: str) -> List[Token]:
    tokens = []
    line = 1
    pos = 0
    while pos < len(source):
        if source[pos] in ' \t\r':
            pos += 1
            continue
        if source[pos] == '\n':
            line += 1
            pos += 1
            continue
        matched = False
        for pattern, ttype in TOKEN_PATTERNS:
            m = re.match(pattern, source[pos:])
            if m:
                value = m.group(0)
                if ttype == TokenType.IDENT and value in KEYWORDS:
                    ttype = KEYWORDS[value]
                tokens.append(Token(ttype, value, line))
                pos += len(value)
                matched = True
                break
        if not matched:
            raise SyntaxError(f"予期しない文字 '{source[pos]}' (行 {line})")
    tokens.append(Token(TokenType.EOF, '', line))
    return tokens

# 使用例
src = "let x = 10 + 20 * 3;"
for tok in tokenize(src):
    print(tok)

from dataclasses import dataclass
from typing import Union

# ASTノード定義
@dataclass
class NumberLit:
    value: float

@dataclass
class Identifier:
    name: str

@dataclass
class BinOp:
    op: str
    left: 'Expr'
    right: 'Expr'

@dataclass
class LetStmt:
    name: str
    value: 'Expr'

Expr = Union[NumberLit, Identifier, BinOp]

# 再帰下降パーサ
class Parser:
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0

    def peek(self) -> Token:
        return self.tokens[self.pos]

    def consume(self, expected: Optional[TokenType] = None) -> Token:
        tok = self.tokens[self.pos]
        if expected and tok.type != expected:
            raise SyntaxError(f"期待: {expected}, 実際: {tok.type} (行 {tok.line})")
        self.pos += 1
        return tok

    def parse_let(self) -> LetStmt:
        self.consume(TokenType.LET)
        name = self.consume(TokenType.IDENT).value
        self.consume(TokenType.ASSIGN)
        expr = self.parse_expr()
        self.consume(TokenType.SEMICOL)
        return LetStmt(name=name, value=expr)

    def parse_expr(self) -> Expr:
        return self.parse_additive()

    def parse_additive(self) -> Expr:
        left = self.parse_multiplicative()
        while self.peek().type in (TokenType.PLUS, TokenType.MINUS):
            op = self.consume().value
            right = self.parse_multiplicative()
            left = BinOp(op=op, left=left, right=right)
        return left

    def parse_multiplicative(self) -> Expr:
        left = self.parse_primary()
        while self.peek().type in (TokenType.STAR, TokenType.SLASH):
            op = self.consume().value
            right = self.parse_primary()
            left = BinOp(op=op, left=left, right=right)
        return left

    def parse_primary(self) -> Expr:
        tok = self.peek()
        if tok.type == TokenType.NUMBER:
            self.consume()
            return NumberLit(value=float(tok.value))
        if tok.type == TokenType.IDENT:
            self.consume()
            return Identifier(name=tok.value)
        if tok.type == TokenType.LPAREN:
            self.consume()
            expr = self.parse_expr()
            self.consume(TokenType.RPAREN)
            return expr
        raise SyntaxError(f"予期しないトークン: {tok}")

# ソース: x = (a + b) * (c - d)
t1 = a + b
t2 = c - d
x  = t1 * t2

# 通常コード
x = 1
x = x + 2   # xが2回代入される

# SSA変換後
x_1 = 1
x_2 = x_1 + 2

if cond:
    x_1 = 10
else:
    x_2 = 20
# 合流点
x_3 = phi(x_1, x_2)

; 関数定義: int add(int a, int b)
define i32 @add(i32 %a, i32 %b) {
entry:
  %result = add i32 %a, %b
  ret i32 %result
}

; 条件分岐 + SSA phi関数
define i32 @max(i32 %a, i32 %b) {
entry:
  %cmp = icmp sgt i32 %a, %b    ; a > b ?
  br i1 %cmp, label %then, label %else

then:
  br label %merge

else:
  br label %merge

merge:
  %result = phi i32 [ %a, %then ], [ %b, %else ]
  ret i32 %result
}

# 最適化前
x = 2 * 3 + 4

# 最適化後 (コンパイル時に計算済み)
x = 10

def constant_fold(node):
    """簡単な定数畳み込みの実装"""
    if isinstance(node, BinOp):
        left  = constant_fold(node.left)
        right = constant_fold(node.right)
        if isinstance(left, NumberLit) and isinstance(right, NumberLit):
            ops = {
                '+': lambda a, b: a + b,
                '-': lambda a, b: a - b,
                '*': lambda a, b: a * b,
                '/': lambda a, b: a / b,
            }
            return NumberLit(value=ops[node.op](left.value, right.value))
        return BinOp(op=node.op, left=left, right=right)
    return node

def dead_code_eliminate(stmts, used_vars):
    """使われない変数への代入を削除する後方パス"""
    live = set()
    for stmt in reversed(stmts):
        if isinstance(stmt, LetStmt):
            if stmt.name in used_vars or stmt.name in live:
                live.update(collect_vars(stmt.value))
    return [s for s in stmts
            if not isinstance(s, LetStmt) or s.name in live]

// 元のコード (ループ4回)
for (int i = 0; i < 4; i++) {
    a[i] = b[i] + c[i];
}

// アンローリング後 (ループ消去)
a[0] = b[0] + c[0];
a[1] = b[1] + c[1];
a[2] = b[2] + c[2];
a[3] = b[3] + c[3];

// インライン前
inline int square(int x) { return x * x; }
int y = square(5);

// インライン後 (定数畳み込みで25に)
int y = 25;

; C関数: int add(int a, int b) { return a + b; }
; Linux System V ABI: 引数はrdi, rsi順; 返り値はrax

add:
    push    rbp
    mov     rbp, rsp
    mov     DWORD PTR [rbp-4], edi
    mov     DWORD PTR [rbp-8], esi
    mov     edx, DWORD PTR [rbp-4]
    mov     eax, DWORD PTR [rbp-8]
    add     eax, edx
    pop     rbp
    ret

; -O2最適化版 (1命令のみ!)
add:
    lea     eax, [rdi + rsi]
    ret

import numpy as np
from numba import njit, prange
import time

# 純粋Pythonベースライン
def matmul_python(A, B):
    N = A.shape[0]
    C = np.zeros((N, N))
    for i in range(N):
        for j in range(N):
            for k in range(N):
                C[i, j] += A[i, k] * B[k, j]
    return C

# Numba JIT版 (並列化あり)
@njit(parallel=True)
def matmul_numba(A, B):
    N = A.shape[0]
    C = np.zeros((N, N))
    for i in prange(N):       # 並列ループ
        for j in range(N):
            for k in range(N):
                C[i, j] += A[i, k] * B[k, j]
    return C

N = 256
A = np.random.rand(N, N).astype(np.float32)
B = np.random.rand(N, N).astype(np.float32)

# 初回呼び出しでコンパイル (ウォームアップ)
_ = matmul_numba(A, B)

t0 = time.perf_counter()
C_py = matmul_python(A, B)
print(f"Python:    {time.perf_counter()-t0:.3f}s")

t0 = time.perf_counter()
C_nb = matmul_numba(A, B)
print(f"Numba JIT: {time.perf_counter()-t0:.4f}s")
# 典型的な結果: Python ~30s vs Numba ~0.01s (数百倍高速)

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128, 256)
        self.fc2 = nn.Linear(256, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = SimpleNet().cuda()
x = torch.randn(32, 128).cuda()

# torch.compile適用
compiled = torch.compile(model, mode='max-autotune')

# 初回呼び出し: コンパイル発生 (数秒かかる)
out = compiled(x)

# 以降の呼び出し: コンパイル済みカーネルを再利用
for _ in range(100):
    out = compiled(x)

import torch._dynamo as dynamo

def my_func(x, y):
    return torch.sin(x) + torch.cos(y)

# FXグラフの確認
explanation = dynamo.explain(my_func)(
    torch.randn(4), torch.randn(4)
)
print(explanation.graphs[0].print_tabular())
# opcode  | name   | target       | args
# --------|--------|--------------|----------
# call_fn | sin    | torch.sin    | (x,)
# call_fn | cos    | torch.cos    | (y,)
# call_fn | add    | operator.add | (sin, cos)
# output  | output | output       | (add,)

import tvm
from tvm import relay
import torch, torchvision
import numpy as np

# PyTorchモデルから開始
model = torchvision.models.resnet18(pretrained=False).eval()
input_shape = [1, 3, 224, 224]
trace_input = torch.randn(input_shape)

# TorchScriptを経由してTVM Relay IRへ変換
scripted = torch.jit.trace(model, trace_input)
mod, params = relay.frontend.from_pytorch(
    scripted, [("input0", input_shape)]
)

target = tvm.target.Target("cuda")

# AutoScheduler: 最適カーネルパラメータの探索
from tvm import auto_scheduler
tasks, task_weights = auto_scheduler.extract_tasks(
    mod["main"], params, target
)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
    num_measure_trials=200,
    runner=auto_scheduler.LocalRunner(repeat=10),
    measure_callbacks=[auto_scheduler.RecordToFile("tuning.json")],
)
tuner.tune(tune_option)

# 最良スケジュールでコンパイル
with auto_scheduler.ApplyHistoryBest("tuning.json"):
    with tvm.transform.PassContext(opt_level=3):
        lib = relay.build(mod, target=target, params=params)

print("TVMコンパイル完了!")

# フュージョン前: 3回のGPUカーネル起動
y1 = conv2d(x)        # カーネル1 — HBMへ書き込み
y2 = batch_norm(y1)   # カーネル2 — HBM読み書き
y3 = relu(y2)         # カーネル3 — HBM読み書き

# フュージョン後: 1回のGPUカーネル起動
y3 = fused_conv_bn_relu(x)
# 中間結果はL1/L2キャッシュに留まる
# HBM帯域幅の節約 = 大幅な高速化

import jax
import jax.numpy as jnp

@jax.jit
def compute(x, W, b):
    # XLAが自動的にmatmul + add + reluをフューズ
    return jax.nn.relu(x @ W + b)

x = jnp.ones((32, 128))
W = jnp.ones((128, 256))
b = jnp.zeros(256)

# 初回呼び出し: XLAコンパイル発生
result = compute(x, W, b)
print(result.shape)  # (32, 256)

TensorFlow Graph
     ↓  TF Dialect
MLIR HLO Dialect
     ↓  Linalg Dialect
ループネスト IR
     ↓  Affine Dialect
アフィンループ IR
     ↓  LLVM Dialect
LLVM IR
     ↓
機械語