4. 算子原型定义

—— 输入输出、属性与数据排布格式

好，咱们进入正题。算子原型定义，说白了就是给算子画个「标准像」。你想想看，一个算子要跑在昇腾芯片上，总得告诉CANN三件事：吃什么数据、吐什么结果、有什么脾气（属性）。我刚开始接触CANN时，觉得这步就是填个表格，后来踩了坑才明白——原型定义没写好，后面编译、调试全是泪。

这一章，我就带你把这层窗户纸捅破。咱们从输入输出定义开始，再到属性定义，最后聊聊数据排布格式。嗯，内容有点干，但都是硬通货。

4.1 输入输出定义

先聊输入输出。在CANN里，每个算子都得用OpInput和OpOutput把接口说清楚。我个人习惯先把算子的「数据流」画出来——输入几个张量，输出几个张量，每个张量的形状、数据类型、是否可选，都得白纸黑字写明白。

举个例子，一个简单的Add算子，两个输入一个输出，都是float32，形状相同。但如果你要开发一个BatchNorm，输入就有特征图、均值、方差、gamma、beta，足足5个输入。我当年第一次写BatchNorm原型时，漏了一个输入，编译过了，跑推理时直接崩——嗯，从那以后我养成了「先列清单，再写代码」的习惯。

在op_proto.h里，定义方式大致是这样：

// 伪代码示例
class AddCustom : public OpProto {
    void DefineInputs() override {
        AddInput("x1", TensorType({DT_FLOAT16, DT_FLOAT}))
            .Describe("第一个输入张量");
        AddInput("x2", TensorType({DT_FLOAT16, DT_FLOAT}))
            .Describe("第二个输入张量");
    }
    void DefineOutputs() override {
        AddOutput("y", TensorType({DT_FLOAT16, DT_FLOAT}))
            .Describe("输出张量");
    }
};

注意看，TensorType里可以传多个数据类型，表示这个输入支持多种精度。我个人建议，能支持多精度的尽量支持，别把路走窄了。但也要小心——支持太多类型，后面算子实现的工作量会翻倍。

4.2 属性定义

属性是什么？就是算子的「旋钮」和「开关」。比如卷积的stride、padding，激活函数的alpha值，归一化的epsilon——这些不随输入变化，但又直接影响计算逻辑的参数，都归属性管。

我见过不少新手，喜欢把属性塞到输入里，比如搞个额外的Tensor传参数。千万别这么干！属性就是属性，用Attr定义，清晰又高效。CANN里属性支持的类型挺全：int、float、bool、string、list、tensor等。

属性类型	CANN关键字	常用场景
整型	AttrInt	stride, kernel_size
浮点	AttrFloat	alpha, epsilon
布尔	AttrBool	is_train, use_bias
字符串	AttrString	data_format, activation_type
列表	AttrListInt / AttrListFloat	padding, dilation

定义属性时，我强烈建议给每个属性加默认值。为什么？因为上层框架（比如TensorFlow、PyTorch）在调用算子时，很多参数走默认值。你没给默认值，人家就得显式传，稍微不匹配就报错。我曾经被一个没给默认值的epsilon属性坑过，排查了半天，最后发现是框架侧没传这个参数。

// 属性定义示例
void DefineAttrs() override {
    AddAttr("alpha", AttrValue::CreateFrom(1.0f))
        .Describe("LeakyReLU的负斜率");
    AddAttr("keep_dims", AttrValue::CreateFrom(false))
        .Describe("是否保持维度");
}

4.3 数据排布格式

这块是CANN算子开发里最容易出幺蛾子的地方。数据排布格式，说白了就是张量在内存里怎么摆。昇腾芯片对数据排布有特殊要求，不像GPU那么随意。

常见的排布格式有：

• ND：最普通的排布，按维度顺序连续存储。适合全连接、逐元素算子。
• NCHW：经典的四维格式，batch、通道、高、宽。卷积的标配。
• NHWC：TensorFlow偏爱的格式，通道在最后。昇腾也支持，但性能不如NCHW。
• 5D (NC1HWC0)：昇腾特有的格式，为了适配AI Core的向量化计算。C0通常对齐到16或32。

我刚开始接触5D格式时，整个人是懵的。明明一个四维的卷积，怎么到内存里变成了五维？后来才明白，这是为了数据搬运和向量化计算做的优化。C0维度把通道拆成小块，每个小块正好塞进向量寄存器，计算效率直接拉满。

在原型定义里，数据排布格式通过OpFormat来指定：

// 指定输入输出格式
AddInput("x", TensorType({DT_FLOAT16}))
    .SetFormat(FORMAT_NCHW)
    .Describe("输入特征图");
AddOutput("y", TensorType({DT_FLOAT16}))
    .SetFormat(FORMAT_NC1HWC0)
    .Describe("输出特征图（5D排布）");

这里有个细节：输入和输出的格式可以不一样。CANN会在算子内部自动插入格式转换吗？不会。你得自己在算子实现里处理格式转换，或者用CANN提供的FormatTransform接口。我个人建议，尽量让输入输出格式统一，减少不必要的转换开销。如果实在要转，记得在原型注释里写清楚，方便后面维护的人（包括你自己）理解。

4.4 原型定义的整体流程

说了这么多，咱们串一下。写一个算子原型定义，大致分四步：

1. 确定算子语义：这个算子干什么？输入几个？输出几个？有什么属性？
2. 定义输入输出：用AddInput/AddOutput，指定数据类型和形状约束。
3. 定义属性：用AddAttr，给默认值，加描述。
4. 指定排布格式：用SetFormat，明确每个张量的内存排布。

嗯，看起来不复杂，但每一步都有坑。我的经验是：写原型定义时，多花10分钟想清楚，后面能省10小时。别急着往下走，先把接口定稳了。