TorchInductor Pattern Matcher

import torch


# Simple module for demonstration
class MyModule(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.param = torch.nn.Parameter(torch.rand(3, 4))
        self.linear = torch.nn.Linear(4, 5)

    def forward(self, x):
        return self.linear(x + self.param).clamp(min=0.0, max=1.0)


module = MyModule()

from torch.fx import symbolic_trace

# Symbolic tracing frontend - captures the semantics of the module
symbolic_traced: torch.fx.GraphModule = symbolic_trace(module)

# High-level intermediate representation (IR) - Graph representation
# 由一个列表组成 代表函数输入、调用点（函数、方法、 或 torch.nn.Module 实例），以及返回值。
print(symbolic_traced.graph)
"""
graph():
    %x : [num_users=1] = placeholder[target=x]
    %param : [num_users=1] = get_attr[target=param]
    %add : [num_users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {})
    %linear : [num_users=1] = call_module[target=linear](args = (%add,), kwargs = {})
    %clamp : [num_users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0})
    return clamp
"""

# Code generation - valid Python code
# 使 FX 成为 Python 到 Python（或 模块到模块）转换工具包。对于每个 Graph IR，我们可以 创建与图语义匹配的有效 Python 代码。
print(symbolic_traced.code)
"""
def forward(self, x):
    param = self.param
    add = x + param;  x = param = None
    linear = self.linear(add);  add = None
    clamp = linear.clamp(min = 0.0, max = 1.0);  linear = None
    return clamp
"""

if post_grad_custom_post_pass := config.post_grad_custom_post_pass:
    GraphTransformObserver(gm, "post_grad_custom_post_pass").apply_graph_pass(
        post_grad_custom_post_pass
    )

def register_replacement(
    pattern: Callable,
    replacement: Callable,
    example_inputs: List[Tensor],
    trace_fn: Callable = fwd_only,
    pm_pass: Optional[PatternMatcherPass] = None,
) -> PatternMatcherPass:
    # 实现细节...

class PatternMatcherPass:
    def apply(self, graph: torch.fx.GraphModule) -> int:
        ...
        for node in reversed(graph.nodes):
            target = extract_target(node)
            if (
                node.op in ["call_function", "call_method", "call_module"]
                and target in self.patterns
            ):
                ... # fallback_node_due_to_unsupported_type      
                for entry in self.patterns[target]:
                    if node._erased:
                        break
                    m = entry.pattern.match(node)
                    ... # pattern match crosses mutation barrier - discard
                    if is_match(m) and entry.extra_check(m):
                        count += 1
                        entry.apply(m, graph, node)  # type: ignore[arg-type]
                        ... # counters

# 原始操作
foo_inplace(x)    # 修改张量 x
bar_out(x, out)   # 使用修改后的 x，产生 out

# 对应的 FX 图（缺少依赖！）
%x = placeholder()
%out = placeholder()
%foo_result = call_function(foo_inplace, (%x,))
%bar_result = call_function(bar_out, (%x, %out))  # 缺少依赖！

# 使用 auto_functionalized 后的结果
def pattern(...):
    at = auto_functionalized(
        torch.ops.vllm.some_op.default,
        result=result,
        input=input,
        ...
    )
    return at[1], at[2]  # 返回原始返回值和修改后的参数

(original_return, mutated_arg1, mutated_arg2, ...)

class VllmPatternMatcherPass(VllmInductorPass):
    matched_count: int = 0
    match_table: ClassVar[defaultdict[str, int]] = defaultdict(int)
    
    def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None:
        # 输出模式调试信息

class VllmFusionPatternMatcherPass(VllmPatternMatcherPass):
    def __init__(self, config: VllmConfig, pass_name: str) -> None:
        super().__init__(config)
        self.pm_pass = PatternMatcherPass(pass_name=pass_name)
        
    def register(self, pr: VllmPatternReplacement) -> None:
        pm.register_replacement(
            pr.pattern,
            pr.replacement,
            pr.get_inputs(),
            self._trace_fn,
            self.pm_pass,
        )

def pattern(
    qkv: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	# 分割 QKV
	q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
	# 转换精度
	q_fp32 = q.to(torch.float32)
	k_fp32 = k.to(torch.float32)
	# 计算方差
	q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
	k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
	# 拼接方差
	qk_var = torch.cat([q_var, k_var], dim=-1)
	# AllReduce 操作
	qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
	# 分割方差
	q_var, k_var = qk_var.chunk(2, dim=-1)
	# 归一化和应用权重
	q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
	k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
	return q_out, k_out, v

def replacement(
    qkv: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    assert _MINIMAX_QK_NORM_FUSED_OP is not None
    # 调用融合内核
    q_out, k_out = torch.ops.vllm.minimax_qk_norm_fused(
        qkv,
        q_weight,
        k_weight,
        q_size,
        kv_size,
        tp_rank,
        tp_world,
        eps,
        max_tokens,
    )
    _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
    return q_out, k_out, v

pm.register_replacement(
    pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
)

class MiniMaxQKNormPass(VllmPatternMatcherPass):
	def __init__(self, config: VllmConfig) -> None:
		self.patterns: PatternMatcherPass = PatternMatcherPass(
            pass_name="minimax_qk_norm_pass"
        )
        self._register_patterns(q_size, kv_size, eps, tp_world, tp_rank)
        self.dump_patterns(config, self.patterns)
        self.disabled = False
	def _register_patterns(...):
	    MiniMaxQKNormPattern(
            q_size=q_size,
            kv_size=kv_size,
            eps=eps,
            tp_world=tp_world,
            tp_rank=tp_rank,
            max_tokens=self.max_token_num,
            dtype=self.model_dtype,
            device=self.device,
        ).register(self.patterns)
        
    def is_applicable_for_range(self, compile_range: Range) -> bool:
        if self.disabled:
            return False
		# 控制适用条件
        return bool(compile_range.end <= self.max_token_num)