基于Ascend C开发的Vector算子模板库-ATVOSS 技术深度解读

ATVOSS 技术深度解读

一、项目概述

ATVOSS（Ascend C Templates for Vector Operator Subroutines）是一套基于Ascend C开发的Vector算子模板库，于2025年11月首次上线。该项目致力于为昇腾硬件上的Vector类融合算子提供极简、高效、高性能、高拓展的编程方式，大幅降低AI算子开发的复杂度。

核心定位

• 面向硬件：专为昇腾AI处理器（Ascend 910B/910C）设计
• 技术栈：基于Ascend C和CANN生态
• 开发范式：采用C++模板元编程和表达式模板技术
• 应用场景：Vector类融合算子的快速开发与性能优化

二、核心技术架构

ATVOSS采用了五层分层设计，从上至下形成了完整的编程模型：

2.1 五层架构设计

┌─────────────────────────────────────┐
│     Device 层 (Host侧调用)           │  ← ACL资源管理、数据管理、Kernel调用
├─────────────────────────────────────┤
│     Kernel 层 (多核调度)             │  ← 多核任务分解、Block调度
├─────────────────────────────────────┤
│     Block 层 (单核流水)              │  ← Tile分解、数据搬运、流水编排
├─────────────────────────────────────┤
│     Tile 层 (大块计算)               │  ← 封装Ascend C API、大块操作
├─────────────────────────────────────┤
│     Basic 层 (基础API)               │  ← Ascend C基础数据搬运与计算
└─────────────────────────────────────┘

(1) Device层

• 职责：Host侧调用总入口
• 功能：参数校验、ACL资源管理、Host与Device数据管理、切分计算、Workspace管理、Kernel调用
• 实现方式 ：通过DeviceAdapter模板统一封装

(2) Kernel层

• 职责：Kernel函数总入口
• 功能：多核间任务分解、Block调度控制
• 核心特性：支持均匀切分和整尾块切分策略
• 配置参数：最大核数、切分策略

(3) Block层

• 职责：单核任务处理
• 功能：单核任务分解到多个Tile块、数据搬运/计算流水编排
• 内存管理：UB（Unified Buffer）空间分配与管理
• 流水支持：可使用TPipe进行内存管理和同步

(4) Tile层

• 职责：大块数据操作
• 功能：对Ascend C基础API进行封装，提供更大Tile块的搬运、计算能力
• 支持的操作：加减乘除、开方、指数、广播、规约求和、类型转换等

(5) Basic层

• 职责：基础操作
• 功能：直接使用Ascend C基础API完成数据搬运计算等基础操作

2.2 分层优势

1. 关注点分离：每层专注于特定粒度的任务，降低开发复杂度
2. 模板化设计：各层均采用模板类，支持灵活组合与扩展
3. 性能优化：每层都有针对性的优化策略（多核并行、流水、内存管理）
4. 易于维护：清晰的层次结构便于问题定位和功能扩展

三、表达式模板技术

ATVOSS的核心创新在于采用了**C++表达式模板（Expression Templates）**技术，这是一种高级的C++元编程技巧。

3.1 表达式模板原理

表达式模板通过运算符重载 和类型计算 实现延迟计算：

// 当写下 out = in1 + in2 时
// 不会立即执行加法运算，而是构建一颗类型化的抽象语法树
Expression<OpAssign<Param<3>, OpAdd<Param<1>, Param<2>>>>

关键要素：

1. 延迟计算：操作符调用时仅记录操作，不立即执行
2. 类型化AST：形成编译期可分析的抽象语法树
3. 零运行时开销：所有类型计算在编译期完成

3.2 核心模板组件

(1) Expression模板 - 基础表达式节点

template <typename T>
struct Expression {
    using Type = T;
    static constexpr bool hasData = HasDataTrait<T>::value;
    T const data{};
    
    template <typename U>
    constexpr auto operator=(Expression<U> u);
};

作用：统一包装器，将任何元素（字面量、参数、运算）包装为表达式系统的一部分

(2) Param模板 - 参数占位符

template <std::size_t N, typename T, ParamUsage V = ParamUsage::in>
struct Param {
    using Type = T;
    static constexpr std::size_t number = N;
    static constexpr ParamUsage usage = V;  // in/out/in_out
};

作用：表示尚未赋值的输入/输出参数，编译期记录类型和用途

(3) LocalVar模板 - 局部变量

