CANN pyasc 的 Python 语法到硬件指令的编译流程

在当今数字化业务场景中，网页自动化与智能数据采集已成为提升研发效率、实现业务监控和构建AI训练数据集的关键能力。然而，传统方案如Selenium与Requests的割裂使用，常导致开发复杂度高、维护成本大、执行效率低等问题。开发者迫切需要一个**统一、高效、易用**的自动化工具。

咕咕日志

418人浏览 · 2026-02-06 20:47:07

咕咕日志 · 2026-02-06 20:47:07 发布

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

前言

在 AI 算子开发领域，开发者长期面临“性能”与“易用性”的两难困境：底层 C/C++ 或专用 DSL（如 CUDA、Ascend C）虽能榨取硬件性能，但开发门槛高；而 Python 虽简洁灵活，却难以直接部署到加速器。CANN（Compute Architecture for Neural Networks）开源项目中的 pyasc 仓库（https://atomgit.com/cann/pyasc）提出了一种创新解决方案——将 Python 语法直接编译为高效硬件指令，实现“Python 原生写法，硬件级性能”。

1. 编译架构概览：四阶段流水线

pyasc 采用经典的编译器分层架构，分为四个阶段：

该流程由 pyasc/compiler/compiler.py 驱动，最终生成可被 CANN 运行时直接调用的 .so 文件。

关键组件

Frontend：Python AST 解析与类型标注；
Middle-end：MLIR 表示与优化；
Backend：目标硬件代码生成；
Runtime：Kernel 加载与执行。

2. 前端：Python AST 解析与受限语法支持

pyasc 并非支持完整 Python，而是聚焦于算子开发所需的子集。

2.1 支持的语法特性（截至 v1.2）

基本控制流：if/else、for（仅静态范围）
变量声明与赋值
函数定义（@kernel 装饰器）
Tensor 操作：load()、store()、alloc()
内置函数：ceil_div()、get_block_id()

❌ 不支持：动态类型、递归、异常处理、全局变量。

2.2 AST 转换与类型推导

以一个简单 ReLU 算子为例：

# tutorials/relu.py
from pyasc import kernel, Tensor, float16

@kernel
def relu_kernel(input: Tensor[float16], output: Tensor[float16]):
    idx = get_block_id()
    val = input.load([idx])
    output.store([idx], max(val, 0.0))

pyasc 首先通过 ast.parse() 获取 AST，然后进行类型标注：

# pyasc/python/src/ast_transformer.py
class PyAscTypeInfer(ast.NodeVisitor):
    def visit_FunctionDef(self, node):
        # 从注解提取 Tensor 类型
        for arg in node.args.args:
            if hasattr(arg, 'annotation'):
                dtype = self.parse_dtype_annotation(arg.annotation)
                self.symbol_table[arg.arg] = TensorType(dtype)
    
    def visit_Call(self, node):
        if node.func.id == "load":
            # 推导 load 返回类型 = Tensor 元素类型
            tensor_name = node.args[0].id
            elem_type = self.symbol_table[tensor_name].element_type
            node.inferred_type = elem_type

✅ 优势：静态类型推导确保后续 MLIR 生成无歧义。

3. 中端：MLIR 表示与 asc.dialect 定义

pyasc 将类型标注后的 AST 转换为 MLIR（Multi-Level Intermediate Representation）。

3.1 自定义 Dialect：`asc`

在 pyasc/lib/Dialect/ASCDialect.cpp 中定义核心操作：

// pyasc/lib/Dialect/ASCOps.td
def ASC_LoadOp : ASC_Op<"load", [NoSideEffect]> {
  let arguments = (ins ASC_Tensor:$tensor, I64ArrayAttr:$indices);
  let results = (outs AnyType:$result);
}

def ASC_StoreOp : ASC_Op<"store"> {
  let arguments = (ins ASC_Tensor:$tensor, I64ArrayAttr:$indices, AnyType:$value);
}

这些操作直接对应硬件内存访问原语。

3.2 AST 到 MLIR 转换

// pyasc/lib/Target/ASTToMLIR.cpp
mlir::Value MLIRBuilder::VisitLoad(const LoadNode* node) {
    auto tensor = symbol_table_[node->tensor_name];
    auto indices = ConvertIndices(node->indices);
    
    // 创建 asc.load 操作
    return rewriter_.create<ASC::LoadOp>(
        loc_, tensor, indices
    );
}

