# torch/fx/graph.py (L1287-L1306)
class Graph:
    def __init__(
        self,
        owning_module: Optional["GraphModule"] = None,
        tracer_cls: Optional[type["Tracer"]] = None,
        tracer_extras: Optional[dict[str, Any]] = None,
    ):
        self._root: Node = Node(self, "", "root", "", (), {})
        self._used_names: dict[str, int] = {}
        self._insert = self._root.prepend  # 默认插入点：_root 之前（即链表末尾）
        self._len = 0
        self._graph_namespace = _Namespace()
        self._owning_module = owning_module

# torch/fx/graph.py (L268-L284)
class _node_list:
    def __init__(self, graph: "Graph", direction: Literal["_prev", "_next"] = "_next"):
        self.graph = graph
        self.direction = direction

    def __len__(self):
        return self.graph._len

    def __iter__(self):
        return _NodeIter(self.graph._root, self.direction == "_prev")

    def __reversed__(self):
        return _node_list(self.graph, "_next" if self.direction == "_prev" else "_prev")

# target = 参数名（str），args = () 或 (default_value,)
# 对应生成代码中的函数参数
%x : [num_users=1] = placeholder[target=x]

# target = 属性全限定路径（str），args = (), kwargs = {}
%weight : [num_users=1] = get_attr[target=linear.weight]

# target = Callable（函数对象本身），args/kwargs = 函数参数
%add : [num_users=1] = call_function[target=torch.add](args = (%x, %y), kwargs = {})

# target = 方法名（str），args[0] = self 对象，其余为方法参数
%relu : [num_users=1] = call_method[target=relu](args = (%add,), kwargs = {})

# target = 模块全限定名（str），args/kwargs 不含 self
%linear : [num_users=1] = call_module[target=layers.0](args = (%x,), kwargs = {})

# target = "output", args[0] = 返回值（Node 或 Node 的 tuple）
return (%relu,)

# torch/fx/node.py (L339-L352)
if op == "call_function":
    if not callable(target):
        raise ValueError(...)
else:
    if op not in _legal_ops:
        raise AssertionError(f"op '{op}' is not in _legal_ops")
    if not isinstance(target, str):
        raise ValueError(...)

# torch/fx/node.py (L265-L300)
class Node(_NodeBase):
    _args: tuple["Argument", ...]        # 位置参数
    _kwargs: dict[str, "Argument"]       # 关键字参数
    graph: "Graph"                        # 所属图
    name: str                             # 唯一名称，如 "add", "add_1"
    op: str                               # opcode，6 种之一
    target: "Target"                      # 操作目标
    _input_nodes: dict["Node", None]      # 当前节点的所有输入节点（有序集合）
    users: dict["Node", None]             # 使用本节点输出的所有节点（有序集合）
    type: Optional[Any]                   # 输出类型注解
    _sort_key: Any                        # 拓扑排序键
    _repr_fn: Optional[Callable]          # 自定义 repr 函数
    meta: dict[str, Any]                  # 元数据字典

// torch/csrc/fx/node.cpp (L307-L350)
static PyObject* NodeBase__update_args_kwargs(
    PyObject* self, PyObject* const* args, Py_ssize_t nargs) {
  auto node = reinterpret_cast<NodeBase*>(self);
  auto input_nodes = node->_input_nodes;

  // Step 1: 清除旧的 _input_nodes，同时从旧输入节点的 users 中移除自身
  if (PyDict_GET_SIZE(input_nodes) > 0) {
    PyObject *key = nullptr, *value = nullptr;
    Py_ssize_t pos = 0;
    while (PyDict_Next(input_nodes, &pos, &key, &value)) {
      PyDict_DelItem(reinterpret_cast<NodeBase*>(key)->users, self);
    }
    PyDict_Clear(input_nodes);
  }

  // Step 2: 遍历新的 args/kwargs，收集所有 Node 类型的值
  //         将它们加入 _input_nodes，同时在对应节点的 users 中注册自身
  auto visit_fn = [self, input_nodes](PyObject* x) {
    if (is_node(x)) {
      PyDict_SetDefault(input_nodes, x, Py_None);    // self._input_nodes[x] = None
      PyDict_SetDefault(
          reinterpret_cast<NodeBase*>(x)->users, self, Py_None);  // x.users[self] = None
    }
    return Py_NewRef(x);
  };
  // ... map_aggregate(new_args, visit_fn), map_aggregate(new_kwargs, visit_fn)
}

