存储对齐

本文介绍HIVM中的存储对齐(Stride Alignment)机制,包括 hivm-pre-mark-stride-alignhivm-mark-stride-alignhivm-enable-stride-align 三个Pass的硬件背景、算法原理、接口说明和使用约束。

存储对齐分为"标记"(mark)与"使能"(enable)两个阶段:标记阶段分析每个HIVM算子的memref操作数,确定需要对齐的维度和对齐字节数,并以 annotation.mark 注解的形式写入IR;使能阶段读取这些注解,将标记信息传播到根分配点(memref.alloc),对分配的内存形状进行填充(padding),使各行起始地址满足硬件对齐要求,再通过 memref.subview 切回原始逻辑形状。

硬件背景

昇腾硬件的片上内存(UB、L1、L0C等)对数据访问有严格的地址对齐要求。各类Buffer的对齐要求与功能如下:

Buffer

对齐要求

功能

Unified Buffer (UB)

32字节对齐

通用缓存空间,主要用于向量和标量运算

L1 Buffer

32字节对齐

暂存feature map等卷积使用到的数据

L0A Buffer

512字节对齐

暂存矩阵运算的左矩阵(feature map)

L0B Buffer

512字节对齐

暂存矩阵运算的右矩阵(weight)

L0C Buffer

512字节对齐

暂存矩阵运算的中间结果和输出矩阵

BT Buffer

64字节对齐

BiasTable Buffer,存放矩阵运算中的Bias

FP Buffer

64字节对齐

Fixpipe Buffer,存放量化参数、Relu参数等

除了Buffer自身的起始地址需要对齐外,当memref的最低维(last dim)不连续时——即次低维(sub-tail dim)的stride不等于最低维的size——硬件在逐行访问数据时,每一行的起始地址也需要满足对齐约束。例如,对于 memref<3x13xi8, strided<[16, 1]>>,最低维size为13但stride为16,最低维不连续;16字节的行宽不满足硬件对齐要求(32字节),就需要在分配时将最低维填充到对齐大小32。

DMA搬运指令和Vector计算指令对stride有各自的约束:

  • DMA指令(copy_gm_to_ubcopy_ub_to_gm)要求被搬运的UB memref各行起始地址对齐到32字节(或硬件对应的对齐粒度)。

  • Fixpipe(L0C→UB搬运)有更复杂的M/N轴对齐约束,取决于数据位宽、ChannelSplit模式、NZ2ND/NZ2DN转换模式等。

  • vsstb(stride store with broadcast)指令建议stride对齐到 (512*N + 32) 字节,以避免Bank Conflict。

相关术语说明

术语

含义

stride_align_dims / stride_align_value_in_byte

两个注解属性,分别记录需要对齐的维度索引(从0开始)和对齐字节数。通过 annotation.mark op附加到memref Value上。例如 {hivm.stride_align_dims = array<i32: 1>, hivm.stride_align_value_in_byte = array<i32: 32>} 表示对第1维的stride对齐到32字节。

last discontinuous dim(最后不连续维度)

从最低维(最右侧维度)向高维方向扫描,第一个stride≠size的维度。该维度的stride即为行宽,需要满足对齐要求。如果最低维stride已经是unit stride(即stride==size,最低维连续),则需要继续向高维查找。

对齐传播(Align Propagation)

标记可能落在 memref.subviewmemref.cast 等view-like op的结果上,而非直接落在 memref.alloc 上。对齐传播将标记沿着view链向上回溯到根分配点,同时向下传播到叶子操作数,并在同一算子的多个操作数间取并集,确保所有相关Buffer使用一致的对齐信息。

SkipStrideAlignForVLoad

预标记阶段(PreMarkStrideAlign)在根alloc上设置的跳过标记。某些Buffer由DMA加载后会在Vector Function中被 vlds(仅支持连续加载)读取,如果对其做stride对齐会产生行间间隙,vlds 无法处理非连续布局,因此需要跳过。

算法原理

软件背景

在PlanMemory完成地址分配之前,输入IR中的 memref.alloc 仅声明了逻辑形状和所需内存大小,不包含地址信息。如果某个Buffer的逻辑形状在次低维未对齐(如 memref<37x5x3xi32, strided<[15, 3, 1]>>,次低维stride=15),DMA模版实现中转为2D处理时,各行的起始地址不对齐,会导致功能异常。

存储对齐机制通过以下方式解决该问题:

  1. 在标记阶段,分析每个HIVM算子的memref操作数,找出"最后一个不连续维度"(last discontinuous dim),即从最低维向上第一个stride≠size的维度。该维度的stride就是需要被对齐的行宽。

  2. 在使能阶段,将根 memref.alloc 的对应维度填充到满足硬件对齐要求的size(如将 memref<37x5x3xi32> 的最低维从3填充到8,使stride=8*4字节=32字节满足对齐),再通过 memref.subview 切回原始逻辑形状。

