def compile(
    model: Callable | None = None,
    *,
    fullgraph: bool = False,
    dynamic: bool | None = None,
    backend: str | Callable = "inductor",
    mode: str | None = None,
    options: dict | None = None,
    disable: bool = False,
) -> Callable:

# torch/__init__.py 第 2758 行
return torch._dynamo.optimize(
    backend=backend,
    nopython=fullgraph,
    dynamic=dynamic,
    disable=disable,
    guard_filter_fn=guard_filter_fn,
)(model)

import torch

# 方式一：装饰器
@torch.compile
def fn_a(x):
    return torch.sin(x) + torch.cos(x)

# 方式二：显式调用
def fn_b(x):
    return torch.sin(x) + torch.cos(x)

compiled_fn_b = torch.compile(fn_b)

# 方式三：编译 nn.Module
model = torch.nn.Linear(128, 64)
compiled_model = torch.compile(model)

class _TorchCompileInductorWrapper:
    def apply_mode(self, mode: str | None):
        if mode and mode != "default":
            from torch._inductor import list_mode_options
            self.apply_options(list_mode_options(mode, self.dynamic))

# 使用不同模式
model = torch.nn.Sequential(
    torch.nn.Linear(1024, 512),
    torch.nn.ReLU(),
    torch.nn.Linear(512, 256),
)

compiled_default = torch.compile(model, mode="default")
compiled_ro = torch.compile(model, mode="reduce-overhead")
compiled_ma = torch.compile(model, mode="max-autotune")

# 查看稳定后端
torch._dynamo.list_backends()
# 查看所有后端（含实验性后端）
torch._dynamo.list_backends(None)

def my_function(x):
    y = torch.sin(x)
    torch._dynamo.graph_break()  # 强制中断
    z = torch.cos(y)
    return z

if nopython and not config.debug_force_graph_break_on_leaf_return:
    return optimize_assert(
        backend,
        dynamic=dynamic,
        hooks=hooks,
        rebuild_ctx=rebuild_ctx,
        package=package,
    )

# 确保无 Graph Break，适用于性能关键路径
@torch.compile(fullgraph=True)
def critical_path(x, weight):
    return torch.nn.functional.linear(x, weight).relu()

# 如果函数内有 print 等不可追踪操作，fullgraph=True 会报错
@torch.compile(fullgraph=True)
def will_fail(x):
    print("debug")  # 这里会导致编译错误
    return x + 1

def explain(f, *extra_args, **extra_kwargs):
    def inner(*args, **kwargs):
        graphs = []
        break_reasons = []
        op_count = 0
        ops_per_graph = []

        def dynamo_graph_accumulating_compiler(gm, example_inputs):
            # 累积每个子图的信息
            ...
            return gm.forward

        opt_f = optimize(
            dynamo_graph_accumulating_compiler,
            nopython=False,
            guard_export_fn=guard_export_print,
        )(f)
        opt_f(*args, **kwargs)

        return ExplainOutput(
            graphs, graph_count, graph_break_count,
            break_reasons, op_count, ops_per_graph,
            out_guards, compile_time,
        )
    return inner

import torch

def my_model(x):
    y = torch.sin(x)
    print("intermediate:", y.shape)  # 会导致 Graph Break
    z = torch.cos(y)
    return z

explanation = torch._dynamo.explain(my_model)(torch.randn(8, 8))
print(explanation)
# 输出：
# Graph Count: 2
# Graph Break Count: 1
# Op Count: 2
# Break Reasons:
#   Break Reason 1:
#     Reason: builtin: print ...

# 作为上下文管理器
with torch._dynamo.config.patch(verbose=True):
    compiled_fn(x)

# 作为装饰器
@torch._dynamo.config.patch(recompile_limit=2)
def test_something():
    ...

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

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

model = SimpleModel().cuda()
x = torch.randn(32, 784, device="cuda")

# 编译模型
compiled_model = torch.compile(model)

# 首次调用触发编译
output = compiled_model(x)
print(f"Output shape: {output.shape}")  # torch.Size([32, 10])

import torch

def model_with_breaks(x):
    y = torch.sin(x) + torch.cos(x)
    # 数据依赖的控制流会导致 Graph Break
    if y.sum().item() > 0:
        z = y * 2
    else:
        z = y * 3
    return torch.relu(z)

x = torch.randn(4, 4)
explanation = torch._dynamo.explain(model_with_breaks)(x)
print(explanation)
# 查看 Graph Break 数量和原因
print(f"共 {explanation.graph_count} 个子图, "
      f"{explanation.graph_break_count} 次 Graph Break")

import torch
import torch.nn as nn
import time

class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, nhead=8):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 2048),
            nn.GELU(),
            nn.Linear(2048, d_model),
        )

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        x = self.norm2(x + self.ffn(x))
        return x

device = "cuda" if torch.cuda.is_available() else "cpu"
model = TransformerBlock().to(device)
x = torch.randn(16, 128, 512, device=device)

def benchmark(fn, x, warmup=10, repeat=100):
    """简单的计时基准函数"""
    # 预热
    for _ in range(warmup):
        fn(x)
    if device == "cuda":
        torch.cuda.synchronize()

    start = time.perf_counter()
    for _ in range(repeat):
        fn(x)
    if device == "cuda":
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    return elapsed / repeat * 1000  # 转换为毫秒

# Eager 基准
eager_time = benchmark(model, x)
print(f"Eager:              {eager_time:.3f} ms")

# torch.compile default 模式
compiled_default = torch.compile(model, mode="default")
default_time = benchmark(compiled_default, x)
print(f"Compiled (default): {default_time:.3f} ms")

# torch.compile reduce-overhead 模式
torch._dynamo.reset()
compiled_ro = torch.compile(model, mode="reduce-overhead")
ro_time = benchmark(compiled_ro, x)
print(f"Compiled (reduce):  {ro_time:.3f} ms")

# torch.compile max-autotune 模式
torch._dynamo.reset()
compiled_ma = torch.compile(model, mode="max-autotune")
ma_time = benchmark(compiled_ma, x)
print(f"Compiled (max-at):  {ma_time:.3f} ms")

print(f"\n加速比 (default):       {eager_time / default_time:.2f}x")
print(f"加速比 (reduce-overhead): {eager_time / ro_time:.2f}x")
print(f"加速比 (max-autotune):    {eager_time / ma_time:.2f}x")

import torch

# 一个可以完整捕获的函数
@torch.compile(fullgraph=True, backend="inductor")
def clean_function(x, weight):
    h = torch.nn.functional.linear(x, weight)
    h = torch.nn.functional.gelu(h)
    return h.sum()

x = torch.randn(32, 128, device="cuda")
w = torch.randn(64, 128, device="cuda")
result = clean_function(x, w)
print(f"fullgraph 编译成功, result = {result.item():.4f}")

# 验证 fullgraph 模式下 Graph Break 会报错
try:
    @torch.compile(fullgraph=True)
    def broken_function(x):
        print("this will cause an error")
        return x + 1

    broken_function(torch.randn(4))
except Exception as e:
    print(f"预期的错误: {type(e).__name__}")