SIMD与SIMT混合编程

抽象硬件架构

AI Core上SIMD（Single Instruction Multiple Data，单指令多数据）与SIMT（Single Instruction Multiple Thread，单指令多线程）混合编程，结合了SIMD多数据并行计算能力与SIMT离散访存的优势，实现了向量级并行与线程级并行的高效协同。目前，该技术仅支持Atlas 350 加速卡。

整个执行过程以Vector Function（VF）为基本调度单位，VF为一个基本函数块。SIMD与SIMT混合编程支持在同一算子中灵活切换SIMD与SIMT执行方式，两种不同类型的VF可以快速切换，每个VF代表一个独立的计算任务片段，通常对应算子中的一段可并行处理的逻辑，从而在性能、能效与开发效率之间取得更优平衡。在SIMD与SIMT混合编程中：

在SIMD与SIMT混合编程中，SIMT能够简化复杂算子与不规则控制流的开发；而SIMD基于向量寄存器与指令，实现高效的数据并行处理，即单指令处理多数据，提升每周期的吞吐量。SIMD与SIMT混合编程支持开发者根据算子特征进行精细化映射：规则的逐元素elementwise操作通过SIMD获得高带宽和高算力利用率，不规则或包含分支的计算通过SIMT来缓解发散和控制复杂度。在系统层面，这有利于提高硬件利用率和能效；同时，也更便于进行算子融合和数据复用等优化。同一个算子中既包含SIMD擅长的连续规整计算，也包含SIMT擅长的离散访问等任务，从而在同一算子中同时利用SIMD和SIMT的优势。

如图1所示，SIMD和SIMT的内部执行流程为：

图1 SIMD与SIMT混合编程硬件架构图

SIMD与SIMT编程存在以下差异：

表1 SIMD与SIMT核心差异点

维度

SIMD

SIMT

编程模型

单指令多数据（SIMD），基于向量寄存器与向量指令。

单指令多线程（SIMT），以线程为单位并行执行。

数据搬运方式

通过显式Load/Store将数据从Unified Buffer搬运到向量寄存器。

不支持直接从Global Memory搬运数据到SIMD的向量寄存器。

支持直接读写Global Memory或Unified Buffer中的数据。

适用场景

规则、连续的逐元素操作（elementwise），如卷积、矩阵乘法、向量操作等。

不规则、含分支、动态访问等复杂逻辑，如注意力机制、稀疏操作等。

尽管SIMD与SIMT在编程模型和执行机制上有显著差异，但在硬件层面上共享以下关键资源：

内存层级

在SIMD与SIMT混合编程场景下，可以访问多种内存空间，下表汇总了常见内存类型的作用域及其生命周期。

内存类型

作用域

生命周期

物理位置

特点

全局内存

所有核函数

应用程序

Device

大容量，低带宽

共享内存

单核核函数

核函数

Vector Core

小容量，高带宽

SIMT寄存器

Thread

SIMT VF函数

Device

极小容量，极高带宽

SIMD寄存器

Register File

SIMD VF函数

Vector Core

极小容量，极高带宽

整体内存架构如下图所示：

图2 SIMD与SIMT混合编程内存模型示意图

在SIMT工作模式下，各个内存的工作流程如下：

在SIMD工作模式下，各个内存的工作流程如下：

UB内存分配

UB（即Unified Buffer）内存空间总大小为256KB，参考图3，按功能划分为四个主要区域，从低地址向高地址依次为静态内存、动态内存、 预留空间 、Data Cache，具体结构如下：

  1. 静态内存：从内存的起始地址分配一段指定大小的内存空间，其大小在编译时确定，不可动态修改。 
    1
2
    		 
    // 静态内存通过数组分配，例如：
__ubuf__ char staticBuf[1024];
  2. 动态内存（该方式将在后续版本中支持）：位于静态内存之后，通过<<<>>>中参数dynUBufSize指定的动态内存大小空间，可通过以下方式申请使用：
    • 通过TPipe的相关接口申请。
    • 通过LocalMemAllocator的Alloc接口申请。
    • 使用动态数组分配。 
      // 动态内存通过动态数组分配，例如：
extern __ubuf__ char dynamicBuf[];

    由于上述三种方法申请动态内存时均从静态内存结束位置之后开始分配，如果同时使用可能会导致地址空间重叠，从而引发未定义行为，因此只能选择其中一种方法进行申请。

  3. 预留空间：编译器和Ascend C预留空间，大小固定为8KB。
  4. Data Cache：SIMT专有的Data Cache空间，内存大小必须大于或等于32KB。

