---

• 个人首页： 永远都不秃头的程序员(互关)
• C语言专栏:从零开始学习C语言
• C++专栏:C++的学习之路
• K-Means专栏:K-Means深度探索系列
• 本章所属专栏:CANN系列

---

一、AIGC模型对自定义算子的迫切需求

AIGC模型的创新往往伴随着对传统神经网络架构的突破，例如扩散模型中的U-Net变体、Transformer中的各类注意力机制、或者一些新颖的激活函数、归一化层等。这些创新结构在计算模式上可能具有以下特点：

1. 高度融合的计算模式：多个逻辑上独立的算子（如矩阵乘法、元素级操作、激活函数）在模型中紧密相连，形成一个计算密集型的数据流。将它们融合成一个自定义算子，可以显著减少数据搬运和内存访问，降低延迟。
2. 不规则的计算逻辑：部分AIGC算法可能涉及非标准的采样、稀疏计算、动态裁剪或条件逻辑，这些无法简单通过现有通用算子高效实现。
3. 特定硬件优化需求：为了最大限度地发挥昇腾AI处理器的异构计算能力，开发者有时需要编写专门针对硬件特性进行优化的算子，以实现极致性能。

CANN的TBE正是解决这些挑战的利器。它允许开发者直接在昇腾AI处理器上定义和实现这些定制化的计算逻辑，从而为AIGC模型提供独一无二的性能优势。

---

二、深度实践：基于cann-operator-sample的AIGC定制算子开发

cann-operator-sample仓库是CANN社区中一个非常宝贵的资源，它提供了丰富的自定义算子开发示例，涵盖了从简单元素级操作到复杂融合算子的各种场景。通过学习这些示例，我们可以掌握为AIGC模型开发高性能算子的精髓。

我们将以一个在AIGC，特别是Transformer-based模型中常见的**“融合门控自注意力输出投影”（Fused Gated Self-Attention Output Projection）**为例。在大型语言模型或图像生成Transformer中，注意力机制的输出往往需要经过一个前馈网络，其中可能包含门控机制（例如Gated Linear Unit, GLU的变体）。将矩阵乘法、门控激活和元素级乘法融合，能有效提升性能。

1. 算子原型定义（op_proto）

首先，我们需要在cann-operator-sample的结构下，在op_proto目录中定义新算子的输入、输出和属性。这部分使用C++描述，向CANN框架声明算子的接口。

// 示例：定义一个名为 "FusedGatedAttention" 的算子原型
// 位于 cann-operator-sample/op_proto/fused_gated_attention_op.cc (概念性代码)
#include "graph/operator_reg.h"

namespace ge {
REG_OP(FusedGatedAttention)
    .INPUT(query, TensorType({DT_FLOAT16, DT_FLOAT})) // 注意力查询输出
    .INPUT(weight_proj, TensorType({DT_FLOAT16, DT_FLOAT})) // 第一个投影层权重
    .INPUT(weight_gate, TensorType({DT_FLOAT16, DT_FLOAT})) // 门控层权重
    .OUTPUT(output, TensorType({DT_FLOAT16, DT_FLOAT})) // 融合后的输出
    .ATTR(activation_type, String, "sigmoid") // 门控激活函数类型，例如 "sigmoid" 或 "gelu"
    .OP_END_DEFINE();
} // namespace ge

这个C++代码片段定义了一个名为FusedGatedAttention的算子。它有三个输入（注意力输出、两个投影权重）和一个输出，并包含一个可配置的门控激活函数类型属性。这是CANN识别、校验和调度该算子的基础。

2. TBE核函数实现（tbe/impl）

接下来，是算子性能的核心实现：TBE核函数。这部分通常使用Python和TBE DSL编写，描述底层的计算逻辑和调度策略。我们将在tbe/impl目录下创建相应的Python文件。

# 示例：FusedGatedAttention 算子的TBE核函数实现
# 位于 cann-operator-sample/tbe/impl/fused_gated_attention_impl.py (概念性代码)
import te.lang.cce as cce
from te import platform as tbe_platform
from te import tvm
from topi.cce import util