template <std::size_t N, typename T, typename L = void>
struct LocalVar {
    using Type = T;
    using Like = L;  // 参照某个参数的规格初始化
    static constexpr std::size_t number = N;
};

作用：提供表达式内部的临时变量支持，类似函数局部变量

(4) 运算符模板 - 操作节点

// 二元运算基类
template <typename T, typename U>
struct BinaryOp {
    constexpr const T& GetLhs() const;  // 左操作数
    constexpr const U& GetRhs() const;  // 右操作数
};

// 具体运算（加减乘除等）
template<typename T, typename U>
struct OpAdd : BinaryOp<T, U> { };

template<typename T, typename U>
struct OpMul : BinaryOp<T, U> { };

3.3 表达式求值机制

通过Evaluator模板实现递归求值：

template <typename T>
struct Evaluator {
    template <typename ArgTup, typename LocalVarTup>
    __aicore__ decltype(auto) operator()(
        const T& value, 
        ArgTup& args, 
        LocalVarTup& localVars) const
    {
        // 递归求值逻辑
    }
};

求值策略：

1. 深度优先遍历AST
2. 对每个节点特化Evaluator
3. 递归终止条件：Param、LocalVar或常量
4. 最终调用Ascend C API执行实际计算

3.4 编译期优化

利用C++模板元编程实现编译期优化：

// 编译期参数收集
template <typename T>
struct Params {
    using Type = typename Detail::UniqueParams<T>::Type;
};

// TypeList操作（函数式编程风格）
- Append_t     // 追加类型
- Filter_t     // 过滤类型
- Map_t        // 映射转换
- Unique_t     // 去重
- ForEach      // 遍历操作

优势：

• 编译期确定内存布局
• 自动推导参数类型和数量
• 零运行时类型检查开销
• 编译期错误检测

四、编程范式实践

4.1 RMS Norm算子示例

以下是一个完整的RMS Norm算子实现，展示了ATVOSS的编程范式：

static constexpr int32_t HEIGHT = 1;
static constexpr int32_t WIDTH = 32;

struct RmsNormCompute {
    template <template <typename> class Tensor>
    __host_aicore__ constexpr auto Compute() const
    {
        // 定义输入输出
        auto in1 = Atvoss::PlaceHolder<1, Tensor<T>, Atvoss::ParamUsage::in>();
        auto in2 = Atvoss::PlaceHolder<2, Tensor<T>, Atvoss::ParamUsage::in>();
        auto out = Atvoss::PlaceHolder<3, Tensor<T>, Atvoss::ParamUsage::out>();
        auto temp = Atvoss::PlaceHolderTmpLike<1>(in1);
        
        // 计算表达式（用逗号分隔的操作序列）
        return (temp = in1 * in1,                          // 1. 平方
                out = ReduceSum<Pattern::AR>(temp),        // 2. 求和
                out = Broadcast<Pattern::AB>(out),         // 3. 广播
                temp = Divs<WIDTH>(out),                   // 4. 除以常数
                out = Sqrt(temp),                          // 5. 开方
                temp = in1 / out,                          // 6. 归一化
                out = in2 * temp);                         // 7. 缩放
    }
};

4.2 层级模板组装

// 配置Block层策略
static constexpr Atvoss::EleWise::BlockPolicy<TileShape> blockPolicy {
    190 * 1024,    // UB空间大小
    TileShape{}    // Tile形状
};

// 配置Kernel层策略
static constexpr Atvoss::EleWise::KernelPolicy kernelPolicy {
    48,            // 最大核数
    KernelPolicySegment::UniformSegment  // 均匀切分策略
};

// 组装Block层模板
using BlockOp = Atvoss::EleWise::BlockBuilder<
    RmsNormCompute,
    blockPolicy,
    BlockConfig,
    BlockScheduleWithTPipe
>;

// 组装Kernel层模板
using KernelOp = Atvoss::EleWise::KernelBuilder<
    BlockOp,
    kernelPolicy
>;

// 封装Device适配器
using DeviceOp = Atvoss::DeviceAdapter<KernelOp>;

4.3 调用方式

// 准备数据
std::vector<float> v1(32 * 32, 1.0F);
std::vector<float> v2(32 * 32, 2.0F);
std::vector<float> v3(32 * 32);
uint32_t shape[2] = {32, 32};

// 构造Tensor
Atvoss::Tensor<float> t1(v1.data(), shape);
Atvoss::Tensor<float> t2(v2.data(), shape);
Atvoss::Tensor<float> t3(v3.data(), shape);