这样,底层分配的内存布局满足硬件对齐约束,而上层算子仍按原始逻辑形状访问数据。

Pass Pipeline

三个Pass在Pipeline中按以下顺序执行:

hivm-align-alloc-size          // 调整特殊算子的alloc大小
  └─ hivm-pre-mark-stride-align   // [仅reg-based] 预分析:标记需跳过对齐的alloc
  └─ hivm-mark-stride-align       // 标记阶段:为算子操作数添加stride_align注解
hivm-enable-stride-align        // 使能阶段:传播注解 + 重分配内存
  • 标记阶段在FuncOp级别运行(func::FuncOp),可并行处理不同函数。

  • 使能阶段在ModuleOp级别运行(mlir::ModuleOp),因为对齐传播可能跨 func.call 边界影响Vector Function被调用者,FuncOp并行嵌套会引发数据竞争。

hivm-pre-mark-stride-align

功能概述

该Pass仅对reg-based架构(如Ascend 950PR/Ascend 950DT)生效,是标记阶段的前置分析步骤。

某些Buffer由DMA(hivm.hir.load)从GM搬运到UB后,会在Vector Function(通过 func.call 调用的VF)中被 vector.transfer_read(即vload)读取。如果对这些Buffer做stride对齐,会在各行之间插入填充间隙,使UB中的memref布局变为非连续。而 vlds 指令要求连续元素的扁平指针加载,无法处理非连续布局。

因此,PreMarkStrideAlign提前检测这类Buffer,在其根alloc上添加 hivm.skip_stride_align_for_vload 注解,供后续MarkStrideAlign跳过。

核心流程

  1. 遍历函数内的所有 hivm.hir.load(DMA加载)操作。

  2. 对每个load的结果Buffer,通过数据流追踪(经过 memref.subviewmemref.cast 等view-like op,以及 func.call 调用边界),判断其是否最终流向某个VF被调用者中的 vector.transfer_read

  3. 若确认流向vload,则回溯到根alloc(memref.alloc),在其后插入 annotation.mark 并设置 hivm.skip_stride_align_for_vload 属性。

注:当VF被调用者仅有声明(declaration,无函数体,如SplitSimtModule后的情况)时,无法判断参数是否被vload,因此跳过该Buffer的标记。

hivm-mark-stride-align

功能概述

该Pass遍历所有HIVM结构化算子(HIVMStructuredOp),自动分析其memref操作数的布局,找出最后不连续维度,并为UB空间的操作数创建 annotation.mark 注解,记录对齐维度和对齐字节数。

核心流程

1. 通用结构化算子标记

对每个 HIVMStructuredOp,执行以下步骤:

  1. 跳过非Buffer化(tensor类型)或仅含GM操作数的算子(GM无需对齐)。

  2. 收集算子的memref操作数类型,跳过全rank-0或全shape为1的操作数。

  3. 判断是否为UB DMA操作(hivm.hir.load/store),DMA操作的对齐判断规则与普通计算op不同。

  4. 根据架构走不同分析路径:

    • reg-based(A5/950):由于A5的 hfusion-flatten 已合并轴,直接将memref视为已flatten的状态,用 getLastDiscontinuousDimRegBased 查找最后不连续维度。对Fixpipe有专门的对齐约束计算。

    • 非reg-based(A2/A3):通过 FlattenInterfacegetFlattened 获取flatten后的关联组和类型,用 getLastDiscontinuousDim 在flatten后的类型上查找,再映射回原始维度。

  5. 取UB空间的操作数(getTargetSpaceOperands(UB)),调用 markAlignedDim 为每个操作数创建对齐注解。

2. 最后不连续维度查找

getLastDiscontinuousDim 的查找逻辑:

  • 如果任何memref的最低维stride不为1(最低维不连续),则最低维就是对齐目标维度。

  • 如果最低维stride为1(最低维连续),则在更高维中查找最后一个stride≠size的维度。

  • 对UB DMA操作(copy_gm_to_ub等),在Ascend 950PR/Ascend 950DT上只考虑最低维的stride,如果次低维已对齐且最低维无tail-jump问题,则无需对齐。

  • 对1D memref,不存在不连续维度,返回 nullopt

3. Fixpipe特殊处理

Fixpipe(L0C→UB搬运)的对齐约束更为复杂,取决于:

参数

说明

DataWidthType

数据位宽:B4(4bit)、B8、B16、B32

ChannelSplit

是否通道分割(CS_N/CS_Y)

NZ2ND

是否NZ→ND转换

NZ2DN

是否NZ→DN转换

