CANN opbase 中 GetUbBlockSize 接口解析:获取 UB 块单元大小(32 Bytes)的源码级详解
CANN opbase 中 GetUbBlockSize 接口解析获取 UB 块单元大小32 Bytes的源码级详解【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase本文围绕 op_common 平台工具接口GetUbBlockSize展开介绍它的功能定位、函数原型、参数与返回值约定并结合 CANN/opbase 仓库中的真实源码说明该接口在 Tiling 编译信息CompileInfo收集链路与 ATVOSS 公共 Tiling 模板reduce、broadcast 等中的具体用法帮助开发者在自研算子 Tiling 开发中正确获取并运用 UB block 对齐单位。1. 接口功能定位在 AscendC 算子的 Tiling 开发中UBUnified Buffer统一缓冲区是 AI Core 片上高速存储。UB 的数据搬运和内存访问以“块block”为基本对齐/划分的参考单位GetUbBlockSize的作用正是获取平台 UB 的 block 单元大小单位为字节bytes。该接口属于Ops::Base命名空间下的平台信息获取接口族与 GetUbSize、GetCacheLineSize、GetVRegSize 等接口并列完整清单可参考 platform 文档页英文对照版见 GetUbBlockSize。理解它的最佳方式是把它看作一组“平台硬件常量查询 API”中的一员开发者在 Tiling 侧写对齐、切块、计算 UB 空间开销时不直接硬编码32而是统一调用该接口取值保证与平台约定一致、便于后续平台演进的统一替换。2. 函数原型、参数与返回值官方文档给出的函数原型如下见 GetUbBlockSizetemplate typename T uint32_t GetUbBlockSize([[maybe_unused]] const T *context)参数说明参数名输入/输出说明context输入Tiling 的上下文信息类型为TilingContext/TilingParseContext。返回值说明uint32_t返回平台 UB 的 block 单元大小当前为 32 bytes。约束说明无。2.1 Host 侧实现Host 侧实现在 platform_util.h/** * Get the block size of unified buffer in bytes */ template typename T uint32_t GetUbBlockSize([[maybe_unused]] const T* context) { return 32U; // will using AscendC api later }从源码结构看当前实现直接返回编译期常量32U参数标记为[[maybe_unused]]表明 context 暂时不参与取值注释// will using AscendC api later暗示后续版本可能改为通过PlatformAscendC从平台描述中动态查询接口签名则会保持稳定——这正是使用模板 context 参数的原因与同文件中GetUbSize、GetAivCoreNum等真正依赖context-GetPlatformInfo()的接口保持统一的调用形态对比 GetUbSize 的实现它通过platform_ascendc::PlatformAscendC读取CoreMemType::UB的实际大小。需要注意一个与同文件其他接口的差异GetAivCoreNum、GetUbSize、GetWorkspaceSize等在实现里带有static_assert约束T必须是gert::TilingParseContext或gert::TilingContext而GetUbBlockSize当前实现不做该静态断言因为取值与具体 context 无关。2.2 Kernel 侧实现与 Host 侧对应Kernel 侧__aicore__编译单元提供无参版本定义在 platform_util.h/** * Get the block size of unified buffer in bytes */ __aicore__ inline constexpr uint32_t GetUbBlockSize() { return 32U; }该版本是constexpr内联函数可在 Kernel 侧直接用于static/constexpr常量表达式例如 ATVOSS Tiling 模板中即以常量形式引用它见第 4 节。这说明同一个语义在 Host/Kernel 两个编译域各有一套实现但当前返回值一致32 bytes开发者在 Tiling 代码中按所在编译环境包含对应的头文件即可。3. 典型调用示例与校验链路官方给出的关键调用示例如下仅供参考不支持直接拷贝运行compileInfo-ubBlockSize GetUbBlockSize(context_);这一示例并非示意性写法它几乎就是仓库内 ATVOSS reduce 公共 Tiling 框架的真实代码。在 reduce_tiling.cpp 的编译信息采集阶段compileInfo_-ubBlockSize GetUbBlockSize(context_); OP_CHECK_IF(compileInfo_-ubBlockSize 0, OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_-GetNodeName(), ubBlockSize, std::to_string(compileInfo_-ubBlockSize).c_str(), The value of ub block size must be greater than 0), return ge::GRAPH_FAILED);可以看到框架对返回值做了防御性校验取值为 0 时记录错误日志并返回ge::GRAPH_FAILED与相邻的ubSize、cacheLineSize、vRegSize等字段的校验模式完全一致reduce_tiling.cpp 中依次调用GetAivCoreNum、GetUbSize、GetCacheLineSize、GetUbBlockSize、GetVRegSize填充同一份compileInfo_。该结构成员定义在 reduce_tiling.huint64_t ubBlockSize 0;自研算子在 Tiling 阶段如果也需要 UB block 对齐能力参考这一“取值 非零校验 存入 CompileInfo”的模式即可与平台接口族保持一致的健壮性。4. 在公共 Tiling 模板中的真实用途ubBlockSize取到值之后其核心价值体现在元素数口径的对齐换算上GetUbBlockSize返回的是字节数32B而 Tiling 中处理的是元素数因此需要除以元素大小dSize得到“每个 block 可容纳的元素数”再配合CeilAlign/FloorAlign做上下取整。仓库中这一换算出现得非常密集例如 reduce_tiling.cpp 中多处形如uint64_t ubBlockSize compileInfo_-ubBlockSize / dSize; ... r r * CeilAlign(shape[i], ubBlockSize); // 按 block 向上取整计算切块数 ... step FloorAlign(step, ubBlockSize); // 步长按 block 向下取整参见 reduce_tiling.cpp 中CalculateUBBlockNum、AdjustStepToDivisor等函数以及 reduce_tiling_batch_invariant.h 中的同类换算另有按最小输入元素字节数归一化的ubBlockSize compileInfo_-ubBlockSize / opDag_.minInputBytes口径见 reduce_tiling.cpp。在 Kernel 侧调度模板中该常量则直接以编译期形式参与 UB 地址/划分的计算例如reduce_sch.hconstexpr static uint64_t UB_BLOCK Ops::Base::GetUbBlockSize();reduce_sch_aux_base.h同样以UB_BLOCK常量形式引用reduce_operator.hconstexpr uint16_t BLOCK_SIZE Ops::Base::GetUbBlockSize();reduce_tensor_empty.hconstexpr uint32_t UB_SIZE Ops::Base::GetUbBlockSize();broadcast 非连续场景 broadcast_sch_nlast_transpose_ub_broadcast.huint32_t blockSize Ops::Base::GetUbBlockSize();这些调用点集中体现了两点工程收益一是对齐单位收敛到单一接口避免各 Tiling 模板各自硬编码魔数二是constexpr语义使其能嵌入编译期常量表达式减少 Kernel 侧运行时开销。从源码结构看只要后续平台 block 粒度发生变化仅需修改GetUbBlockSize的实现所有引用该接口的公共 Tiling 模板即可自动适配。5. 使用建议与边界说明结合上文源码证据给出以下实践要点返回值约定当前平台实现固定返回32U32 bytes。Host 侧模板实现带// will using AscendC api later注释表明存在向平台动态查询演进的可能因此不建议在业务代码中绕过接口直接写32。字节到元素的换算将ubBlockSize用于元素维度对齐时必须先除以元素字节数dSize或minInputBytes口径并选对CeilAlign向上取整用于估算占用与FloorAlign向下取整用于步长/块数的语义方向可参考 reduce_tiling.cpp 的既有实现。Host/Kernel 两个版本Tiling 侧Hostgert::TilingContext/gert::TilingParseContext可用使用带 context 的模板版本op_host/util/platform_util.hKernel 侧使用__aicore__无参constexpr版本op_kernel/platform_util.h。健壮性参照 ATVOSS 的写法取值后做非零校验reduce_tiling.cpp失败路径统一走错误码返回便于问题定位。适用前提以上结论基于当前仓库版本代码32 bytes为当前取值事实若切换平台/架构应以实际运行环境对应头文件中的实现为准。综上GetUbBlockSize虽然是一个仅返回常量的小接口但它是 opbase 平台信息接口族中服务于 Tiling 对齐计算的“基准刻度”。掌握其取值口径字节、Host/Kernel 双实现形态以及在 reduce/broadcast 公共 Tiling 模板中的换算与对齐用法是开展基于 op_common 的自研算子 Tiling 开发的基础前置知识。【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考