第4章 · DSL开发入门

从这一章开始，我们正式进入 DSL（Domain Specific Language） 的世界。说白了，DSL 就是昇腾专门为算子开发设计的一套“方言”，让你用更接近数学表达的方式写算子，而不是跟底层汇编死磕。我个人觉得，DSL 是新手通往熟练工最快的一条路——你想想看，不用管寄存器分配，不用操心流水线，多舒服。

4.1 DSL 编程模型 —— 没那么玄乎

DSL 编程模型，本质上就是一套 计算描述 + 数据映射 的框架。你只需要告诉芯片“我要算什么”，剩下的调度由编译器搞定。我记得第一次接触 DSL 时，心里还犯嘀咕：这玩意儿能跟手写汇编比性能？后来真香了——开发效率提升好几倍，性能损失几乎可以忽略。

在 CANN 里，DSL 算子通常继承 tbe.dsl.compute 或直接使用 tbe.dsl 下的接口。核心就三个东西：张量 (Tensor)、计算表达式、调度策略。不过入门阶段，调度可以先放一放，咱们先把计算写对。

DSL 编程口诀： 定义输入 → 描述计算 → 绑定输出。就这么简单。

4.2 张量 (Tensor) 的定义与操作

张量就是带形状的多维数组。在 DSL 里，你不需要手动管理内存，只需要声明 shape 和 dtype。我习惯把 Tensor 想象成一个“智能盒子”——你告诉它大小和类型，它自己会找地方住。

import tbe.dsl as dsl
import tbe.common.dtype as dtype

# 定义一个 float16 类型的 1D 张量，长度 1024
tensor_a = dsl.tensor([1024], dtype.float16)
tensor_b = dsl.tensor([1024], dtype.float16)
# 输出张量，形状与输入一致
tensor_c = dsl.tensor([1024], dtype.float16)

你还可以对张量做切片、reshape、广播。不过要注意，DSL 里的操作大多是 符号化 的，不会立即执行，而是构建一张计算图。嗯，这里有个坑：千万别在 DSL 里写 Python 循环去操作每个元素，要用矢量化的表达。

我在项目里遇到过有人试图用 for 循环初始化张量，结果编译报错“不支持动态索引”。记住：DSL 里张量操作要“批量思维”，一次操作一整片数据。

4.3 简单的矢量加法算子实现

光说不练假把式。咱们直接撸一个矢量加法算子：C = A + B。这个算子虽然简单，但五脏俱全——输入、计算、输出、注册，一个不少。

@dsl.operator
def vector_add(a, b, c):
    # a, b, c 都是 Tensor，shape 必须一致
    c_result = a + b          # 矢量加法，一行搞定
    dsl.assign(c, c_result)   # 把结果写回输出张量

# 注册算子信息
vector_add.register("vector_add_fp16", 
                     input_dtypes=[dtype.float16, dtype.float16],
                     output_dtypes=[dtype.float16])

你看，核心计算就一行 a + b。DSL 会自动把它映射到矢量单元。我曾经在调试一个加法算子时，发现输出全是零，查了半天——原来是输入张量的 shape 没对齐。DSL 不会帮你做广播，shape 必须严格匹配。

我曾经因为偷懒，把输入 shape 写成 [1024] 和 [1024, 1]，结果加法变成了广播，性能直接崩了。所以，务必保证参与计算的张量 shape 一致，除非你明确需要广播语义。

4.4 完整示例：从定义到运行

下面是一个可以在 NPU 上跑通的完整示例（省略了环境初始化部分）。你把它复制到你的开发环境里，稍微改改路径就能用。

import tbe.dsl as dsl
from tbe.common.register import register_op

@register_op("vector_add_demo")
def vector_add_demo(a, b, c, kernel_name="vector_add_demo"):
    # a, b, c 是 TVM 张量，DSL 层封装
    c_out = dsl.compute(a.shape, 
                        lambda *idx: a[idx] + b[idx],
                        name="c_out")
    dsl.assign(c, c_out)
    # 调度部分（入门可省略）
    # dsl.auto_schedule([c])

我个人建议，刚开始不要碰调度，先用默认调度跑通。等你对数据流有感觉了，再回来调流水线、tile 大小。

4.5 避坑 & 小技巧

数据类型要统一：fp16 + fp32 会触发隐式转换，性能下降。我一般统一用 fp16。
不要用 Python 原生 print 调试：DSL 是编译执行的，要用 dsl.debug 或者打印计算图。
善用 shape 推导：如果输出 shape 和输入不同，用 dsl.compute 手动指定 lambda 表达式。
多看编译日志：DSL 编译失败时，错误信息里会明确提示是 shape 不匹配还是类型错误。

我刚开始学 DSL 时，总觉得它“黑魔法”太多。后来养成了一个习惯：每写一个算子，先用小 shape（比如 [8]）跑通，再放大。这样定位问题快很多。

接口 / 函数	作用	备注
`dsl.tensor(shape, dtype)`	声明一个张量占位符	不分配实际内存
`dsl.compute(shape, lambda)`	定义计算逻辑	核心 DSL 表达式
`dsl.assign(dst, src)`	将计算结果绑定到输出张量	必须显式调用
`dsl.operator`	装饰器，标记算子函数	用于注册

好了，DSL 入门就这些。说白了，它就是一层“语法糖”，让你用更少的代码写出高效的矢量算子。下一节我们会深入 TIK 编程，那才是真正贴近硬件的玩法。不过别急，先把 DSL 的基础打牢——你想想看，如果连加法都写不顺，后面怎么优化性能？