LoopEnhance

是否循环增强

DualDstMode

双目的模式:N、SplitN、SplitM

N轴对齐要求(部分示例):

数据位宽

NZ2ND

对齐元素数

B16

ND_N

16

B8

ND_N

16

B8

ND_Y

32

B4

ND_Y

64

B32

ND_N

16

B32

ND_Y

8

M轴对齐要求(部分示例):

数据位宽

NZ2DN

对齐元素数

B16

DN_Y

16

B8

DN_Y

32

B4

DN_Y

64

B32

DN_N(SplitM)

2

B32

DN_Y

8

此外,对3D-dot的Fixpipe,还需对第0维(tiling维度)进行对齐。

4. vsstb Bank Conflict避免

当VF函数中存在 vector.transfer_write 且适合stride store(vsstb)时,MarkStrideAlign会对 func.call 的对应参数额外添加更高的对齐要求,以避免Bank Conflict。对齐大小计算为:

alignSize = shape[subTailDim] * 16 / gcd(shape[subTailDim], 16) + 1

其中 subTailDim = rank - 2(倒数第二维),16 = 512字节 / 32字节(对齐粒度与最低维size的比值)。最终对齐字节数需乘以数据类型宽度。

5. rank-reducing subview修复

当Store操作的UB操作数经过trailing-unit rank-reduced subview(如 memref<3x1x1xi32>memref<1x1xi32>)时,对齐维度可能落在被丢弃的维度上。此时需要将标记目标从subview结果回溯到subview的source,并调整对齐维度索引。

6. Vector Function参数标记(仅reg-based)

func.call 的参数,如果尚未被标记且属于本地Buffer,则走与结构化算子相同的最后不连续维度分析,为参数添加对齐注解。同时跳过被PreMarkStrideAlign标记为 SkipStrideAlignForVLoad 的参数。

标记效果示例

// 标记前:UB alloc的最低维不连续
%alloc = memref.alloc() : memref<37x5x3xi32, #hivm.address_space<ub>>
hivm.hir.load ins(%gm : ...) outs(%alloc : memref<37x5x3xi32, strided<[15, 3, 1]>, ...>)

// 标记后:在第2维(stride=15)添加32字节对齐注解
annotation.mark %alloc {hivm.stride_align_dims = array<i32: 1>, hivm.stride_align_value_in_byte = array<i32: 32>}

hivm-enable-stride-align

功能概述

该Pass读取标记阶段产生的 stride_align 注解,通过多轮传播将注解汇总到根分配点,然后对 memref.alloc/memref.alloca 进行形状填充(padding),使各行起始地址满足硬件对齐要求,最后通过 memref.subview 切回原始逻辑形状。