// 构建参数（支持流式API）
auto arguments = Atvoss::ArgumentsBuilder{}
    .input(t1, t2)
    .output(t3)
    .build();

// 执行算子
DeviceOp deviceOp;
deviceOp.Run(arguments);

五、Tile层API支持

ATVOSS当前支持丰富的Tile层API，均采用统一的接口范式：

API接口名	实现接口	功能描述	支持类型
+	AddAssign	按元素求和	int32_t/float
-	SubAssign	按元素求差	int32_t/float
*	MulAssign	按元素求积	int32_t/float
/	DivAssign	按元素求商	float
Divs	DivsAssign	矢量除以标量	int32_t/float
Exp	ExpAssign	自然指数	float
Sqrt	SqrtAssign	开方运算	float
Power	PowerAssign	幂运算（目前仅平方）	float
Broadcast	BroadcastAssign	广播操作（二维）	float
ReduceSum	ReduceSumAssign	规约求和（二维）	float
Cast	CastAssign	类型转换	half/float/int32_t

特点：

• 双层接口设计：用户层的表达式API + 实现层的Assign接口
• 类型安全：编译期类型检查
• 零开销抽象：表达式最终内联为Ascend C API调用

六、性能优化机制

6.1 多核并行策略

Kernel层支持两种切分策略：

1. 均匀切分（UniformSegment）

   • 每个核处理尽量相同的数据量
   • 适合负载均衡场景

2. 整尾块切分（FullAddTail）

   • 多数核处理整块，最后一个核处理尾块
   • 适合数据量不能均匀分配的场景

6.2 流水编排

Block层支持使用TPipe进行流水编排：

using BlockOp = Atvoss::EleWise::BlockBuilder<
    Compute,
    blockPolicy,
    BlockConfig,
    BlockScheduleWithTPipe  // 使用TPipe流水
>;

流水优势：

• 数据搬运与计算重叠
• 多级缓存利用
• 提高硬件利用率

6.3 内存优化

UB（Unified Buffer）管理：

struct BlockPolicy {
    uint32_t ubSizeMax = 190 * 1024;  // 最大UB空间
    Shape tileShape{};                 // Tile对齐策略
};

struct UbAssign {
    uint32_t ubInNum;         // 输入UB Tensor数量
    uint32_t ubOutCnt;        // 输出UB Tensor数量
    uint32_t ubTmpCnt;        // 临时UB Tensor数量
    uint32_t eleNumSingleTensor;  // 单个Tensor元素数
};

优化措施：

• 编译期计算存活节点
• 自动内存分配与复用
• Tile大小自动调优

6.4 编译期优化

• 空基类优化 ：CompressedPair避免空类内存开销
• constexpr计算：尽可能在编译期完成计算
• 类型擦除优化：避免不必要的类型信息传递

七、扩展性设计

7.1 自定义Tile层API

用户可以方便地扩展新的Tile层API：

步骤1：定义表达式

// 在 math.h 中
DeclareUnaryOp(MyCustomOp);   // 一元操作
DeclareBinaryOp(MyBinaryOp);  // 二元操作

步骤2：实现Evaluator

// 在 tile_evaluator_math.h 中
template<typename T, typename U>
struct Evaluator<OpAssign<T, OpMyCustomOp<U>>> {
    template<typename ArgTup, typename LocalVarTup>
    __aicore__ auto operator()(/*...*/) const {
        return MyCustomOpAssign(/*...*/);
    }
};

步骤3：调用Ascend C API

// 在 tile_ascendc_math.h 中
template <typename OperationShape, typename T>
__aicore__ inline void MyCustomOpAssign(
    LocalTensor<T>& dst, 
    const LocalTensor<T>& src, 
    OperationShape& operationShape)
{
    AscendC::CustomOp(dst, src, operationShape.axis0);
}

7.2 自定义调度策略

继承BaseBlockSchedule或BaseKernelSchedule可实现自定义调度：

template <typename Compute, const auto& Policy, typename ScheduleCfg>
class MyCustomSchedule : public BaseBlockSchedule</*...*/> {
    // 自定义Run方法
    template <typename ArgTup>
    __aicore__ inline void Run(ScheduleCfg& cfg, ArgTup& argTuple) {
        // 自定义调度逻辑
    }
};

最后欢迎加入CANN社区：https://atomgit.com/cann，更多的算子都可以在cann社区看到