# torch/fx/node.py (L476-L488)
@property
def all_input_nodes(self) -> list["Node"]:
    """
    Return all Nodes that are inputs to this Node.
    """
    return list(self._input_nodes.keys())

import torch
import torch.fx

graph = torch.fx.Graph()
x = graph.placeholder("x")
y = graph.placeholder("y")
add = graph.call_function(torch.add, (x, y))

# 自动建立的 Use-Def 关系
print(add.all_input_nodes)  # [x, y]
print(x.users)              # {add: None}
print(y.users)              # {add: None}

# 修改 args 后，关系自动更新
z = graph.placeholder("z")
add.args = (x, z)           # 用 z 替换 y
print(add.all_input_nodes)  # [x, z]
print(y.users)              # {} -- y 不再被使用
print(z.users)              # {add: None}

# torch/fx/graph.py (L1421-L1483)
def create_node(
    self,
    op: str,
    target: "Target",
    args: Optional[tuple["Argument", ...]] = None,
    kwargs: Optional[dict[str, "Argument"]] = None,
    name: Optional[str] = None,
    type_expr: Optional[Any] = None,
) -> Node:
    candidate = name if name is not None else self._target_to_str(target)
    name = self._graph_namespace.create_name(candidate, None)   # 确保名称唯一
    n = Node(self, name, op, target, args, kwargs, type_expr)   # 创建节点
    self._insert(n)                                              # 插入到当前插入点
    self._find_nodes_lookup_table.insert(n)                      # 加入快速查找表
    self._len += 1
    return n

# 各便捷方法的签名与 target 类型

graph.placeholder("x")                                  # op="placeholder", target="x"
graph.get_attr("layer.weight")                           # op="get_attr", target="layer.weight"
graph.call_function(torch.relu, (x,))                    # op="call_function", target=torch.relu
graph.call_method("reshape", (x, (2, 3)))                # op="call_method", target="reshape"
graph.call_module("linear", (x,))                        # op="call_module", target="linear"
graph.output(result)                                     # op="output", target="output"

import torch
import torch.fx
import operator

graph = torch.fx.Graph()

# Step 1: 创建输入
x = graph.placeholder("x")
y = graph.placeholder("y")

# Step 2: 构建计算
add = graph.call_function(operator.add, (x, y))
relu = graph.call_function(torch.relu, (add,))

# Step 3: 设置输出
graph.output(relu)

# 查看生成的代码
print(graph)
# def forward(self, x, y):
#     add = operator.add(x, y)
#     relu = torch.relu(add)
#     return relu

# torch/fx/graph.py (L1537-L1585)
def inserting_before(self, n: Optional[Node] = None):
    if n is None:
        return self.inserting_after(self._root)     # 插入到图的开头
    return _InsertPoint(self, n.prepend)

def inserting_after(self, n: Optional[Node] = None):
    if n is None:
        return self.inserting_before(self._root)    # 插入到图的末尾
    return _InsertPoint(self, n.append)

# torch/fx/graph.py (L256-L265)
class _InsertPoint:
    def __init__(self, graph, new_insert):
        self.graph = graph
        self.orig_insert, graph._insert = graph._insert, new_insert

    def __enter__(self):
        pass

    def __exit__(self, type, value, tb):
        self.graph._insert = self.orig_insert

# 在 add 节点之后插入新操作
with graph.inserting_after(add):
    neg = graph.call_function(torch.neg, (add,))

# 在 relu 节点之前插入新操作
with graph.inserting_before(relu):
    mul = graph.call_function(torch.mul, (neg, x))

# torch/fx/graph.py (L1497-L1535)
def erase_node(self, to_erase: Node) -> None:
    if len(to_erase.users) > 0:
        raise RuntimeError(
            f"Tried to erase Node {to_erase} but it still had "
            f"{len(to_erase.users)} users in the graph: {to_erase.users}!"
        )

    self._find_nodes_lookup_table.remove(to_erase)
    to_erase._remove_from_list()               # 从链表中摘除
    to_erase._erased = True
    self._len -= 1

    # 清空 args/kwargs，自动触发 Use-Def 更新
    to_erase._update_args_kwargs(
        map_arg(to_erase._args, lambda n: None),
        map_arg(to_erase._kwargs, lambda n: None),
    )