@util.check_input_type(dict, dict, dict, dict, dict, str, str)
def fused_gated_attention_compute(query, weight_proj, weight_gate, output, activation_type="sigmoid", kernel_name="fused_gated_attention"):
    """
    Compute for Fused Gated Self-Attention Output Projection.
    Steps:
    1. proj_out = matmul(query, weight_proj)
    2. gate_val = matmul(query, weight_gate)
    3. gate_act = activation(gate_val) (e.g., sigmoid)
    4. output = proj_out * gate_act (element-wise multiplication)
    """
    shape_query = query.get("shape")
    shape_weight_proj = weight_proj.get("shape")
    shape_weight_gate = weight_gate.get("shape")
    dtype = query.get("dtype").lower()

    # 创建tvm placeholder
    data_query = tvm.placeholder(shape_query, name="data_query", dtype=dtype)
    data_weight_proj = tvm.placeholder(shape_weight_proj, name="data_weight_proj", dtype=dtype)
    data_weight_gate = tvm.placeholder(shape_weight_gate, name="data_weight_gate", dtype=dtype)

    # 1. Matmul for projection
    proj_out = cce.matmul(data_query, data_weight_proj)

    # 2. Matmul for gate
    gate_val = cce.matmul(data_query, data_weight_gate)

    # 3. Apply activation for gate
    if activation_type == "sigmoid":
        gate_act = cce.sigmoid(gate_val)
    elif activation_type == "gelu":
        gate_act = cce.gelu(gate_val)
    else:
        raise RuntimeError(f"Unsupported activation_type: {activation_type}")

    # 4. Element-wise multiplication
    res = cce.mul(proj_out, gate_act)

    with tvm.tag_scope(tag=util.set_ ভারত_tag): # 这里可以自定义tag或使用通用tag
        output_tensor = cce.cast_to(res, dtype) # 确保输出数据类型一致

    return output_tensor

@util.check_input_type(dict, dict, dict, dict, str, str)
def fused_gated_attention(query, weight_proj, weight_gate, output, activation_type="sigmoid", kernel_name="fused_gated_attention"):
    """
    Entry function for Fused Gated Self-Attention Output Projection.
    """
    # ... (这里会进行输入形状、数据类型等的详细校验，确保算子使用正确) ...
    util.check_shape_rule(query.get("shape"))
    util.check_dtype_rule(query.get("dtype"), ("float16", "float32"))
    # ... 其他校验 ...

    res = fused_gated_attention_compute(query, weight_proj, weight_gate, output, activation_type, kernel_name)

    with tvm.target.cce():
        sch = generic.auto_schedule(res) # 自动调度优化，也可手动编写调度策略

    config = {"name": kernel_name,
              "tensor_list": [query, weight_proj, weight_gate, output],
              "op_pattern": 1} # 1表示融合算子

    return sch, config

这个Python代码段详细展示了如何使用TBE DSL来描述“融合门控自注意力输出投影”的计算过程。通过cce.matmul实现矩阵乘法，cce.sigmoid/cce.gelu实现门控激活，再通过cce.mul完成元素级乘法。最关键的是，所有这些操作都被描述在一个TBE核函数中，最终会被CANN编译成一个高度优化的计算单元，从而避免了多次数据在片内和片外内存之间搬运的开销。 这对于AIGC模型中常见的巨大中间特征图和大量计算至关重要。

3. 编译、注册与集成

完成自定义算子开发后，将其通过CANN的编译工具链进行编译，并注册到CANN运行时环境。当AIGC模型通过ATC（Ascend Tensor Compiler）进行转换时，如果模型中识别到该自定义算子，CANN将自动调用我们实现的TBE核函数，从而在昇腾AI处理器上以最优性能执行。

---

三、自定义算子对AIGC性能的深远影响

通过上述自定义算子的开发实践，我们能够为AIGC模型带来显著的性能提升：

• 降低计算延迟：通过算子融合，减少了核函数启动次数和数据搬运，从而显著缩短了AIGC推理的整体延迟，尤其对于实时生成场景（如直播美颜、实时语音合成）意义重大。
• 提升吞吐量：在单位时间内处理更多的生成请求，这对于批量生成内容（如大规模图像生成、视频渲染）的应用至关重要。
• 实现架构创新：开发者不再受限于标准算子库，可以自由地设计和实现前沿的AIGC模型结构，推动技术边界。
• 优化资源利用：TBE允许开发者精细控制计算流程，最大限度地利用昇腾AI处理器的计算、存储和带宽资源。

这种深度定制化的能力，是CANN赋能AIGC从“可用”走向“高效”、“卓越”的关键一步。

---

四、展望未来：CANN与AIGC共绘智能蓝图

AIGC技术仍在飞速发展，新的模型架构和优化方法层出不穷。CANN作为昇腾AI全栈软件的核心，将持续演进，提供更强大、更易用的自定义算子开发工具和更智能的编译优化能力。我们相信，通过CANN的赋能，开发者将能够更高效地探索AIGC的无限可能，共同构建一个由AI驱动的智能内容新世界！

---

CANN组织链接：https://atomgit.com/cann
本文实践参考仓库链接：https://atomgit.com/cann/cann-operator-sample

---