动态内存的动态数组分配方式目前开发中，将在后续版本中支持，请关注后续版本。

  • DataCache = UB总大小（256KB） – 静态内存 – 动态内存 – 预留空间(8KB）
  • 若DataCache小于32KB，会出现校验报错。
  • 在SIMD与SIMT混合编程的场景下，算子内部不能使用全部的Unified Buffer空间，除了预留8KB空间外，还需至少为SIMT预留32KB的Data Cache空间。
图3 UB内存分配图

核函数的定义

调用层级

具体支持的调用关系图如下所示。

图4 函数调用关系图

编程示例

样例中介绍的算子完整代码请参见SIMD与SIMT混合编程实现gather&adds算子样例。更多SIMT函数声明的相关约束请参考语法限制。 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
 
__simt_vf__ __launch_bounds__(THREAD_COUNT) inline void simt_gather(
    __gm__ float* input,
    __gm__ uint32_t* index,
    __ubuf__ float* gather_output,
    uint32_t input_total_length,
    uint32_t index_total_length,
    uint32_t output_total_length)
{
    if (threadIdx.x >= output_total_length) {
        return;
    }
    // blockIdx will be supported later.
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= index_total_length) {
        return;
    }

    uint32_t gather_idx = index[idx];
    if (gather_idx >= input_total_length) {
        return;
    }

    gather_output[threadIdx.x] = input[gather_idx];
}

__simd_vf__ inline void simd_adds(__ubuf__ float* output, __ubuf__ float* input,
    uint32_t count, uint32_t one_repeat_size, uint16_t repeat_times)
{
    AscendC::Reg::RegTensor<float> src_reg0;
    AscendC::Reg::RegTensor<float> dst_reg0;
    AscendC::Reg::MaskReg mask_reg; 

    for (uint16_t i = 0; i < repeat_times; i++) {
        mask_reg = AscendC::Reg::UpdateMask<float>(count);
        AscendC::Reg::LoadAlign(src_reg0, input + i * one_repeat_size);
        AscendC::Reg::Adds(dst_reg0, src_reg0, ADDS_ADDEND, mask_reg);
        AscendC::Reg::StoreAlign(output + i * one_repeat_size, dst_reg0, mask_reg);
    }
}

__global__ __aicore__ void gather_and_adds_kernel(__gm__ float* input, __gm__ uint32_t* index, __gm__ float* output,
    uint32_t input_total_length, uint32_t index_total_length)
{
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    AscendC::LocalMemAllocator<AscendC::Hardware::UB> ub_allocator;
    // 1. gather numbers from input.
    uint32_t index_total_length_per_block = index_total_length / AscendC::GetBlockNum();
    AscendC::LocalTensor<float> gather_output = ub_allocator.Alloc<float>(index_total_length_per_block);
    asc_vf_call<simt_gather>(dim3(THREAD_COUNT),input, index,
                             (__ubuf__ float *)gather_output.GetPhyAddr(),
                             input_total_length,
                             index_total_length,
                             index_total_length_per_block);

    // 2. use reg compute api to do addition.
    AscendC::LocalTensor<float> adds_output = ub_allocator.Alloc<float>(index_total_length_per_block);
    constexpr uint32_t one_repeat_size = AscendC::GetVecLen() / sizeof(float);
    uint16_t repeat_times = (index_total_length_per_block + one_repeat_size - 1) / one_repeat_size;
    asc_vf_call<simd_adds>((__ubuf__ float *)adds_output.GetPhyAddr(),
        (__ubuf__ float *)gather_output.GetPhyAddr(), index_total_length_per_block, one_repeat_size, repeat_times);

    AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0);
    AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0);

    // 3. copy data to global memory.
    AscendC::GlobalTensor<float> output_global_tensor;
    output_global_tensor.SetGlobalBuffer(output + index_total_length_per_block * AscendC::GetBlockIdx());
    AscendC::DataCopy(output_global_tensor, adds_output, index_total_length_per_block);
}

int main(int argc, char *argv[])
{
     …
     //numBlocks only supports one dimension currently.
     gather_and_adds_kernel<<<numBlocks, dynUBufSize, stream>>>(input_device, index_device, output_device, input_total_length, index_total_length);
     …
}
父主题： AI Core SIMD编程