GMCP调度算法优化矩阵乘法内存管理

1. 项目概述

在计算密集型应用中，矩阵乘法是最基础也是最耗时的操作之一。特别是在神经网络推理、图像处理等场景中，如何高效地调度矩阵乘法计算图的内存使用，直接关系到整体性能表现。GMCP（Graph Memory Constraint Priority）调度算法正是为解决这一问题而设计的一种创新性解决方案。

这个调度器的核心思想是通过优先级队列管理计算图中的节点执行顺序，在满足依赖关系的前提下，尽可能减少内存峰值使用量。这对于资源受限的嵌入式设备、AI加速芯片等场景尤为重要，因为内存带宽和容量往往是性能瓶颈所在。

2. 核心数据结构设计

2.1 NodeInfo类：节点信息封装

NodeInfo类使用Python的__slots__特性优化内存占用，避免了普通Python类的动态字典开销。这种设计在需要处理大量节点的场景下能显著减少内存消耗。

python复制class NodeInfo:
    __slots__ = ["op", "buf_id", "size", "bufs", "type", "succs"]
    
    def __init__(self, op: str, buf_id: Optional[str], size: int, 
                 bufs: List[str], type_: str, succs: List[str]):
        self.op = op  # 节点操作类型
        self.buf_id = buf_id  # 关联的缓冲区ID
        self.size = size  # 缓冲区大小(字节)
        self.bufs = bufs  # 操作涉及的缓冲区列表
        self.type = type_  # 缓冲区类型
        self.succs = succs  # 后继节点ID列表

每个节点的操作类型(op)包括：

• ALLOC：内存分配操作
• FREE：内存释放操作
• Matmul：矩阵乘法计算
• COPY_IN：数据拷贝输入

2.2 BufInfo类：缓冲区生命周期管理

BufInfo类完整记录了缓冲区的整个生命周期，从分配到释放，以及期间的所有使用情况：

python复制class BufInfo:
    __slots__ = ["a_node", "f_node", "size", "type", "o_nodes"]
    
    def __init__(self):
        self.a_node: Optional[str] = None  # 分配节点ID
        self.f_node: Optional[str] = None  # 释放节点ID
        self.size: int = 0  # 缓冲区大小
        self.type: str = ""  # 缓存类型
        self.o_nodes: List[str] = []  # 使用该缓冲区的操作节点

这种设计使得我们可以精确跟踪每个缓冲区的使用情况，为后续的调度决策提供依据。

3. 图解析模块实现

3.1 图解析流程

parse_matmul_graph函数负责解析描述计算图的JSON文件，构建完整的内存数据结构：

python复制def parse_matmul_graph(json_path: str) -> Tuple[Dict[str, NodeInfo], Dict[str, BufInfo], Dict[str, int]]:
    node_dict = {}  # 节点ID到NodeInfo的映射
    buf_dict = {}   # 缓冲区ID到BufInfo的映射
    in_degree = {}  # 节点入度表
    adj = {}        # 邻接表
    
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    
    # 初始化节点信息
    for node in data["Nodes"]:
        node_id = str(node["Id"])
        # 解析节点各字段...
        
        # 初始化入度和邻接表
        in_degree[node_id] = 0
        adj[node_id] = []
        
        # 处理不同类型的节点
        if op == "ALLOC":
            # 记录分配节点信息
            pass
        elif op == "FREE":
            # 记录释放节点信息
            pass