cann组织链接：https://atomgit.com/cann
pypto仓库链接：https://atomgit.com/cann/pypto

前言

在 AI 编译器与高性能算子开发领域，如何在保持编程灵活性的同时实现底层硬件的极致优化，始终是核心挑战。CANN（Compute Architecture for Neural Networks）开源生态推出的 PyPTO（Parallel Tensor/Tile Operation）编程范式，通过融合 Python 的易用性与 Tile 级控制能力，为开发者提供了一条从算法描述到高效执行的完整路径。其核心在于一套自定义的中间表示（IR），该 IR 不仅承载了 Tile 操作语义，还支持复杂的依赖分析、内存规划与硬件映射。

本文基于 pypto 仓库（https://atomgit.com/cann/pypto），深入剖析其 IR 构建流程 与 图表示机制，涵盖 AST 解析、IR 节点设计、依赖图构建、内存生命周期建模等关键技术，并通过 pypto/ir/、pypto/frontend/ 与 pypto/backend/ 中的核心代码片段，揭示其工程实现细节。

1. PyPTO 编程模型概览

PyPTO 允许开发者以 Python 风格编写 Tile 级 Kernel：

# examples/gemm/gemm.py
def gemm_kernel(A: Tile, B: Tile, C: Tile):
    with parallel(M=128, N=128, K=64):
        C[M, N] += A[M, K] * B[K, N]

该代码并非直接执行，而是被 PyPTO 前端 解析为结构化 IR，再经由后端编译为 pto-isa 指令序列或 C++ Kernel。

2. IR 节点设计：面向 Tile 操作的抽象

PyPTO 定义了一套层次化的 IR 节点体系，位于 pypto/ir/nodes.py。

2.1 核心节点类型

# pypto/ir/nodes.py
class IRNode:
    pass

class Tensor(IRNode):
    def __init__(self, name: str, shape: Tuple[int], dtype: str):
        self.name = name
        self.shape = shape
        self.dtype = dtype

class TileOp(IRNode):
    """ 表示一个 Tile 级操作 """
    def __init__(self, op_type: str, inputs: List[IRNode], outputs: List[IRNode]):
        self.op_type = op_type  # e.g., "matmul", "load", "store"
        self.inputs = inputs
        self.outputs = outputs
        self.attrs = {}  # 存储 Tile Shape、Layout 等属性

class ParallelScope(IRNode):
    """ 表示 parallel(...) 上下文 """
    def __init__(self, tile_shape: Dict[str, int]):
        self.tile_shape = tile_shape
        self.body: List[IRNode] = []

✅ 设计原则：每个 TileOp 对应一条 pto-isa 指令，确保 IR 与目标 ISA 语义对齐。

2.2 内存操作显式建模

PyPTO 强制区分 Global 与 Local 内存访问：

# 在 IR 构建时自动插入 Load/Store
class LoadOp(TileOp):
    def __init__(self, src: Tensor, dst_tile: Tile, mem_space="global"):
        super().__init__("load", [src], [dst_tile])
        self.attrs["mem_space"] = mem_space

class StoreOp(TileOp):
    def __init__(self, src_tile: Tile, dst: Tensor, mem_space="global"):
        super().__init__("store", [src_tile], [dst])
        self.attrs["mem_space"] = mem_space

这使得后续 Pass 可精确分析数据流与内存占用。

---

3. IR 构建：从 Python AST 到结构化图

3.1 AST 访问器解析

前端使用 ast.NodeVisitor 遍历 Python AST：

# pypto/frontend/parser.py
class PyPTOParser(ast.NodeVisitor):
    def visit_FunctionDef(self, node: ast.FunctionDef):
        # 创建函数作用域
        func_ir = FunctionIR(node.name)
        
        # 解析参数（应为 Tile 类型）
        for arg in node.args.args:
            tensor = Tensor(arg.arg, ..., ...)
            func_ir.add_input(tensor)
        
        # 递归解析函数体
        for stmt in node.body:
            ir_node = self.visit(stmt)
            func_ir.body.append(ir_node)
        
        return func_ir

    def visit_With(self, node: ast.With):
        # 处理 with parallel(...) 语句
        if node.items[0].context_expr.func.id == "parallel":
            tile_shape = self._parse_parallel_args(node.items[0].context_expr)
            scope = ParallelScope(tile_shape)
            
            for stmt in node.body:
                op = self.visit(stmt)  # 如 Assign 节点
                scope.body.append(op)
            return scope

3.2 表达式到 TileOp 的转换

对于赋值语句 C[M, N] += A[M, K] * B[K, N]，解析为 MatMul：

