CANN ops-math stateless_bernoulli 算子详解伯努利分布随机采样原理、参数与 aclnn 调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 random/stateless_bernoulli/README.md 展开系统讲解 CANN ops-math 开源算子库中stateless_bernoulli算子的功能语义、参数规格、产品支持矩阵以及基于 aclnn 两段式接口的完整调用方法。读者学完后将掌握在 Ascend 950 平台上从伯努利分布高效生成 0/1 随机张量的实战能力并能理解该算子从算子定义、Shape 推导、Tiling 到 Philox 随机数内核的完整实现链路。功能与数学原理stateless_bernoulli是一个无状态stateless随机采样算子其核心功能是从伯努利分布中提取二进制随机数0 或 1。其中prob为生成 1 的概率输入张量shape用于指定输出随机数的形状。计算公式为$$ out \sim Bernoulli(prob) $$即输出张量中的每个元素独立地以概率prob取值为 1以概率1 - prob取值为 0。这一操作在神经网络中常用于 dropout 掩码生成、蒙特卡洛采样、数据增强等需要按概率保留/丢弃的场景。需要说明的是该算子属于无状态随机算子随机数生成完全由用户显式传入的seed随机种子与offset偏移量驱动不依赖算子内部的隐藏状态因此同一组(seed, offset, shape, prob)可以保证结果可复现便于调试与测试。产品支持情况根据 README.md 的产品支持表该算子当前的支持矩阵如下产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从源码可以印证这一支持范围在 算子定义文件 中AICore 内核仅注册了ascend950与ascend350两种芯片配置分别对应 Ascend 950 系列与 arch35 架构Tiling 与 Kernel 实现也统一放在arch35/目录下参见 stateless_bernoulli_tiling_arch35.cpp 与 stateless_bernoulli_simt.h。参数说明stateless_bernoulli共包含 5 个参数均为 ND 数据格式。完整规格如下参数名输入/输出/属性描述数据类型数据格式shape输入随机数的 shape。INT32、INT64NDprob输入伯努利分布的保活系数取 1 的概率满足 0 ≤ prob ≤ 1。FLOAT16、BF16、FLOATNDseed输入随机种子驱动随机数生成。INT64NDoffset输入获取值的步长随机数偏移量。INT64NDy输出输出的 tensor。INT8、UINT8、INT16、UINT16、INT32、UINT32、INT64、UINT64、BOOL、FLOAT16、BFLOAT16、FLOATND以上规格与 stateless_bernoulli_def.cpp 中的算子定义完全一致shapeinputDataType {DT_INT32, DT_INT64}必选输入且标记为ValueDepend(OPTIONAL)即其取值会影响输出形状推导下文 Shape 推导一节会展开。probprobDataType {DT_FLOAT16, DT_FLOAT, DT_BF16}必选输入。seed、offsetbaseDataType {DT_INT64}其中offset为可选输入ParamType(OPTIONAL)。youtputDataType覆盖 12 种数据类型含整型、无符号整型、BOOL 与浮点类型。另有可选属性dtypeAttr(dtype).AttrType(OPTIONAL).Int()用于显式指定输出数据类型。关键取值约束prob 取值范围必须满足0 ≤ prob ≤ 1。该约束在 aclnn 第一段接口的入参校验中强制执行超出范围会报ACLNN_ERR_PARAM_INVALID。offset 取值约束必须满足offset % 4 0例如可取 0、4、8……不满足约束会调用失败。这一约束同时体现在两处源码中aclnn 接口文档见 aclnnBernoulli 文档与 Tiling 实现——stateless_bernoulli_tiling_arch35.cpp 中通过OFFSET_MULTIPLE 4检查 offset若offset % 4 ! 0则返回GRAPH_FAILED并输出offset value must be a multiple of 4的错误日志。其原因是底层 Philox 随机算法按 4 个随机数为一组STEP 4进行批量生成见下文内核实现。约束说明原文档约束说明为无。从 aclnn 接口层面看存在如下通用约束输入/输出 tensor 维度不超过 8 维0~8 维self与out的 shape、数据类型必须一致支持非连续非稠密tensor 输入。确定性计算aclnnBernoulli与aclnnInplaceBernoulli默认采用确定性实现相同(seed, offset)下输出可复现。调用说明aclnn 两段式接口stateless_bernoulli算子对外通过 aclnnAscend CANN 的算子单算子调用接口暴露且每个算子都采用两段式接口先调用xxxGetWorkspaceSize接口获取计算所需 workspace 大小与包含算子计算流程的执行器executor再调用xxx接口真正执行计算。README 中给出了两种调用方式调用方式样例代码说明aclnn 调用test_aclnn_bernoulli_tensor.cpp通过 aclnnBernoulliTensor / aclnnInplaceBernoulliTensor 接口调用prob以张量形式传入。aclnn 调用test_aclnn_bernoulli.cpp通过 aclnnBernoulli / aclnnInplaceBernoulli 接口调用prob以标量形式传入。两种方式的区别在于prob的传参形态aclnnBernoulli系列接受aclScalar*标量全张量共用同一个概率aclnnBernoulliTensor系列接受aclTensor*张量每个元素可对应独立的概率但张量尺寸需与输出一致或为单元素标量张量。函数原型aclnnBernoulli与aclnnInplaceBernoulli实现相同的功能区别在于前者需新建一个输出张量对象存储计算结果后者无需新建输出张量直接在输入张量的内存中覆盖写入计算结果inplace。aclnnStatus aclnnBernoulliGetWorkspaceSize( const aclTensor* self, const aclScalar* prob, int64_t seed, int64_t offset, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnBernoulli( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream) aclnnStatus aclnnInplaceBernoulliGetWorkspaceSize( const aclTensor* selfRef, const aclScalar* prob, int64_t seed, int64_t offset, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnInplaceBernoulli( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)aclnnBernoulliGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorselfaclTensor*输入输入 tensor。支持空 Tensor数据类型需与 out 一致shape 需与 out 一致。FLOAT16、FLOAT、DOUBLE、UINT8、INT8、INT16、INT32、INT64、BOOL、BFLOAT16ND0-8√probaclScalar*输入公式中的 prob。满足 0 ≤ prob ≤ 1。FLOAT16、FLOAT、DOUBLE、BFLOAT16---seedint64_t输入设置随机数生成器的种子。-INT64---offsetint64_t输入设置随机数偏移量。取值约束offset % 4 0例如 0、4、8……不满足会调用失败。INT64---outaclTensor*输出公式中的 out。支持空 Tensor数据类型需与 self 一致shape 需与 self 一致。FLOAT16、FLOAT、DOUBLE、UINT8、INT8、INT16、INT32、INT64、BOOL、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含算子计算流程。-----平台差异说明Atlas 训练系列产品910 平台上self/out数据类型不支持 BFLOAT16。从源码看op_api 层提供了 AICPU 内核作为非 AICore 平台的降级路径见下文双层内核选择具体支持范围请以实际 CANN 版本发布说明为准。返回值与错误码第一段接口GetWorkspaceSize完成入参校验出现以下场景时返回对应错误返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、prob 或 out 是空指针。ACLNN_ERR_PARAM_INVALID161002self 或 out 的数据类型或数据格式不在支持范围之内。ACLNN_ERR_PARAM_INVALID161002prob 的数据类型不在支持范围之内。ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型不一致。ACLNN_ERR_PARAM_INVALID161002prob 不满足 0 ≤ prob ≤ 1。ACLNN_ERR_PARAM_INVALID161002self 或 out 的维度大于 8 维。ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不一致。第二段接口aclnnBernoulli参数为 workspace、workspaceSize、executor、stream 四项前两者来自第一段接口的计算结果stream指定执行任务的 Stream。返回值同样为aclnnStatus具体返回码语义可参考仓库中 aclnn 返回码说明。aclnnInplaceBernoulli系列接口的参数与上述基本一致区别在于用selfRef输入/输出替代self与out计算结果直接写回selfRef指向的内存因此入参校验错误码表中仅涉及selfRef与prob两类参数。BernoulliTensor 系列prob 为张量当需要为输出张量的不同位置指定不同概率时使用 aclnnBernoulliTensor aclnnInplaceBernoulliTensor 接口其原型与标量版一一对应仅将aclScalar* prob替换为aclTensor* probaclnnStatus aclnnBernoulliTensorGetWorkspaceSize( const aclTensor *self, // 用于指定输出 out 的 shapeDevice 侧 aclTensor const aclTensor *prob, // 公式中的 probDevice 侧 aclTensor满足 0 ≤ prob ≤ 1 int64_t seed, int64_t offset, // offset % 4 0 aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)其入参校验规则与标量版类似self/prob/out的数据类型与数据格式须在支持范围内、self与out的 shape 和数据类型须一致、维度不超过 8错误码同样为ACLNN_ERR_PARAM_NULLPTR161001与ACLNN_ERR_PARAM_INVALID161002。平台差异方面Atlas 训练系列产品910不支持 UINT8、INT8、INT16、INT32、INT64、BOOL、BFLOAT16 等输出类型Atlas A2/A3 系列不支持 UINT8、INT8、INT16、INT32、INT64、BOOL。完整调用流程示例仓库在 examples/arch35/test_aclnn_bernoulli.cpp 中提供了可直接参考的完整示例prob1.0、shape[4,4]、seed0、offset0。核心调用链路可概括为七步以下为aclnnBernoulli关键片段#include acl/acl.h #include aclnnop/aclnn_bernoulli.h int main() { // 1. device/stream 初始化固定写法 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // aclInit / aclrtSetDevice / aclrtCreateStream // 2. 构造输入与输出 std::vectorint64_t selfShape {4, 4}; std::vectorint64_t outShape {4, 4}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclScalar* prob nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, ..., 15}; std::vectorfloat outHostData(16, 0); int64_t seed 0; int64_t offset 0; // 必须满足 offset % 4 0 float probValue 1.0f; // 满足 0 prob 1 CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, ACL_FLOAT, self); prob aclCreateScalar(probValue, aclDataType::ACL_FLOAT); CreateAclTensor(outHostData, outShape, outDeviceAddr, ACL_FLOAT, out); // 3. 调用两段式接口 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnBernoulliGetWorkspaceSize(self, prob, seed, offset, out, workspaceSize, executor); void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } ret aclnnBernoulli(workspaceAddr, workspaceSize, executor, stream); // 4. 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); // 5. 将 Device 侧结果拷贝回 Host ret aclrtMemcpy(resultData.data(), size * sizeof(float), outDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); // 6. 释放 aclTensor / aclScalar aclDestroyTensor(self); aclDestroyScalar(prob); aclDestroyTensor(out); // 7. 释放 device 资源 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }aclnnInplaceBernoulli的差异点在于不创建out张量将结果直接写回self内存接口参数为selfRef结果拷贝时也从selfDeviceAddr读取。test_aclnn_bernoulli_tensor.cpp中对应的 Tensor 版本示例仅需把aclScalar* prob替换为通过aclCreateTensor构造的aclTensor* probprob 张量尺寸为 1 或与输出一致。具体编译与执行过程请参考仓库的 编译与运行样例。源码级实现原理要深入理解该算子可以从以下五个层次阅读仓库源码构成一条定义 → 推导 → 切分 → 计算 → 调度的完整链路。1. 算子定义与注册Host 侧 OpDefstateless_bernoulli_def.cpp 通过OpDef类完成算子原型注册OP_ADD(StatelessBernoulli)借助randomdef::RandomDtypeFmtGen统一生成输入输出的数据类型序列与 ND 格式shape、seed均标记ValueDepend(OPTIONAL)表明这些输入的值在编译/推导期会被读取AICore 配置开启DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)、PrecisionReduceFlag(true)即支持动态 rank、动态 shape 与精度降低模式通过this-AICore().AddConfig(ascend950)与AddConfig(ascend350)注册支持的芯片配置。2. Shape 推导以 shape 输入为准stateless_bernoulli_infershape.cpp 实现InferShapeStatelessBernoulli通过inputMap/outputMap建立输入输出索引映射调用randomCommon::CommonInferShape(context, inputMap, outputMap, MODE_DEPENDENCY)采用依赖模式——即输出 shape 完全由shape输入张量的取值决定这也是shape被标记ValueDepend的原因IMPL_OP_INFERSHAPE(StatelessBernoulli).InputsDataDependency({0})显式声明输入 0shape在编译期为数据依赖项框架会将其值下发到推导逻辑中。图侧 stateless_bernoulli_graph_infer.cpp 则通过MODE_ATTR与dtype属性完成输出数据类型的推导。3. TilingSIMT 任务切分与参数提取stateless_bernoulli_tiling_arch35.cpp 是 arch35Ascend 950平台的 Tiling 实现通过inputCheckRules/outputCheckRules校验各输入的 dtype 与维度shape 支持任意维 -1prob 支持 FLOAT/FLOAT16/BF16 等getOutputSize复用RandomUtils::GetAndCheckOutputSize从shape输入读取并校验输出规模getSeedAndOffset从 seed/offset 张量中读取标量值并执行offset % 4 0的强校验kernelMode RandomKernelMode::SIMTDcacheSize 128KB并计算每个核的负载CORE_MINIEST_NUM 256对齐与usedCoreNum关键信息extraInt64Param1记录 prob 是否为单元素标量张量probTensorSize 1供内核选择标量/张量采样分支。4. Kernel 内核Philox 算法与 SIMT 并行采样stateless_bernoulli_simt.h 与 stateless_bernoulli.cpp 构成计算核心随机数生成采用Philox 算法RandomKernelBase::PhiloxAlgParsInit初始化 key 与 counterPhiloxRandomSimt生成随机浮点数内核按STEP 4批量生成随机数PhiloxBernoulliSample中每次迭代处理 4 个元素这正是 offset 必须为 4 的倍数的根本原因——counter 以 4 为粒度推进offset 步进需与之一致每个线程以blockIdx.x * blockDim.x threadIdx.x为起点、按 grid 总线程数步进512 线程启动PHILOX_THREAD_LAUNCH 512利用magic/shiftGetUintDivMagicAndShift做高效的整数除法索引映射采样判定逻辑为outGm[i j] results[j] probFp32 ? 1 : 0生成的 [0,1) 均匀随机数小于等于 prob 则输出 1否则输出 0PROB_MODE模板参数区分PROB_SCALARprob 取首元素与PROB_TENSORprob 按元素索引两种模式由 Tiling 下发的extraInt64Param1决定内核入口 stateless_bernoulli.cpp 按DTYPE_PROBfloat/half/bfloat16_t实例化模板并针对bool输出类型将其映射为int8_t写回ConditionalIsSameTypeDTYPE_Y, bool, int8_t, DTYPE_Y。5. op_api 调度AICore 优先、AICPU 兜底op_api/stateless_bernoulli.cpp 实现了 aclnn 接口到算子的调度逻辑StatelessBernoulli入口先将 Host 侧的seed/offset标量转换为 INT64 张量将input的 shape 转为 INT32/INT64 的shapeTensor元素总个数超过INT32_MAX时使用 INT64并依据输入 shape 与 dtype 分配输出out通过IsAiCoreSupport检查输出 dtype 与 prob dtype 是否落在 AICore 支持列表内若是走StatelessBernoulliAiCoreADD_TO_LAUNCHER_LIST_AICORE否则降级到StatelessBernoulliAiCpuADD_TO_LAUNCHER_LIST_AICPU两者均通过OP_ATTR_NAMES({dtype})将输出数据类型作为属性传递给内核这与算子定义中的dtype属性一一对应。6. 图融合 Passop_graph/fusion_pass/stateless_bernoulli_fusion_pass.cpp 实现了StatelessBernoulliV2 → StatelessBernoulli的图融合并针对不同 GE 编译器版本9.0.0 及以上使用kCompatibleInherited阶段做了编译期宏与运行时版本双重保护对应的融合 Pass 单测位于 tests/ut/op_graph/test_stateless_bernoulli_fusion_pass.cpp可作为理解融合逻辑的参考。测试与验证资源仓库为该算子提供了完整的测试覆盖可用于自行验证行为tests/ut/op_host/test_stateless_bernoulli_infershape.cpp验证 shape 依赖推导逻辑tests/ut/op_host/arch35/test_stateless_bernoulli_tiling.cpp验证 SIMT Tiling 参数与 offset 校验tests/ut/op_kernel/test_stateless_bernoulli.cpp验证 Philox 采样内核的输出分布与正确性。总结stateless_bernoulli是 CANN ops-math 中面向 Ascend 950 系列提供的高效伯努利采样算子接口层面同时提供标量 probaclnnBernoulli与张量 probaclnnBernoulliTensor两种形态并支持 inplace 变体实现层面采用shape 值依赖推导 SIMT 任务切分 Philox 随机算法的组合通过seed/offset实现完全可复现的无状态随机采样。读者在实战中需重点注意两类约束prob 必须在 [0,1] 区间内offset 必须为 4 的倍数。深入阅读本文引用的源码与测试文件即可完整掌握从接口调用到底层随机数生成的全链路细节。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
