torch.compile 里 Inductor 是怎么把编译结果包进 CUDA Graph

from torch._dynamo import register_backend

@register_backend   # 把 "inductor" 注册到 _COMPILER_FNS 字典
def inductor(*args, **kwargs):
    from torch._inductor.compile_fx import compile_fx  # 懒加载
    return compile_fx(*args, **kwargs)

_COMPILER_FNS: dict[str, CompilerFn] = {}
_default_backend: str | CompilerFn = "inductor"  # 默认后端就是 inductor

def lookup_backend(compiler_fn: str | CompilerFn) -> CompilerFn:
    if isinstance(compiler_fn, str):
        compiler_fn = _COMPILER_FNS[compiler_fn]  # "inductor" → 上面注册的函数
    return compiler_fn

graph = GraphLowering(gm, example_inputs=example_inputs, ...)

with V.set_graph_handler(graph):
    graph.run(*example_inputs)     # FX node → Inductor IR (TensorBox, Pointwise, Reduction...)

with dynamo_timed("GraphLowering.compile_to_fn"):
    compiled_fn = graph.compile_to_fn()  # Inductor IR → Triton kernel → 编译

# config.py，triton 子配置
cudagraphs = os.environ.get("TORCHINDUCTOR_CUDAGRAPHS") == "1"  # 默认关闭
cudagraph_trees = True   # 树形内存池共享，默认开启

# 开启方式
torch.compile(model, options={"triton.cudagraphs": True})
# 或
TORCHINDUCTOR_CUDAGRAPHS=1 python train.py

# compile_fx.py ~900
if graph_kwargs.get("cudagraphs") is None:
    graph_kwargs["cudagraphs"] = BoxedBool(config.triton.cudagraphs)

if cudagraphs:
    if self.disabled_cudagraphs_reason:
        BoxedBool.disable(cudagraphs)   # 有 skip 原因，放弃
    else:
        if config.graph_partition and policy is None:
            cudagraph_partition_post_compile(...)  # 按 partition 分别包装
        else:
            cudagraph_post_compile(...)            # 整图包装

compiled_graph.current_callable = cudagraphify(
    current_callable,
    static_input_idxs=...,
    device_index=...,
    is_backward=...,
    is_inference=...,
    ...
)

def cudagraphify(model, static_input_idxs=(), *, device_index, ...):
    if config.triton.cudagraph_trees:
        cudagraphify_fn = functools.partial(new_cudagraphify_impl, ...)  # 树形版
    else:
        cudagraphify_fn = cudagraphify_impl   # 简单版

    compiled_fn = None

    def run(new_inputs):   # 懒惰录制：第一次真正调用时才录制
        nonlocal compiled_fn
        if compiled_fn is None:
            compiled_fn = cudagraphify_fn(model, new_inputs, static_input_idxs)
        return compiled_fn(new_inputs)

    return run

def cudagraphify_impl(model, inputs, static_input_idxs=()):
    # 1. 分配 static_inputs
    static_inputs = [static_input(x) if idx not in static_input_idxs else x.detach()
                     for idx, x in enumerate(inputs)]

    # 2. warmup（side stream）
    torch.cuda.synchronize()
    stream = torch.cuda.Stream()
    with torch.cuda.stream(stream):
        model(list(static_inputs))
    torch.cuda.synchronize()

    # 3. 录制
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"):
        static_outputs = model(list(static_inputs))

    # 4. replay 闭包
    def run(new_inputs):
        for idx in copy_indices:
            index_expanded_dims_and_copy_(static_inputs[idx], new_inputs[idx], ...)
        new_inputs.clear()
        graph.replay()
        return static_outputs

    return run

def cudagraphify_impl(model, inputs, static_input_idxs, *args, **kwargs):
    fn_cache: dict[tuple[int,...], Callable] = {}

    def deferred_cudagraphify(inputs):
        int_key = get_ints(inputs)
        fn = fn_cache.get(int_key)
        if fn is not None:
            return fn(inputs)
        # 第一次：录制，结果缓存
        fn, out = cudagraphify(model, inputs, ...)  # → CUDAGraphTreeManager.add_function()
        fn_cache[int_key] = fn
        return out

    return deferred_cudagraphify

def __init__(self, device_index):
    with torch.cuda.device(device_index):
        torch.cuda.synchronize()
        self.stream = torch.cuda.Stream()
        self.graph = torch.cuda.CUDAGraph()
        self.cuda_graphs_thread_pool = torch.cuda.graph_pool_handle()  # 共享内存池

        # 用空图激活 pool，保持其存活
        with torch.cuda.graph(self.graph, pool=self.cuda_graphs_thread_pool, ...):
            pass

def _record(self, model, inputs):
    with (
        preserve_rng_state(),
        torch.cuda.device(self.device),
        clear_cublas_manager(),
        torch.cuda.graph(
            self.graph,
            stream=self.stream,
            pool=self.cuda_graphs_pool,           # 共享内存池
            capture_error_mode="thread_local",
        ),
        CUDAGraphCaptureControlFlowOpDispatchMode(),
        get_history_recording(),
    ):
        static_outputs = model(inputs)
    return static_outputs

def run(self, new_inputs):
    self.check_static_inputs_are_stable(new_inputs)   # 检查 static ptr 没变
    self._copy_inputs_and_remove_from_src(            # 新 input 拷到 static slot
        self.reconstructed_inputs, new_inputs
    )
    self.run_graph()                                   # graph.replay()
    return self.reconstruct_outputs()                  # 从 metadata 重建 tensor

if is_backward and config.triton.cudagraph_trees:
    def compiled_artifact(new_inputs):
        manager.set_to_running_backward()  # 通知 manager 进入 backward 模式
        return compiled_graph_callable(new_inputs)
    compiled_graph.current_callable = compiled_artifact

# 主开关（默认关）
torch.compile(model, options={"triton.cudagraphs": True})

# 树形 pool 共享（默认开）
torch.compile(model, options={"triton.cudagraph_trees": True})

# 动态 shape 下跳过（默认不跳过，会为每个 shape 录制一个图）
torch.compile(model, options={"triton.cudagraph_skip_dynamic_graphs": True})

# 最多允许多少次 re-record（默认 128）
torch.compile(model, options={"triton.cudagraph_unexpected_rerecord_limit": 128})

# 强制要求 cudagraph，不兼容时报错而非 skip
torch.compile(model, options={"triton.cudagraph_or_error": True})