自动块化¶
背景¶
Auto Blockify Pass通过高效地将逻辑块映射到硬件物理块，是昇腾兼容算子执行链路的核心优化手段。当前架构下调度效率直接决定算子性能，逻辑块与物理块一对一映射可消除调度开销，实现性能提升。
结合AscendNPU IR架构使用实践，可用物理块数量通常远少于计算所需的逻辑块数量（物理块 < 50，逻辑块可能达到500+）。在这种10倍差距的场景下，加速效果可超过原始速度的两倍。
在运行Triton内核（通过triton-ascend）时，激活Auto Blockify逻辑的方式是添加以下标志：TRITON_ALL_PARALLEL。
对于AscendNPU IR开发者，可在bishengir-compile命令中添加以下标志：--enable-auto-blockify-loop。
算法原理¶
Auto Blockify Pass（AutoBlockifyParallelLoop）通过引入额外的循环层来变换IR，具体逻辑如下：
for outer from 0,...,ceildiv(logical_block_dim, physical_block_dim)
for inner from 0,...,physical_block_dim <- 作为 block.idx 使用
use(min(outer * physical_block_dim + inner, logical_block_dim))
逻辑说明：
原始调度
原始模式通常如下所示：
block.idx = hivm.get_block_idx use(block.idx) -------等价于-------------- for block.idx from 0,...,logical_block_num use(block.idx)使用
TRITON_ALL_PARALLEL的示例当用户在triton adapter中添加
TRITON_ALL_PARALLEL标志时，内核将被限制为仅使用最大物理块数量启动（假设逻辑块数 > 物理块数）。因此执行被限制为：for block.idx from 0,...,physical_block_num <- 来自 get_block_idx use(block.idx)仅依靠该循环逻辑无法覆盖全部计算索引，存在索引缺失问题。这也是引入Auto Blockify Pass补齐逻辑的原因：通过自动添加一层外部循环/块化来完善。
注：若不通过triton adapter接入，需要自行确保块维度的设置与上述一致。
使用Auto Blockify后的最终逻辑
经Auto Blockify Pass自动补全循环结构，最终执行逻辑如下：
for outer from 0,...,ceildiv(logical_block_dim, physical_block_dim) for inner from 0,...,physical_block_dim <- 作为 block.idx 使用 use(min(outer * physical_block_dim + inner, logical_block_dim))
接口说明：
该功能通过bishengir-compile中的--enable-auto-blockify-loop标志控制，也可通过bishengir-opt的--auto-blockify-parallel-loop标志直接调用。
为正确使用此功能，需注意以下几点：
Pass获取逻辑块数量的方式是查找标有
kLogicalBlockNumAttr属性（IR中为logical_block_num）的值，用户需确保该值可用，否则Pass调用时将失败。Pass还需要找到一个
hivm get_block_idx操作，该操作返回从0到块维度的块索引。使用AutoBlockify时，用户需要在调用设备内核时修改块维度（以最大物理块维度启动，与上述算法一致），使得blockidx操作返回0到physical_block_num范围内的值。
Triton Adapter：
该Pass已在triton adapter流水线中广泛使用。在此情况下正确使用AutoBlockify特性的方式是从前端（triton）通过TRITON_ALL_PARALLEL=1启用，该环境变量会同时完成准备工作（锁定块数量），然后自动以正确的标志调用相应的编译器命令。在triton流水线中有一个名为TritonGlobalKernelArgsToHIVMOpPass的Pass，会自动确保存在标记了logical_block_num的值，并创建所需的get_block_idx操作。
输入示例：
module attributes {dlti.target_system_spec = #dlti.target_system_spec<"NPU" : #hacc.target_device_spec<#dlti.dl_entry<"AI_CORE_COUNT", 20 : i32>, #dlti.dl_entry<"CUBE_CORE_COUNT", 20 : i32>, #dlti.dl_entry<"VECTOR_CORE_COUNT", 40 : i32>, #dlti.dl_entry<"UB_SIZE", 1572864 : i32>, #dlti.dl_entry<"L1_SIZE", 4194304 : i32>, #dlti.dl_entry<"L0A_SIZE", 524288 : i32>, #dlti.dl_entry<"L0B_SIZE", 524288 : i32>, #dlti.dl_entry<"L0C_SIZE", 1048576 : i32>, #dlti.dl_entry<"UB_ALIGN_SIZE", 256 : i32>, #dlti.dl_entry<"L1_ALIGN_SIZE", 256 : i32>, #dlti.dl_entry<"L0C_ALIGN_SIZE", 4096 : i32>>>, hivm.module_core_type = #hivm.module_core_type<AIV>} {
func.func @add_kernel(%arg0: i64 {hacc.arg_type = #hacc.arg_type<ffts_base_address>}, %arg1: memref<?xi8> {hacc.arg_type = #hacc.arg_type<workspace>}, %arg2: memref<?xf32> {tt.divisibility = 16 : i32}, %arg3: memref<?xf32> {tt.divisibility = 16 : i32}, %arg4: memref<?xf32> {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32, %arg7: i32, %arg8: i32) attributes {WorkspaceArgIdx = 0 : i64, func_dyn_memref_args = dense<[false, true, true, true, true, false, false, false, false]> : vector<9xi1>, hacc.entry, hacc.function_kind = #hacc.function_kind<DEVICE>, hivm.func_core_type = #hivm.func_core_type<AIV>} {
%c1024 = arith.constant 1024 : index
%c1024_i32 = arith.constant 1024 : i32
%c0 = arith.constant 0 : index
hivm.hir.set_mask_norm
%0 = arith.muli %arg6, %arg7 : i32
%1 = arith.muli %0, %arg8 : i32
annotation.mark %1 {logical_block_num} : i32 // 此 logical_block_num 是原始的大数值
%2 = hivm.hir.get_block_idx -> i64 // for block.idx from 0,...,block_num
%3 = arith.trunci %2 : i64 to i32
%4 = arith.muli %arg8, %arg7 : i32
%5 = arith.divsi %3, %4 : i32
%6 = arith.remsi %5, %arg6 : i32
%7 = arith.muli %6, %c1024_i32 : i32
%8 = arith.index_cast %7 : i32 to index
%reinterpret_cast = memref.reinterpret_cast %arg2 to offset: [%8], sizes: [1024], strides: [1] : memref<?xf32> to memref<1024xf32, strided<[1], offset: ?>>
%alloc = memref.alloc() : memref<1024xf32>
%9 = arith.addi %8, %c1024 : index
%10 = arith.index_cast %arg5 : i32 to index
%11 = arith.maxsi %8, %10 : index
%12 = arith.minsi %9, %11 : index
%13 = arith.subi %12, %8 : index
%subview = memref.subview %reinterpret_cast[0] [%13] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
%subview_0 = memref.subview %alloc[0] [%13] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
hivm.hir.load ins(%subview : memref<?xf32, strided<[1], offset: ?>>) outs(%subview_0 : memref<?xf32, strided<[1]>>) left_padding_num = %c0 : index init_out_buffer = false
%14 = bufferization.to_tensor %alloc restrict writable : memref<1024xf32>
%reinterpret_cast_1 = memref.reinterpret_cast %arg3 to offset: [%8], sizes: [1024], strides: [1] : memref<?xf32> to memref<1024xf32, strided<[1], offset: ?>>
%alloc_2 = memref.alloc() : memref<1024xf32>
%subview_3 = memref.subview %reinterpret_cast_1[0] [%13] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
%subview_4 = memref.subview %alloc_2[0] [%13] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
hivm.hir.load ins(%subview_3 : memref<?xf32, strided<[1], offset: ?>>) outs(%subview_4 : memref<?xf32, strided<[1]>>) left_padding_num = %c0 : index init_out_buffer = false
%15 = bufferization.to_tensor %alloc_2 restrict writable : memref<1024xf32>
%16 = tensor.empty() : tensor<1024xf32>
%17 = hivm.hir.vadd ins(%14, %15 : tensor<1024xf32>, tensor<1024xf32>) outs(%16 : tensor<1024xf32>) -> tensor<1024xf32>
%reinterpret_cast_5 = memref.reinterpret_cast %arg4 to offset: [%8], sizes: [1024], strides: [1] : memref<?xf32> to memref<1024xf32, strided<[1], offset: ?>>
%extracted_slice = tensor.extract_slice %17[0] [%13] [1] : tensor<1024xf32> to tensor<?xf32>
%subview_6 = memref.subview %reinterpret_cast_5[0] [%13] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
hivm.hir.store ins(%extracted_slice : tensor<?xf32>) outs(%subview_6 : memref<?xf32, strided<[1], offset: ?>>)
return
}
}
输出示例：
module attributes {dlti.target_system_spec = #dlti.target_system_spec<"NPU" : #hacc.target_device_spec<#dlti.dl_entry<"AI_CORE_COUNT", 20 : i32>, #dlti.dl_entry<"CUBE_CORE_COUNT", 20 : i32>, #dlti.dl_entry<"VECTOR_CORE_COUNT", 40 : i32>, #dlti.dl_entry<"UB_SIZE", 1572864 : i32>, #dlti.dl_entry<"L1_SIZE", 4194304 : i32>, #dlti.dl_entry<"L0A_SIZE", 524288 : i32>, #dlti.dl_entry<"L0B_SIZE", 524288 : i32>, #dlti.dl_entry<"L0C_SIZE", 1048576 : i32>, #dlti.dl_entry<"UB_ALIGN_SIZE", 256 : i32>, #dlti.dl_entry<"L1_ALIGN_SIZE", 256 : i32>, #dlti.dl_entry<"L0C_ALIGN_SIZE", 4096 : i32>>>, hivm.module_core_type = #hivm.module_core_type<AIV>} {
func.func @add_kernel(%arg0: i64 {hacc.arg_type = #hacc.arg_type<ffts_base_address>}, %arg1: memref<?xi8> {hacc.arg_type = #hacc.arg_type<workspace>}, %arg2: memref<?xf32> {tt.divisibility = 16 : i32}, %arg3: memref<?xf32> {tt.divisibility = 16 : i32}, %arg4: memref<?xf32> {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32, %arg7: i32, %arg8: i32) attributes {WorkspaceArgIdx = 0 : i64, func_dyn_memref_args = dense<[false, true, true, true, true, false, false, false, false]> : vector<9xi1>, hacc.entry, hacc.function_kind = #hacc.function_kind<DEVICE>, hivm.func_core_type = #hivm.func_core_type<AIV>} {
%0 = arith.muli %arg6, %arg7 : i32
%1 = arith.muli %0, %arg8 : i32
annotation.mark %1 {logical_block_num} : i32 // 此 logical_block_num 是原始的大数值
%c0_i32 = arith.constant 0 : i32
%c40_i32 = arith.constant 40 : i32 // 40 为此处的物理块数
%2 = arith.ceildivsi %1, %c40_i32 : i32 // ceildiv(logical_block_num, physical_block_dim)
%c1_i32 = arith.constant 1 : i32
scf.for %arg9 = %c0_i32 to %2 step %c1_i32 : i32 { // 外层循环
%c1024 = arith.constant 1024 : index
%c1024_i32 = arith.constant 1024 : i32
%c0 = arith.constant 0 : index
hivm.hir.set_mask_norm
%3 = hivm.hir.get_block_idx -> i64 // 内层循环（锁定到物理块）
%4 = arith.trunci %3 : i64 to i32
%5 = arith.muli %arg9, %c40_i32 : i32 // outer_i * physical_block_num
%6 = arith.addi %5, %4 : i32 // outer_i * physical_block_num + inner
%7 = arith.minsi %6, %1 : i32 // (min(outer*physical_block_dim + inner, logical_block_num))
%8 = arith.extsi %7 : i32 to i64
%9 = arith.trunci %8 : i64 to i32
%10 = arith.muli %arg8, %arg7 : i32
%11 = arith.divsi %9, %10 : i32
%12 = arith.remsi %11, %arg6 : i32
%13 = arith.muli %12, %c1024_i32 : i32
%14 = arith.index_cast %13 : i32 to index
%reinterpret_cast = memref.reinterpret_cast %arg2 to offset: [%14], sizes: [1024], strides: [1] : memref<?xf32> to memref<1024xf32, strided<[1], offset: ?>>
%alloc = memref.alloc() : memref<1024xf32>
%15 = arith.addi %14, %c1024 : index
%16 = arith.index_cast %arg5 : i32 to index
%17 = arith.maxsi %14, %16 : index
%18 = arith.minsi %15, %17 : index
%19 = arith.subi %18, %14 : index
%subview = memref.subview %reinterpret_cast[0] [%19] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
%subview_0 = memref.subview %alloc[0] [%19] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
hivm.hir.load ins(%subview : memref<?xf32, strided<[1], offset: ?>>) outs(%subview_0 : memref<?xf32, strided<[1]>>) left_padding_num = %c0 : index
%20 = bufferization.to_tensor %alloc restrict writable : memref<1024xf32>
%reinterpret_cast_1 = memref.reinterpret_cast %arg3 to offset: [%14], sizes: [1024], strides: [1] : memref<?xf32> to memref<1024xf32, strided<[1], offset: ?>>
%alloc_2 = memref.alloc() : memref<1024xf32>
%subview_3 = memref.subview %reinterpret_cast_1[0] [%19] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
%subview_4 = memref.subview %alloc_2[0] [%19] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
hivm.hir.load ins(%subview_3 : memref<?xf32, strided<[1], offset: ?>>) outs(%subview_4 : memref<?xf32, strided<[1]>>) left_padding_num = %c0 : index
%21 = bufferization.to_tensor %alloc_2 restrict writable : memref<1024xf32>
%22 = tensor.empty() : tensor<1024xf32>
%23 = hivm.hir.vadd ins(%20, %21 : tensor<1024xf32>, tensor<1024xf32>) outs(%22 : tensor<1024xf32>) -> tensor<1024xf32>
%reinterpret_cast_5 = memref.reinterpret_cast %arg4 to offset: [%14], sizes: [1024], strides: [1] : memref<?xf32> to memref<1024xf32, strided<[1], offset: ?>>
%extracted_slice = tensor.extract_slice %23[0] [%19] [1] : tensor<1024xf32> to tensor<?xf32>
%subview_6 = memref.subview %reinterpret_cast_5[0] [%19] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
hivm.hir.store ins(%extracted_slice : tensor<?xf32>) outs(%subview_6 : memref<?xf32, strided<[1], offset: ?>>)
}
return
}
}
约束¶
可并行性：Auto Blockify算法仅适用于完全可并行的代码。这意味着各逻辑块的计算与访问必须可以安全地并行执行，块间不存在依赖关系。
使用场景：若逻辑块数量非常小，则此Pass不会带来任何优势。