pkgs/MagiCompiler/magi_compiler/tokenflow/graph_profile.py

# Copyright (c) 2026 SandAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
import re
import weakref
from typing import Any, Dict, Tuple

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
from magi_compiler.config import get_compile_config
from magi_compiler.tokenflow.graph_executor import FX_NODE_OP  # 导入你的枚举类
from magi_compiler.tokenflow.green_ctx import GreenCtxManager, GreenStreamPool
from magi_compiler.tokenflow.sampler import exponential_aligned_sampler
from magi_compiler.utils import magi_logger
from torch._subclasses.fake_tensor import FakeTensor
from torch.fx import GraphModule, Node

MIN_LATENCY_MS = 0.005  # 最小延迟下限（5微秒）


class GraphProfileWrapper:
    def __init__(self, graph_module: GraphModule):
        # 核心属性初始化
        self.graph_module = graph_module
        self.noncompute_node_names = set()
        self.profile_results: Dict[int, Dict[str, Dict[int, float]]] = {}

        # 初始化SM测试列表
        max_sm = GreenCtxManager.get_max_sm()
        min_sm, align = GreenCtxManager.get_min_and_align_sm()
        high_sm_counts = exponential_aligned_sampler(min_sm, max_sm, num_samples=5, align=align)
        low_sm_counts = [max_sm - sm for sm in high_sm_counts]
        self.sm_samples = sorted(list(set(high_sm_counts).union(set(low_sm_counts)).union(set([max_sm]))))
        self.sm_samples.remove(0) if 0 in self.sm_samples else None

        # 初始化SeqLen测试列表
        self.seqlen_samples = exponential_aligned_sampler(2, 4096, num_samples=10, align=1)
        self.seqlen_samples.reverse()  # 从大到小测试

        # 结果保存路径与测试配置
        rank_idx = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
        self.save_base_path = os.path.join(get_compile_config().cache_root_dir, "graph_profile", f"rank_{rank_idx}")
        self.profiled = False

        # 初始化流池
        self.green_stream_pool = GreenStreamPool()
        for sm in self.sm_samples:
            self.green_stream_pool.get_stream(sm)

        self.node_param_ref_dict = {}

    def _resolve_symint_expression(self, sym_expr: Any, seq_len: int) -> int:
        # 普通整数直接返回
        if isinstance(sym_expr, int):
            return sym_expr

        # 解析SymInt表达式
        if isinstance(sym_expr, torch.SymInt):
            expr_str = str(sym_expr)
            symbols = re.findall(r's\d+', expr_str)

            if not symbols:
                return int(sym_expr)

            eval_env = {sym: seq_len for sym in symbols}
            res_str = expr_str
            for sym, val in eval_env.items():
                res_str = res_str.replace(sym, str(val))

            # print(f"解析 SymInt: {expr_str} -> {res_str}")
            return int(eval(res_str))

        raise TypeError(f"不支持的表达式类型: {type(sym_expr)}")

    def _generate_real_tensor(
        self, shape: Tuple[int, ...], stride: Tuple[int, ...], dtype: torch.dtype, device: torch.device, seq_len: int
    ) -> torch.Tensor:
        # 解析shape中的SymInt表达式
        resolved_shape = []
        for dim in shape:
            resolved_dim = self._resolve_symint_expression(dim, seq_len)
            resolved_shape.append(resolved_dim)
        resolved_shape = tuple(resolved_shape)

        # 解析stride中的SymInt表达式
        resolved_stride = []
        for s in stride:
            resolved_s = self._resolve_symint_expression(s, seq_len)
            resolved_stride.append(resolved_s)
        resolved_stride = tuple(resolved_stride)

        # 验证shape与stride长度匹配
        assert len(resolved_stride) == len(
            resolved_shape
        ), f"stride长度({len(resolved_stride)})与shape长度({len(resolved_shape)})不匹配"

        # 创建指定布局的空Tensor并填充随机数据
        tensor = torch.empty_strided(
            size=resolved_shape, stride=resolved_stride, dtype=dtype, device=device, requires_grad=False
        )
        tensor.normal_(mean=0.0, std=1.0)

        # 验证布局（调试用，保留断言）
        assert tensor.shape == resolved_shape, f"生成Tensor形状({tensor.shape})与预期({resolved_shape})不匹配"
        assert tensor.stride() == resolved_stride, f"生成Tensor stride({tensor.stride()})与预期({resolved_stride})不匹配"

        return tensor

    def _prepare_node_inputs(self, node: Node, seq_len: int) -> Tuple[Tuple, Dict]:
        # 处理位置参数与关键字参数
        args = [self._process_arg(arg, seq_len) for arg in node.args]
        kwargs = {key: self._process_arg(value, seq_len) for key, value in node.kwargs.items()}
        return tuple(args), kwargs

    def _process_arg(self, arg: Any, seq_len: int) -> Any:
        """递归处理参数，替换节点引用为真实随机Tensor"""
        # 处理fx.Node类型参数
        if isinstance(arg, Node):
            input_node = arg
            example_value = arg.meta.get("example_value")
            assert example_value is not None, f"无法获取节点 {input_node.name} 的元信息"

            # 处理标量与SymInt类型
            if isinstance(example_value, (int, float, str, bool)):
                return example_value
            if isinstance(example_value, torch.SymInt):
                return self._resolve_symint_expression(example_value, seq_len)

            if isinstance(example_value, torch.nn.Parameter):
                assert input_node.name in self.node_param_ref_dict, f"无法找到节点 {input_node.name} 对应的参数"
                res = self.node_param_ref_dict[input_node.name]()
                assert res is not None, f"节点 {input_node.name} 对应的参数已被释放"
                return res

            # 处理Tensor/FakeTensor类型
            if isinstance(example_value, (torch.Tensor, FakeTensor)):
                shape = example_value.shape
                dtype = example_value.dtype
                device = example_value.device
                stride = example_value.stride()

            # 生成真实Tensor并返回
            assert shape is not None, f"无法获取节点 {input_node.name} 的输出形状"
            return self._generate_real_tensor(shape, stride, dtype, device, seq_len)

        # 嵌套结构递归处理
        elif isinstance(arg, (list, tuple)):
            return type(arg)(self._process_arg(item, seq_len) for item in arg)

        # 其他类型直接返回
        else:
            return arg

    def run_batch_profile(self):
        """批量执行节点性能分析"""
        # 打印批量分析核心配置
        magi_logger.info(f"===== 开始批量性能分析 =====", rank=0)
        magi_logger.info(
            f"测试SeqLen: {self.seqlen_samples} | SM粒度: {self.sm_samples} | 节点总数: {len(list(self.graph_module.graph.nodes))}",
            rank=0,
        )

        # 初始化seq_len结果存储
        for seq_len in self.seqlen_samples:
            if seq_len not in self.profile_results:
                self.profile_results[seq_len] = {}

        # 遍历所有节点进行测试
        for node_idx, node in enumerate(self.graph_module.graph.nodes):
            node_name = node.name
            node_op = node.op
            node_perf = None

            skip_record = False
            if node_op in ["placeholder", "output"]:
                skip_record = True
            elif node_op == "call_function" and any(k in str(node.target) for k in ["getitem", "list"]):
                skip_record = True
            if skip_record:
                for seq_len in self.seqlen_samples:
                    self.profile_results[seq_len][node_name] = {sm: MIN_LATENCY_MS for sm in self.sm_samples}  # 5ns
                magi_logger.info(f"📌 跳过非计算节点: {node_name} ({node_op})", rank=0)
                self.noncompute_node_names.add(node_name)
                continue

            # 匹配重复节点名称（b\d+s\d+_xxx）
            elif re.match(r"b\d+s\d+_.*", node_name) and node_perf is None:
                base_name = re.sub(r"^b\d+s\d+_", "", node_name)
                for seq_len in self.seqlen_samples:
                    if base_name in self.profile_results[seq_len]:
                        node_perf = self.profile_results[seq_len][base_name]
                        self.profile_results[seq_len][node_name] = node_perf
                        skip_record = True
                if skip_record:
                    magi_logger.info(f"📌 重用节点性能: {node_name} ({node_op}) -> 基准节点: {base_name}", rank=0)
                    continue

            # 遍历所有SeqLen进行节点测试
            for seq_len in self.seqlen_samples:
                magi_logger.info(f"--- Profile 节点 {node_name} (SeqLen={seq_len}) ---", rank=0)
                try:
                    # 准备输入并执行性能分析
                    real_args, real_kwargs = self._prepare_node_inputs(node, seq_len)
                    node_perf = self._profile_node_with_data(node=node, args=real_args, kwargs=real_kwargs)

                    # 释放内存并保存结果
                    del real_args, real_kwargs
                    assert isinstance(node_perf, dict), f"节点性能结果格式错误"
                    self.profile_results[seq_len][node_name] = node_perf

                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()

                except Exception as e:
                    magi_logger.info(f"⚠️  节点 {node_name} (SeqLen={seq_len}) 测试失败: {str(e)=}", rank=0)
                    import traceback

                    traceback.print_exc()
                    self.profile_results[seq_len][node_name] = {sm: -1.0 for sm in self.sm_samples}
                    torch.cuda.empty_cache()
                    raise e

        self._incremental_save_results()

        # 打印汇总结果
        magi_logger.info(f"===== 批量性能分析完成 =====", rank=0)

        # 打印各SeqLen节点性能汇总
        for seq_len in self.seqlen_samples:
            magi_logger.info(f"===== SeqLen={seq_len} 性能汇总（单位：ms） =====", rank=0)
            header = f"{'节点名称':<20}" + "".join([f"SM={sm:<10}" for sm in self.sm_samples])
            magi_logger.info(header, rank=0)
            magi_logger.info("-" * len(header), rank=0)

            for node in self.graph_module.graph.nodes:
                if node.name not in self.profile_results[seq_len]:
                    continue
                if node.name in self.noncompute_node_names:
                    continue
                perf_dict = self.profile_results[seq_len].get(node.name, {})
                assert isinstance(perf_dict, dict), f"节点性能结果格式错误"
                row = f"{node.name:<20}" + "".join([f"{perf_dict.get(sm, -1.0):<10.4f}" for sm in self.sm_samples])
                magi_logger.info(row, rank=0)

        self.profiled = True

    def _profile_node_with_data(
        self, node: Node, args: Tuple, kwargs: Dict, warmup_steps: int = 30, test_steps: int = 30
    ) -> Dict[int, float]:
        """单节点性能测试（指定输入数据）"""
        node_perf = {}
        node_name = node.name
        node_op = node.op

        # 遍历所有SM数量进行测试
        for sm in self.sm_samples:
            try:
                stream = self.green_stream_pool.get_stream(sm)

                # 节点执行函数定义
                def execute_node_once():
                    with torch.cuda.stream(stream):
                        if node_op == FX_NODE_OP.CALL_FUNCTION.value:
                            return node.target(*args, **kwargs)
                        elif node_op == FX_NODE_OP.CALL_METHOD.value:
                            obj = args[0]
                            return getattr(obj, node.target)(*args[1:], **kwargs)
                        elif node_op == FX_NODE_OP.CALL_MODULE.value:
                            submod = self.graph_module
                            for mod_name in node.target.split("."):
                                submod = getattr(submod, mod_name)
                            return submod(*args, **kwargs)
                        else:
                            raise NotImplementedError(f"不支持的节点类型: {node_op}")

                # 预热与计时
                torch.cuda.synchronize()
                for _ in range(warmup_steps):
                    execute_node_once()
                torch.cuda.synchronize()

                start_event = torch.cuda.Event(enable_timing=True)
                end_event = torch.cuda.Event(enable_timing=True)

                start_event.record(stream)
                for _ in range(test_steps):
                    execute_node_once()
                end_event.record(stream)
                torch.cuda.synchronize()

                # 计算平均延迟
                total_time = start_event.elapsed_time(end_event)
                node_perf[sm] = total_time / test_steps

            except Exception as e:
                magi_logger.info(f"⚠️  节点 {node_name} (SM={sm}) 测试失败: {str(e)=}", rank=0)
                node_perf[sm] = -1.0
                torch.cuda.empty_cache()
                torch.cuda.synchronize()

        return node_perf

    def _incremental_save_results(self):
        """增量保存性能结果并生成热力图"""
        # 创建保存目录
        os.makedirs(self.save_base_path, exist_ok=True)

        # 增量保存JSON结果
        json_path = os.path.join(self.save_base_path, "profile_results.json")
        existing_results = {}
        if os.path.exists(json_path):
            with open(json_path, "r", encoding="utf-8") as f:
                existing_results = json.load(f)
            existing_results = {int(k): v for k, v in existing_results.items()}

        # 格式化新数据并合并
        new_data = {
            int(seq_len): {
                node_name: {int(sm): float(latency) for sm, latency in sm_data.items()}
                for node_name, sm_data in node_data.items()
                if sm_data
            }
            for seq_len, node_data in self.profile_results.items()
        }
        existing_results.update(new_data)

        # 写入JSON文件
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(existing_results, f, indent=4, ensure_ascii=False)
        magi_logger.info(f"📁 性能结果已保存至: {json_path}", rank=0)

        # 生成热力图（有有效数据时）
        if len(existing_results) > 0 and all(len(v) > 0 for v in existing_results.values()):
            self._generate_heatmaps(existing_results)

    def _generate_heatmaps(self, results: Dict[int, Dict[str, Dict[int, float]]]):
        """生成节点性能热力图"""
        # 配置绘图参数
        plt.rcParams.update(
            {
                'font.sans-serif': ['DejaVu Sans', 'Arial', 'Helvetica'],
                'axes.unicode_minus': False,
                'font.family': 'sans-serif',
                'figure.dpi': 300,
                'savefig.dpi': 300,
            }
        )

        magi_logger.info(f"📊 开始生成热力图（有效SeqLen: {list(results.keys())}）", rank=0)

        # 提取所有唯一数据
        all_seq_lens = sorted([int(k) for k in results.keys()])
        all_sms = sorted(
            {int(sm) for seq_data in results.values() for node_data in seq_data.values() for sm in node_data.keys()}
        )
        all_nodes = sorted({node for seq_data in results.values() for node in seq_data.keys()})

        # 为每个节点生成热力图
        for node_name in all_nodes:
            if node_name in self.noncompute_node_names:
                continue
            # 构建数据矩阵
            data_matrix = np.array(
                [[results.get(seq_len, {}).get(node_name, {}).get(sm, -1.0) for sm in all_sms] for seq_len in all_seq_lens]
            )

            # 绘制并保存热力图
            plt.figure(figsize=(12, 8))
            sns.heatmap(
                data_matrix,
                annot=True,
                fmt=".4f",
                cmap="RdYlBu_r",
                xticklabels=all_sms,
                yticklabels=all_seq_lens,
                cbar_kws={"label": "Latency (ms)"},
                mask=(data_matrix < 0),
                annot_kws={"size": 8},
            )

            plt.title(f"Node {node_name} Latency Heatmap", fontsize=14, pad=20)
            plt.xlabel("SM Count", fontsize=12)
            plt.ylabel("Sequence Length", fontsize=12)
            plt.tight_layout()

            img_path = os.path.join(self.save_base_path, f"{node_name}_latency_heatmap.png")
            plt.savefig(img_path, dpi=300, bbox_inches="tight", facecolor='white', edgecolor='none')
            plt.close()

        magi_logger.info(f"📊 热力图生成完成", rank=0)

    def record_node_param_refs(self, real_args: Tuple):
        """记录节点参数引用，供后续输入数据生成使用"""
        for node in self.graph_module.graph.nodes:
            if node.op == "placeholder":
                input_idx = list(self.graph_module.graph.nodes).index(node)
                if input_idx < len(real_args):
                    arg_value = real_args[input_idx]
                    if isinstance(arg_value, torch.nn.Parameter):
                        self.node_param_ref_dict[node.name] = weakref.ref(arg_value)

    def __call__(self, *args, **kwargs):
        """类可调用入口，未分析则先执行批量分析"""
        if not self.profiled:
            self.record_node_param_refs(args)
            self.run_batch_profile()
        return self.graph_module(*args, **kwargs)


def gen_profile_wrap_func(graph_module: GraphModule) -> GraphProfileWrapper:
    """生成性能分析封装实例"""
    return GraphProfileWrapper(graph_module)