核心流程

  1. 规范化对齐信息(NormalizeAlignInfoPattern) 对每个 annotation.mark 上的 stride_align_dimsstride_align_value_in_byte 属性进行排序,确保维度索引按升序排列,便于后续传播时的并集和比较操作。

  2. 向上传播到根Alloc(PropagateAlignUpToRootAllocationPattern) 将落在 memref.subviewmemref.castmemref.collapse_shapememref.expand_shapememref.reshapememref.view 等view-like op结果上的对齐标记,沿着定义链向上回溯到根 memref.alloc/memref.alloca,并将对齐信息附加到根alloc上。

    约束:memref不应由SCF op返回(如 scf.if yield memref),这类控制流应在标记阶段前由其他pass下沉。

  3. 向下传播到叶子操作数(PropagateAlignDownToLeafOperandsPattern) 将根alloc上的对齐信息向下传播到所有直接或间接使用该alloc的叶子操作数。同时为缺少标记的alloc补充创建 annotation.markAddAlignAnnotationMarkForAlloc)。

  4. 算子操作数间传播(PropagateAlignAmongOperationOperands) 对每个HIVM Vector算子和 hivm.hir.copy,收集其UB空间非临时操作数的对齐信息,取并集后传播到未标记或标记不完全的操作数。例如:

    // op有三个UB操作数A, B, C,A和B有对齐标记
    // 传播后C也获得相同的对齐标记
    %A = ... {stride_align_dims=[1], value=[32]}
    %B = ... {stride_align_dims=[1], value=[32]}
    hivm.hir.vadd ins(%A, %B) outs(%C : ...)  // C获得 {stride_align_dims=[1], value=[32]}
    

    Copy操作的第一操作数(源)不参与传播,仅传播到目的操作数。

  5. 迭代至不动点 步骤3→4→2反复迭代,直到对齐信息不再变化或达到最大迭代次数(10次)。若超过最大迭代仍未收敛,报错退出。

  6. 清除注解(RemoveAlignMarkPattern) 传播完成后,移除所有 annotation.mark 上的 stride_align_dimsstride_align_value_in_byte 属性,保留空的mark op(或由后续pass清理)。该步骤只清除挂在 annotation.mark 上的传播中间信息,不会清除已经传播到 memref.alloc/memref.alloca 上的对齐属性;后续 EnableAlignAllocation 仍依赖alloc上的对齐属性完成内存重分配。

  7. 使能对齐分配(EnableAlignAllocation) 对每个带有 stride_align_dims 属性的 memref.alloc/memref.alloca

    • 7.1 收集对齐单元(collectAlignUnits):将对齐字节数转换为对齐元素数(除以元素位宽),考虑架构差异。

    • 7.2 计算对齐后形状(calculateAlignedShape):对每个需要对齐的维度,将其size向上取整到对齐单元的倍数(AlignUp)。同时计算原始逻辑形状(sub_shape)。

    • 7.3 创建新的alloc,类型为对齐后的形状和identity layout。

    • 7.4 创建 memref.subview,从对齐后的alloc切出原始逻辑形状。

    • 7.5 调用 replaceAndPropagateMemRefType 将原alloc的所有使用替换为新subview,并传播memref类型变化。

    • 7.6 同步更新 buffer_size_in_byte 注解(按对齐后与对齐前的静态体积比缩放)。

    • 7.7 删除原alloc。

    若对齐后形状与原始形状相同(无需padding),仅移除对齐属性,不创建新alloc。

  8. AIC→AIV对齐传播(仅Ascend 950PR/Ascend 950DT) 对于TightlyCoupledBuffer(紧密耦合Buffer),AIC和AIV共享同一物理内存。如果AIC侧的alloc做了stride对齐,AIV侧的对应alloc也需要相同对齐。AddStrideAlignInfoForAiv 通过TightlyCoupledBuffer的ID匹配AIC和AIV的alloc,将对齐信息从AIC传播到AIV。

  9. 后处理

    • handlePropagateFailure:处理传播失败(unrealized_conversion_cast)的op。

    • materializeRemainingStaticUBLayoutCasts(仅reg-based):将stride-align产生的静态UB layout-changing cast物化。

    • 设置 hivm.storage_aligned 属性(仅非reg-based):标记函数已完成存储对齐。

使能效果示例

// 使能前:alloc带对齐标记
%alloc = memref.alloc() {hivm.stride_align_dims = array<i32: 1>, hivm.stride_align_value_in_byte = array<i32: 32>} : memref<37x5x3xi32, #hivm.address_space<ub>>
// 第1维size=3,stride需要32字节对齐 → 32/4(i32)=8元素 → 最低维填充到8

// 使能后:分配对齐后的形状,再subview切回原始形状
%aligned_alloc = memref.alloc() : memref<37x5x8xi32, #hivm.address_space<ub>>
%subview = memref.subview %aligned_alloc[0, 0, 0] [37, 5, 3] [40, 8, 1] : memref<37x5x8xi32, ...> to memref<37x5x3xi32, strided<[40, 8, 1]>, ...>
// stride从[15,3,1]变为[40,8,1],各行起始地址满足32字节对齐

编译选项

选项

默认值

说明

enable-hivm-auto-storage-align

true

Pipeline选项,控制是否在 alignStoragePipeline 中执行标记和使能阶段

三个Pass本身不暴露独立的命令行选项,均通过 enableHIVMAutoStorageAlign 选项统一控制。关闭该选项后,hivm-pre-mark-stride-alignhivm-mark-stride-align 不执行,但 hivm-enable-stride-align 仍会执行(此时因无标记注解,不会产生实际的内存重分配)。

使用约束

  1. 需在Buffer化之后运行:标记阶段要求算子已通过 hasPureBufferSemantics() 检查,即输入IR已完成tensor→memref的Bufferization。

  2. 需在Flatten之后运行:reg-based架构要求 hfusion-flatten 已完成轴合并;非reg-based架构通过 FlattenInterface::getFlattened 在标记时获取flatten信息。

  3. memref不应由SCF op返回:如 scf.if yield memref 的场景,对齐传播无法处理,需在标记阶段前由其他pass将控制流下沉。

  4. 使能阶段可能增大内存占用:对齐填充会使alloc的物理尺寸增大(对齐维度size变大),可能影响后续PlanMemory的内存分配。用户需保证对齐后的总内存需求不超过硬件空间上限。

  5. 传播收敛性:操作数间传播最大迭代10次,若未收敛会编译失败。通常合理的IR结构不会触发此限制。

  6. vsstb与FlattenOps的已知冲突:vsstb的bank conflict对齐可能将连续形状变为非连续形状,与依赖连续形状特征的 memref.collapse_shape/memref.reshape 产生冲突,可能导致后续FlattenOps编译错误。