# torch/fx/node.py (L691-L735)
def replace_all_uses_with(
    self,
    replace_with: "Node",
    delete_user_cb: Optional[Callable[["Node"], bool]] = None,
    *,
    propagate_meta: bool = False,
) -> list["Node"]:
    to_process = [*self.users]
    result = []
    for use_node in to_process:
        if delete_user_cb is not None and not delete_user_cb(use_node):
            continue
        result.append(use_node)
        use_node._replace_input_with(self, replace_with)
    return result

# torch/fx/graph.py (L1820-L1854)
def node_copy(
    self, node: Node, arg_transform: Callable[[Node], "Argument"] = lambda x: x
) -> Node:
    args = map_arg(node.args, arg_transform)
    kwargs = map_arg(node.kwargs, arg_transform)
    result_node = self.create_node(node.op, node.target, args, kwargs, node.name, node.type)
    result_node.meta = copy.copy(node.meta)
    return result_node

# torch/fx/graph.py (L2041-L2118)
def lint(self):
    seen_names: set[str] = set()
    seen_values: set[Node] = set()
    for node in self.nodes:
        # 检查 1: opcode 合法性
        if node.op not in _legal_ops:
            raise RuntimeError(f"Node {node} had unknown opcode {node.op}!")

        # 检查 2: 节点归属
        if node.graph is not self:
            raise RuntimeError(f"Node '{node}' does not belong to this Graph!")

        # 检查 3: 拓扑排序 -- 所有输入节点必须在当前节点之前出现
        for arg in node._input_nodes:
            if arg not in seen_values:
                raise RuntimeError(
                    f"Argument '{arg}' was used before it has been defined!"
                )
        seen_values.add(node)

        # 检查 4: 名称唯一性
        if node.name in seen_names:
            raise RuntimeError(f"Node redefined name {node.name}!")
        seen_names.add(node.name)

    # 检查 5: 如果有 owning_module，验证 target 引用的合法性
    if self.owning_module:
        for node in self.nodes:
            if node.op in ["get_attr", "call_module"]:
                # 逐级验证 target 路径中每个属性都存在
                ...

import torch
import torch.fx
import torch.nn.functional as F

# === Step 1: 通过 symbolic_trace 获取图 ===
class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 20)

    def forward(self, x):
        x = self.linear(x)
        x = torch.relu(x)
        return x

traced = torch.fx.symbolic_trace(MyModule())
graph = traced.graph

# === Step 2: 遍历节点 ===
for node in graph.nodes:
    print(f"op={node.op:16s} name={node.name:12s} target={node.target}")
# op=placeholder       name=x            target=x
# op=call_module       name=linear        target=linear
# op=call_function     name=relu          target=<built-in method relu ...>
# op=output            name=output        target=output

# === Step 3: 将 relu 替换为 gelu ===
for node in graph.nodes:
    if node.op == "call_function" and node.target == torch.relu:
        with graph.inserting_after(node):
            new_node = graph.call_function(F.gelu, node.args, node.kwargs)
        node.replace_all_uses_with(new_node)
        graph.erase_node(node)

# === Step 4: 验证与查看结果 ===
graph.lint()
traced.recompile()
print(traced.code)
# def forward(self, x):
#     linear = self.linear(x)
#     gelu = torch.nn.functional.gelu(linear)
#     return gelu

import torch
import torch.fx

def analyze_graph(graph: torch.fx.Graph):
    """打印图中每个节点的 Use-Def 关系"""
    for node in graph.nodes:
        inputs = [n.name for n in node.all_input_nodes]
        users = [n.name for n in node.users.keys()]
        print(f"{node.name:12s} | inputs: {inputs:30s} | users: {users}")

class Net(torch.nn.Module):
    def forward(self, x, y):
        a = x + y
        b = a * x
        c = torch.relu(b)
        return c

traced = torch.fx.symbolic_trace(Net())
analyze_graph(traced.graph)
# x            | inputs: []                             | users: ['add', 'mul']
# y            | inputs: []                             | users: ['add']
# add          | inputs: ['x', 'y']                     | users: ['mul']
# mul          | inputs: ['add', 'x']                   | users: ['relu']
# relu         | inputs: ['mul']                        | users: ['output']
# output       | inputs: ['relu']                       | users: []