mlir::Value MLIRBuilder::VisitStore(const StoreNode* node) {
    auto tensor = symbol_table_[node->tensor_name];
    auto indices = ConvertIndices(node->indices);
    auto value = VisitExpr(node->value);
    
    rewriter_.create<ASC::StoreOp>(loc_, tensor, indices, value);
    return value;
}

生成的 MLIR 如下：

func.func @relu_kernel(%arg0: !asc.tensor<f16>, %arg1: !asc.tensor<f16>) {
  %0 = asc.get_block_id
  %1 = asc.load %arg0[%0] : f16
  %2 = arith.constant 0.0 : f16
  %3 = arith.maxf %1, %2 : f16
  asc.store %arg1[%0], %3 : f16
  return
}

4. 后端：硬件指令生成与优化

4.1 内存分配：LocalMemAllocator

pyasc 引入 LocalMemAllocator 管理片上内存：

# tutorials/matmul.py
from pyasc import alloc_local

@kernel
def matmul_kernel(A, B, C):
    tile_a = alloc_local((16, 16), dtype=float16)  # ← 片上内存
    tile_b = alloc_local((16, 16), dtype=float16)
    # ...

在 MLIR 中表示为：

%tile_a = asc.alloc_local {shape = [16, 16], dtype = f16}

后端将其映射为硬件 L1/L2 缓存分配指令。

4.2 代码生成：LLVM IR 与二进制

MLIR 经过 asc-lower-to-llvm Pass 转换为 LLVM IR：

// pyasc/lib/Target/ASCToLLVM.cpp
void ASCtoLLVMLowering::lowerLoadOp(ASC::LoadOp op) {
    auto ptr = getTensorDataPtr(op.tensor());
    auto offset = calculateLinearOffset(op.indices());
    auto addr = builder.CreateGEP(ptr, offset);
    replaceOpWithNewOp<LLVM::LoadOp>(op, addr);
}

最终通过 LLVM JIT 编译为二进制：

// pyasc/lib/Target/KernelCompiler.cpp
std::unique_ptr<llvm::Module> module = mlirToLLVMIR(mlir_module);
auto engine = llvm::EngineBuilder(std::move(module)).create();
void* kernel_func = engine->getFunctionAddress("relu_kernel");

5. 运行时：Kernel 加载与执行

5.1 Python 接口封装

编译后的 Kernel 通过 pybind11 暴露给 Python：

// pyasc/python/src/pybind_kernel.cpp
PYBIND11_MODULE(_pyasc_core, m) {
    m.def("launch_kernel", [](const std::string& kernel_name,
                              const std::vector<Tensor>& inputs,
                              const std::vector<Tensor>& outputs) {
        auto func = KernelRegistry::Lookup(kernel_name);
        func(inputs, outputs); // 调用 JIT 生成的函数指针
    });
}

5.2 端到端使用示例

# tutorials/relu_run.py
import numpy as np
from pyasc import compile, Tensor

# 1. 编译 Kernel
compiled_relu = compile(relu_kernel)

# 2. 准备数据
input_data = np.random.randn(1024).astype(np.float16)
output_data = np.zeros_like(input_data)

# 3. 执行
compiled_relu(Tensor(input_data), Tensor(output_data))

print("ReLU executed successfully!")

compile() 内部完成 AST→MLIR→LLVM→JIT 全流程，并缓存至 ~/.cache/pyasc/。

6. 性能与调试支持

6.1 性能对比（Vector Add，1M 元素）

实现方式	延迟（μs）	相对性能
NumPy	280	1.0x
Handwritten C	12	23x
pyasc	15	18.7x

测试环境：CANN 9.0，A3 芯片。

6.2 调试工具

MLIR Dump：PYASC_DUMP_MLIR=1 输出中间表示；
Kernel Disassembly：PYASC_DUMP_ASM=1 查看汇编；
内存检查：自动插入越界访问断言。

结语

CANN pyasc 通过将 Python 语法精准映射到 MLIR 中间表示，并结合硬件亲和的代码生成策略，成功弥合了高级语言易用性与底层硬件性能之间的鸿沟。其不仅保留了 Python 的简洁语法，更通过静态类型推导、自定义 Dialect 与 LLVM 后端，实现了接近手写 C 的执行效率。作为 CANN 算子开发生态的重要一环，pyasc 为算法工程师提供了“零上下文切换”的高性能开发体验，也为社区贡献自定义算子降低了门槛。随着对更多 Python 特性（如 list comprehension、context manager）的支持，pyasc 的能力边界将持续扩展。

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