def visit_AugAssign(self, node: ast.AugAssign):
    if isinstance(node.op, ast.Add):
        # 提取索引变量
        lhs_indices = self._extract_indices(node.target)
        rhs_indices = self._extract_indices(node.value.left, node.value.right)
        
        # 构建 MatMul Op
        matmul_op = TileOp(
            "matmul",
            inputs=[A_tile, B_tile],
            outputs=[C_tile]
        )
        matmul_op.attrs.update({
            "tile_m": lhs_indices["M"],
            "tile_n": lhs_indices["N"],
            "tile_k": rhs_indices["K"]
        })
        return matmul_op

💡 关键点：索引符号（如 M, N, K）在 parallel 作用域中绑定为具体 Tile Size。

---

4. 图表示机制：依赖图与内存生命周期建模

构建的 IR 节点被组织为 有向无环图（DAG），用于后续优化。

4.1 依赖图构建

每个 TileOp 的输入/输出隐式定义数据依赖：

# pypto/ir/graph.py
class IRGraph:
    def __init__(self):
        self.nodes: List[IRNode] = []
        self.edges: Dict[IRNode, List[IRNode]] = defaultdict(list)

    def add_node(self, node: IRNode):
        self.nodes.append(node)
        # 为每个输出建立反向依赖
        for inp in getattr(node, 'inputs', []):
            self.edges[inp].append(node)

    def get_dependencies(self, node: IRNode) -> List[IRNode]:
        # 返回所有必须先于 node 执行的操作
        deps = []
        for inp in node.inputs:
            if hasattr(inp, 'producer'):
                deps.append(inp.producer)
        return deps

例如，StoreOp 依赖于其输入 Tile 的生产者（如 MatMulOp）。

4.2 内存生命周期分析

PyPTO 引入 Lifetime Interval 模型，用于内存复用：

# pypto/analysis/liveness.py
class LifetimeAnalyzer:
    def analyze(self, graph: IRGraph):
        live_ranges = {}
        for node in topological_order(graph):
            for out in node.outputs:
                live_ranges[out] = LiveRange(start=node.id)
            
            for inp in node.inputs:
                if inp in live_ranges:
                    live_ranges[inp].end = node.id
        
        return live_ranges

该信息被 MemoryPlanningPass 用于分配 Local Buffer：

# pypto/passes/memory_planning.py
class MemoryPlanningPass:
    def run(self, graph: IRGraph):
        liveness = LifetimeAnalyzer().analyze(graph)
        allocator = BufferAllocator()
        
        for tile in liveness:
            buf = allocator.allocate(
                size=tile.size,
                lifetime=liveness[tile]
            )
            tile.attrs["buffer_id"] = buf.id

📌 效果：多个不重叠生命周期的 Tile 可共享同一片 Local Memory。

---

5. 后端集成：IR 到 pto-isa 的映射

最终，IR Graph 被遍历生成目标代码。

5.1 pto-isa 代码生成

# pypto/backend/ptoc_codegen.py
class PTOCCodegen:
    def emit(self, graph: IRGraph):
        code = []
        for node in topological_order(graph):
            if isinstance(node, LoadOp):
                code.append(f"TLoad({node.outputs[0]}, {node.inputs[0]}, ...);")
            elif isinstance(node, TileOp) and node.op_type == "matmul":
                code.append(f"TMATMUL({node.inputs[0]}, {node.inputs[1]}, {node.outputs[0]});")
            elif isinstance(node, StoreOp):
                code.append(f"TStore({node.inputs[0]}, {node.outputs[0]}, ...);")
        return "\n".join(code)

5.2 支持 Auto-Tuning

PyPTO 还支持将 IR 与搜索空间结合：

# examples/tune_gemm.py
search_space = {
    "tile_m": [64, 128, 256],
    "tile_n": [64, 128, 256],
    "tile_k": [32, 64]
}

for config in search_space:
    ir_graph = build_gemm_ir(config)
    kernel = codegen(ir_graph)
    latency = benchmark(kernel)
    record(latency, config)

该机制使 PyPTO 成为自动性能调优的理想载体。

---

结语

CANN pypto 通过精心设计的 IR 节点体系与图表示机制，成功将 Python 的高层表达能力与 Tile 级硬件控制相结合。其 IR 不仅精确捕获了并行作用域、内存空间与操作语义，还为内存规划、算子融合与自动调优提供了坚实基础。随着 BiSheng 编译器对 PyPTO IR 的原生支持（见 Roadmap），这一编程范式有望成为 CANN 生态中连接算法创新与硬件加速的关键桥梁。

对于希望深入定制高性能算子的开发者而言，掌握 PyPTO 的 IR 构建与图优化机制，是释放硬件潜力的核心技能。

cann组织链接：https://atomgit.com/cann
pypto仓库链接：https://atomgit.com/cann/pypto