ThresholdV2 阈值算子全解析CANN ops-nn 中的图模式构图与 aclnn 双接口调用实践【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读ThresholdV2 是 CANN ops-nn 算子库中面向 NPU 实现的阈值截断算子其语义与 PyTorch 的torch.nn.Threshold对齐对输入张量 x 逐元素判断大于阈值 threshold 的元素原样保留否则替换为 value。本文以 activation/threshold_v2/README.md 为骨架完整梳理该算子的产品支持矩阵、参数约束与计算公式并深入 examples/test_geir_threshold_v2.cpp、op_graph/threshold_v2_proto.h 等源码讲解图模式构图与 aclnn 两段式接口两种调用路径最终下沉到 tiling 与 Kernel DAG 的实现细节。读完本文你将掌握 ThresholdV2 的完整调用姿势并理解一个 elewise 类算子在 CANN 中从构图、infershape、tiling 到 kernel 执行的完整链路。一、算子功能与产品支持情况1.1 功能语义ThresholdV2 对输入 x 逐元素执行阈值操作当 x 中的元素大于 threshold 时返回该元素本身否则返回 value。计算公式为$$ out(x) \begin{cases} x, x\gt threshold \ value, otherwise \end{cases} $$从算子定义看threshold_v2_proto.h 中的REG_OP(ThresholdV2)明确将其声明为三个输入x、threshold、value、一个输出y的结构且value被标记为OPTIONAL_INPUT注释中写明Compatible with the Pytorch operator Threshold即与 PyTorch Threshold 算子兼容。1.2 产品支持矩阵根据 READMEThresholdV2 的产品支持情况如下产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品√这一支持矩阵与仓库内的实现布局相互印证ThresholdV2 的 host 侧 tiling 实现位于 op_host/arch35/kernel 实现位于 op_kernel/arch35/二进制配置位于 op_host/config/ascend950/threshold_v2_binary.json说明其主要面向 arch35Ascend 950 系列及 Atlas A2/A3 训练与推理系列平台而 Atlas 200I/500 A2 推理产品与 Atlas 推理系列产品不在支持之列。二、参数说明2.1 图模式算子参数README 中给出了 ThresholdV2 算子的完整参数表参数名输入/输出/属性描述数据类型数据格式x输入对应公式中的 xBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDthreshold输入对应公式中的 thresholdBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvalue输入对应公式中的 valueBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDy输出公式中的输出张量BFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8ND结合 threshold_v2_proto.h 的接口注释可以补充几个关键约束xND 格式张量支持 1D~8Dthreshold形状为(1,)的张量作为阈值标量使用value形状为(1,)的张量作为替换值默认值为 0y与输入 x 具有相同的 shape、format 和 dtype。注意README 参数表将value列为输入而源码层面它是OPTIONAL_INPUTthreshold_v2_proto.h、threshold_v2_def.cpp。这一点也体现在 kernel 模板参数上——tiling 会区分有无 value两种编译模式详见后文第四节。2.2 aclnn 接口参数在 aclnn 调用场景下对应 aclnnThresholdaclnnInplaceThreshold.mdthreshold 与 value 不再是以 ND 张量传入而是以aclScalar*标量传入self/out张量的 shape 支持 0~8 维且均支持非连续 Tensor。具体差异对照如下参数aclnn 场景形态说明selfaclTensor*输入张量shape 需与 out 一致支持空 TensorthresholdaclScalar*阈值标量valueaclScalar*替换值标量outaclTensor*输出张量shape 需与 self 一致2.3 数据类型与精度的边界说明README 的约束说明一节声明该算子无约束但在 aclnn 接口文档中有两条需要注意的补充约束aclnnThreshold 与 aclnnInplaceThreshold 为确定性实现当输入为INT32类型时数值超出 [-16777216, 16777216] 范围会存在精度误差因为 INT32 与 float 的转换精度限制在Atlas 训练系列产品上不支持 BFLOAT16 数据类型。三、调用方式一图模式构图调用README 的调用说明一节给出了三种调用入口其中图模式通过算子 IR 构图调用 ThresholdV2。完整示例位于 examples/test_geir_threshold_v2.cpp。3.1 算子 IR 定义op_graph/threshold_v2_proto.h 使用 GE 的REG_OP宏注册算子输入输出全部限定为{DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16}七种数据类型REG_OP(ThresholdV2) .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .INPUT(threshold, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .OPTIONAL_INPUT(value, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .OUTPUT(y, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .OP_END_FACTORY_REG(ThresholdV2);3.2 构图步骤示例主流程分四步初始化 GE → 构图 → 建 Session 并 AddGraph → RunGraph。// 1. 初始化 GE指定 deviceId 与图运行模式 std::mapAscendString, AscendString global_options {{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}}; Status ret ge::GEInitialize(global_options);构图阶段CreateOppInGraph使用op::ThresholdV2(thresholdV21)创建算子节点并通过宏依次设置三个输入auto thresholdV21 op::ThresholdV2(thresholdV21); std::vectorint64_t xShape {8}; // 输入 xshape 为 {8} std::vectorint64_t thresholdShape {8}; // 阈值shape 为 {8} std::vectorint64_t valueShape {8}; // 替换值shape 为 {8} ADD_INPUT(1, x, inDtype, xShape, true, -2.0f, 2.0f, 0.0f); // x 使用 [-2, 2] 随机数 ADD_INPUT(2, threshold, inDtype, thresholdShape, false, 0.0f, 0.0f, 0.5f); // threshold 固定 0.5 ADD_INPUT(3, value, inDtype, valueShape, false, 0.0f, 0.0f, 0.0f); // value 固定 0宏内部为每个输入创建op::Data占位节点与TensorDescFORMAT_ND、kPlacementHost再通过thresholdV21.set_input_x(...)等接口完成算子输入连接最后graph.SetInputs(inputs).SetOutputs(outputs)声明图边界。运行与结果打印ge::Session* session new Session(build_options); ret session-AddGraph(graph_id, graph, graph_options); ret session-RunGraph(graph_id, input, output); // 逐元素打印 x / threshold / value / result 四列结果 LOG_PRINT(x threshold value result\n);示例以 1D shape{8}、FP32 数据为例x 为 [-2, 2] 均匀随机数threshold 固定为 0.5value 固定为 0。按公式大于 0.5 的元素保留原值小于等于 0.5 的元素被替换为 0。四、调用方式二aclnn 两段式接口调用4.1 aclnnThreshold 与 aclnnInplaceThresholdREADME 提供了两条 aclnn 调用路径详见 activation/threshold/docs/aclnnThresholdaclnnInplaceThreshold.mdaclnnThreshold需新建一个输出张量对象存储计算结果out-of-placeaclnnInplaceThreshold无需新建输出张量对象直接在输入张量的内存中就地存储计算结果in-place。两者功能相同按实际场景选择。与 CANN 其他 aclnn 算子一致每个算子都采用两段式接口必须先调用GetWorkspaceSize接口获取计算所需 workspace 大小并得到包含算子计算流程的执行器再调用执行接口完成计算。// 第一段获取 workspace 大小与执行器 aclnnStatus aclnnThresholdGetWorkspaceSize( const aclTensor* self, const aclScalar* threshold, const aclScalar* value, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor) // 第二段执行计算 aclnnStatus aclnnThreshold( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)in-place 版本与之对应aclnnStatus aclnnInplaceThresholdGetWorkspaceSize( aclTensor* selfRef, // 输入/输出共用公式中的 x/out const aclScalar* threshold, const aclScalar* value, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnInplaceThreshold( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)4.2 第一段接口的入参校验与返回码第一段接口会完成入参校验常见错误码如下返回码错误码触发场景aclnnThresholdACLNN_ERR_PARAM_NULLPTR161001self、threshold、value 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 数据类型不在支持范围self 与 threshold、value 不满足数据类型推导规则推导后类型无法转换为 out 的类型self 与 out 的 shape 超过 8 维或两者 shape 不一致in-place 版本的校验更少selfRef的数据类型越界、与 threshold/value 类型推导不满足规则、或 shape 超过 8 维时报 ACLNN_ERR_PARAM_INVALID161002空指针时报 161001。4.3 完整调用示例aclnnThreshold以 FP32、shape{8}为例核心调用序列如下完整代码见 aclnnThresholdaclnnInplaceThreshold.md 的调用示例章节// 1. 初始化 device 与 stream int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // 2. 构造输入与输出 std::vectorfloat selfHostData {0, 1, 2, 3, 4.1, 5, 6, 7}; float thresholdVal 4.1f; // 阈值标量 float valueVal 10.0f; // 替换值标量 CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, ACL_FLOAT, self); threshold aclCreateScalar(thresholdVal, ACL_FLOAT); value aclCreateScalar(valueVal, ACL_FLOAT); CreateAclTensor(outHostData, outShape, outDeviceAddr, ACL_FLOAT, out); // 3. 两段式调用 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnThresholdGetWorkspaceSize(self, threshold, value, out, workspaceSize, executor); if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } ret aclnnThreshold(workspaceAddr, workspaceSize, executor, stream); // 4. 同步等待并取回结果 aclrtSynchronizeStream(stream); aclrtMemcpy(resultData.data(), ..., outDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST);输入{0, 1, 2, 3, 4.1, 5, 6, 7}、threshold4.1、value10 时计算结果为{10, 10, 10, 10, 4.1, 5, 6, 7}——大于 4.1 的元素4.1 自身及 5/6/7保留原值其余替换为 10。in-place 版本aclnnInplaceThreshold唯一的结构差异是第一段接口不接收out参数计算直接写入self张量内存从结果拷贝时同样读取selfDeviceAddr对应的内存。五、源码级实现原理5.1 Host 侧infershape 与 tilingInfershapeThresholdV2 属于逐元素elewise算子其形状推导直接复用Ops::Base::InferShape4Elewisethreshold_v2_infershape.cpp保证输出 y 与输入 x 同形。Tiling 参数校验threshold_v2_tiling_arch35.cpp分三层CalcInputDtype校验 x、threshold、value 的 dtype 均须落在DT_FLOAT16、DT_BF16、DT_INT8、DT_UINT8、DT_INT16、DT_INT32、DT_FLOAT、DT_INT64范围内注意比 README 表多出 INT16 一类tiling 层同样接受CalcOutputDtype要求 y 的 dtype 与 x 完全一致CheckShape要求 y 的 shape 与 x 完全一致若 value 输入缺失则置hasValue false。RunTiling根据输入 dtype 分派到不同的 DAG 模板并按hasValue设置模板参数if (this-inputDtype ge::DT_BF16) { res hasValue ? elewiseBaseTiling.DoTilingThresholdCastDagbfloat16_t, float::OpDag(*tiling) : elewiseBaseTiling.DoTilingThresholdCastDagNoValuebfloat16_t, float::OpDag(*tiling); } else if (this-inputDtype ge::DT_FLOAT16) { res hasValue ? elewiseBaseTiling.DoTilingThresholdDaghalf::OpDag(*tiling) : ... }最后通过SetTilingKey(GET_TPL_TILING_KEY(tiling-scheMode, hasValue ? TPL_HAS_VALUE : TPL_NO_VALUE))与SetBlockDim(tiling-blockNum)下发切分信息。模板参数的合法取值定义在 threshold_tiling_struct.hschMode取TPL_SCH_MODE_0/TPL_SCH_MODE_1两种调度模式valueMode取TPL_NO_VALUE/TPL_HAS_VALUE两种取值模式。5.2 Kernel 侧DAG 计算图Kernel 入口 threshold_v2.cpp 是一个带模板参数的 AIV kernel按valueMode与输入类型选择不同的计算 DAG统一经由ElementwiseSch调度执行。计算 DAG 定义在 threshold_dag.h以最常用的ThresholdDagT为例其数据流为CopyIn(x) ──────────────► Compare(LE) ─► Select ─► CopyOut(y) ▲ ▲ ThresholdScalarHolder ────────┘ │ ValueScalarHolder ───────────────────────────┘实现要点OpCompare使用Vec::Compareuint8_t, T, CMPMODE_LE其中CMPMODE_LE 3表示小于等于比较即先算出(x threshold)的掩码OpSelectResult使用Vec::Selectuint8_t, T, VSEL_TENSOR_SCALAR_MODEVSEL_TENSOR_SCALAR_MODE 1做张量-标量选择掩码为真时取 value否则取 x恰好等价于x threshold ? x : value无 value 的ThresholdDagNoValue变体用MAKE_CONST(T, 0)常量 0 替换 value 标量输入对应源码中value 默认值为 0的语义BF16 特判由于 BF16 缺少直接可用的比较指令ThresholdCastDagbfloat16_t, float先将 x 与 threshold/value 标量Cast到 float 比较再通过CAST_RINT回写到 BF16保证精度与确定性。5.3 二进制配置与测试threshold_v2_binary.json 记录了 ascend950 平台预编译的 8 组 binary 清单bfloat16、float16、float32、int8、uint8、int16、int32、int64每组均声明 x/threshold 为 required、value 为 optional、shape 为[-2]动态 shape 标记格式统一为 ND测试侧host 层单测覆盖 tiling 与 infershapetests/ut/op_host/arch35/test_threshold_v2_tiling.cpp、tests/ut/op_host/test_threshold_infershape.cpp另有 ST 用例 tests/st/arch35/ttk_kernel_threshold_v2_st.csv 与 golden 数据生成脚本 tests/assets/golden.py可从数据侧验证x threshold ? x : value的语义。六、总结ThresholdV2 是 CANN ops-nn 中一个结构简洁但链路完整的 elewise 阈值算子语义简单y (x threshold) ? x : value兼容 PyTorch Threshold双调用入口图模式GE 构图适合整图下沉 NPU 的场景aclnn 两段式接口含 in-place 变体适合在 acl 编程框架中按算子粒度直接调用实现分层清晰proto 定义算子签名 → infershape 保证同形推导 → tiling 负责 dtype/shape 校验与切分含 hasValue 模板分支→ kernel 以 CompareSelect DAG 完成计算并对 BF16 做 float 中转的精度处理。无论你是在做网络算子迁移、自定义构图还是需要在 acl 层逐算子控制计算流程ThresholdV2 的这套实现与调用范式都可直接复用。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
ThresholdV2 阈值算子全解析：CANN ops-nn 中的图模式构图与 aclnn 双接口调用实践
ThresholdV2 阈值算子全解析CANN ops-nn 中的图模式构图与 aclnn 双接口调用实践【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读ThresholdV2 是 CANN ops-nn 算子库中面向 NPU 实现的阈值截断算子其语义与 PyTorch 的torch.nn.Threshold对齐对输入张量 x 逐元素判断大于阈值 threshold 的元素原样保留否则替换为 value。本文以 activation/threshold_v2/README.md 为骨架完整梳理该算子的产品支持矩阵、参数约束与计算公式并深入 examples/test_geir_threshold_v2.cpp、op_graph/threshold_v2_proto.h 等源码讲解图模式构图与 aclnn 两段式接口两种调用路径最终下沉到 tiling 与 Kernel DAG 的实现细节。读完本文你将掌握 ThresholdV2 的完整调用姿势并理解一个 elewise 类算子在 CANN 中从构图、infershape、tiling 到 kernel 执行的完整链路。一、算子功能与产品支持情况1.1 功能语义ThresholdV2 对输入 x 逐元素执行阈值操作当 x 中的元素大于 threshold 时返回该元素本身否则返回 value。计算公式为$$ out(x) \begin{cases} x, x\gt threshold \ value, otherwise \end{cases} $$从算子定义看threshold_v2_proto.h 中的REG_OP(ThresholdV2)明确将其声明为三个输入x、threshold、value、一个输出y的结构且value被标记为OPTIONAL_INPUT注释中写明Compatible with the Pytorch operator Threshold即与 PyTorch Threshold 算子兼容。1.2 产品支持矩阵根据 READMEThresholdV2 的产品支持情况如下产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品√这一支持矩阵与仓库内的实现布局相互印证ThresholdV2 的 host 侧 tiling 实现位于 op_host/arch35/kernel 实现位于 op_kernel/arch35/二进制配置位于 op_host/config/ascend950/threshold_v2_binary.json说明其主要面向 arch35Ascend 950 系列及 Atlas A2/A3 训练与推理系列平台而 Atlas 200I/500 A2 推理产品与 Atlas 推理系列产品不在支持之列。二、参数说明2.1 图模式算子参数README 中给出了 ThresholdV2 算子的完整参数表参数名输入/输出/属性描述数据类型数据格式x输入对应公式中的 xBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDthreshold输入对应公式中的 thresholdBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvalue输入对应公式中的 valueBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDy输出公式中的输出张量BFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8ND结合 threshold_v2_proto.h 的接口注释可以补充几个关键约束xND 格式张量支持 1D~8Dthreshold形状为(1,)的张量作为阈值标量使用value形状为(1,)的张量作为替换值默认值为 0y与输入 x 具有相同的 shape、format 和 dtype。注意README 参数表将value列为输入而源码层面它是OPTIONAL_INPUTthreshold_v2_proto.h、threshold_v2_def.cpp。这一点也体现在 kernel 模板参数上——tiling 会区分有无 value两种编译模式详见后文第四节。2.2 aclnn 接口参数在 aclnn 调用场景下对应 aclnnThresholdaclnnInplaceThreshold.mdthreshold 与 value 不再是以 ND 张量传入而是以aclScalar*标量传入self/out张量的 shape 支持 0~8 维且均支持非连续 Tensor。具体差异对照如下参数aclnn 场景形态说明selfaclTensor*输入张量shape 需与 out 一致支持空 TensorthresholdaclScalar*阈值标量valueaclScalar*替换值标量outaclTensor*输出张量shape 需与 self 一致2.3 数据类型与精度的边界说明README 的约束说明一节声明该算子无约束但在 aclnn 接口文档中有两条需要注意的补充约束aclnnThreshold 与 aclnnInplaceThreshold 为确定性实现当输入为INT32类型时数值超出 [-16777216, 16777216] 范围会存在精度误差因为 INT32 与 float 的转换精度限制在Atlas 训练系列产品上不支持 BFLOAT16 数据类型。三、调用方式一图模式构图调用README 的调用说明一节给出了三种调用入口其中图模式通过算子 IR 构图调用 ThresholdV2。完整示例位于 examples/test_geir_threshold_v2.cpp。3.1 算子 IR 定义op_graph/threshold_v2_proto.h 使用 GE 的REG_OP宏注册算子输入输出全部限定为{DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16}七种数据类型REG_OP(ThresholdV2) .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .INPUT(threshold, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .OPTIONAL_INPUT(value, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .OUTPUT(y, TensorType({DT_FLOAT16, DT_FLOAT32, DT_INT8, DT_INT32, DT_UINT8, DT_INT64, DT_BF16})) .OP_END_FACTORY_REG(ThresholdV2);3.2 构图步骤示例主流程分四步初始化 GE → 构图 → 建 Session 并 AddGraph → RunGraph。// 1. 初始化 GE指定 deviceId 与图运行模式 std::mapAscendString, AscendString global_options {{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}}; Status ret ge::GEInitialize(global_options);构图阶段CreateOppInGraph使用op::ThresholdV2(thresholdV21)创建算子节点并通过宏依次设置三个输入auto thresholdV21 op::ThresholdV2(thresholdV21); std::vectorint64_t xShape {8}; // 输入 xshape 为 {8} std::vectorint64_t thresholdShape {8}; // 阈值shape 为 {8} std::vectorint64_t valueShape {8}; // 替换值shape 为 {8} ADD_INPUT(1, x, inDtype, xShape, true, -2.0f, 2.0f, 0.0f); // x 使用 [-2, 2] 随机数 ADD_INPUT(2, threshold, inDtype, thresholdShape, false, 0.0f, 0.0f, 0.5f); // threshold 固定 0.5 ADD_INPUT(3, value, inDtype, valueShape, false, 0.0f, 0.0f, 0.0f); // value 固定 0宏内部为每个输入创建op::Data占位节点与TensorDescFORMAT_ND、kPlacementHost再通过thresholdV21.set_input_x(...)等接口完成算子输入连接最后graph.SetInputs(inputs).SetOutputs(outputs)声明图边界。运行与结果打印ge::Session* session new Session(build_options); ret session-AddGraph(graph_id, graph, graph_options); ret session-RunGraph(graph_id, input, output); // 逐元素打印 x / threshold / value / result 四列结果 LOG_PRINT(x threshold value result\n);示例以 1D shape{8}、FP32 数据为例x 为 [-2, 2] 均匀随机数threshold 固定为 0.5value 固定为 0。按公式大于 0.5 的元素保留原值小于等于 0.5 的元素被替换为 0。四、调用方式二aclnn 两段式接口调用4.1 aclnnThreshold 与 aclnnInplaceThresholdREADME 提供了两条 aclnn 调用路径详见 activation/threshold/docs/aclnnThresholdaclnnInplaceThreshold.mdaclnnThreshold需新建一个输出张量对象存储计算结果out-of-placeaclnnInplaceThreshold无需新建输出张量对象直接在输入张量的内存中就地存储计算结果in-place。两者功能相同按实际场景选择。与 CANN 其他 aclnn 算子一致每个算子都采用两段式接口必须先调用GetWorkspaceSize接口获取计算所需 workspace 大小并得到包含算子计算流程的执行器再调用执行接口完成计算。// 第一段获取 workspace 大小与执行器 aclnnStatus aclnnThresholdGetWorkspaceSize( const aclTensor* self, const aclScalar* threshold, const aclScalar* value, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor) // 第二段执行计算 aclnnStatus aclnnThreshold( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)in-place 版本与之对应aclnnStatus aclnnInplaceThresholdGetWorkspaceSize( aclTensor* selfRef, // 输入/输出共用公式中的 x/out const aclScalar* threshold, const aclScalar* value, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnInplaceThreshold( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)4.2 第一段接口的入参校验与返回码第一段接口会完成入参校验常见错误码如下返回码错误码触发场景aclnnThresholdACLNN_ERR_PARAM_NULLPTR161001self、threshold、value 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 数据类型不在支持范围self 与 threshold、value 不满足数据类型推导规则推导后类型无法转换为 out 的类型self 与 out 的 shape 超过 8 维或两者 shape 不一致in-place 版本的校验更少selfRef的数据类型越界、与 threshold/value 类型推导不满足规则、或 shape 超过 8 维时报 ACLNN_ERR_PARAM_INVALID161002空指针时报 161001。4.3 完整调用示例aclnnThreshold以 FP32、shape{8}为例核心调用序列如下完整代码见 aclnnThresholdaclnnInplaceThreshold.md 的调用示例章节// 1. 初始化 device 与 stream int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // 2. 构造输入与输出 std::vectorfloat selfHostData {0, 1, 2, 3, 4.1, 5, 6, 7}; float thresholdVal 4.1f; // 阈值标量 float valueVal 10.0f; // 替换值标量 CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, ACL_FLOAT, self); threshold aclCreateScalar(thresholdVal, ACL_FLOAT); value aclCreateScalar(valueVal, ACL_FLOAT); CreateAclTensor(outHostData, outShape, outDeviceAddr, ACL_FLOAT, out); // 3. 两段式调用 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnThresholdGetWorkspaceSize(self, threshold, value, out, workspaceSize, executor); if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } ret aclnnThreshold(workspaceAddr, workspaceSize, executor, stream); // 4. 同步等待并取回结果 aclrtSynchronizeStream(stream); aclrtMemcpy(resultData.data(), ..., outDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST);输入{0, 1, 2, 3, 4.1, 5, 6, 7}、threshold4.1、value10 时计算结果为{10, 10, 10, 10, 4.1, 5, 6, 7}——大于 4.1 的元素4.1 自身及 5/6/7保留原值其余替换为 10。in-place 版本aclnnInplaceThreshold唯一的结构差异是第一段接口不接收out参数计算直接写入self张量内存从结果拷贝时同样读取selfDeviceAddr对应的内存。五、源码级实现原理5.1 Host 侧infershape 与 tilingInfershapeThresholdV2 属于逐元素elewise算子其形状推导直接复用Ops::Base::InferShape4Elewisethreshold_v2_infershape.cpp保证输出 y 与输入 x 同形。Tiling 参数校验threshold_v2_tiling_arch35.cpp分三层CalcInputDtype校验 x、threshold、value 的 dtype 均须落在DT_FLOAT16、DT_BF16、DT_INT8、DT_UINT8、DT_INT16、DT_INT32、DT_FLOAT、DT_INT64范围内注意比 README 表多出 INT16 一类tiling 层同样接受CalcOutputDtype要求 y 的 dtype 与 x 完全一致CheckShape要求 y 的 shape 与 x 完全一致若 value 输入缺失则置hasValue false。RunTiling根据输入 dtype 分派到不同的 DAG 模板并按hasValue设置模板参数if (this-inputDtype ge::DT_BF16) { res hasValue ? elewiseBaseTiling.DoTilingThresholdCastDagbfloat16_t, float::OpDag(*tiling) : elewiseBaseTiling.DoTilingThresholdCastDagNoValuebfloat16_t, float::OpDag(*tiling); } else if (this-inputDtype ge::DT_FLOAT16) { res hasValue ? elewiseBaseTiling.DoTilingThresholdDaghalf::OpDag(*tiling) : ... }最后通过SetTilingKey(GET_TPL_TILING_KEY(tiling-scheMode, hasValue ? TPL_HAS_VALUE : TPL_NO_VALUE))与SetBlockDim(tiling-blockNum)下发切分信息。模板参数的合法取值定义在 threshold_tiling_struct.hschMode取TPL_SCH_MODE_0/TPL_SCH_MODE_1两种调度模式valueMode取TPL_NO_VALUE/TPL_HAS_VALUE两种取值模式。5.2 Kernel 侧DAG 计算图Kernel 入口 threshold_v2.cpp 是一个带模板参数的 AIV kernel按valueMode与输入类型选择不同的计算 DAG统一经由ElementwiseSch调度执行。计算 DAG 定义在 threshold_dag.h以最常用的ThresholdDagT为例其数据流为CopyIn(x) ──────────────► Compare(LE) ─► Select ─► CopyOut(y) ▲ ▲ ThresholdScalarHolder ────────┘ │ ValueScalarHolder ───────────────────────────┘实现要点OpCompare使用Vec::Compareuint8_t, T, CMPMODE_LE其中CMPMODE_LE 3表示小于等于比较即先算出(x threshold)的掩码OpSelectResult使用Vec::Selectuint8_t, T, VSEL_TENSOR_SCALAR_MODEVSEL_TENSOR_SCALAR_MODE 1做张量-标量选择掩码为真时取 value否则取 x恰好等价于x threshold ? x : value无 value 的ThresholdDagNoValue变体用MAKE_CONST(T, 0)常量 0 替换 value 标量输入对应源码中value 默认值为 0的语义BF16 特判由于 BF16 缺少直接可用的比较指令ThresholdCastDagbfloat16_t, float先将 x 与 threshold/value 标量Cast到 float 比较再通过CAST_RINT回写到 BF16保证精度与确定性。5.3 二进制配置与测试threshold_v2_binary.json 记录了 ascend950 平台预编译的 8 组 binary 清单bfloat16、float16、float32、int8、uint8、int16、int32、int64每组均声明 x/threshold 为 required、value 为 optional、shape 为[-2]动态 shape 标记格式统一为 ND测试侧host 层单测覆盖 tiling 与 infershapetests/ut/op_host/arch35/test_threshold_v2_tiling.cpp、tests/ut/op_host/test_threshold_infershape.cpp另有 ST 用例 tests/st/arch35/ttk_kernel_threshold_v2_st.csv 与 golden 数据生成脚本 tests/assets/golden.py可从数据侧验证x threshold ? x : value的语义。六、总结ThresholdV2 是 CANN ops-nn 中一个结构简洁但链路完整的 elewise 阈值算子语义简单y (x threshold) ? x : value兼容 PyTorch Threshold双调用入口图模式GE 构图适合整图下沉 NPU 的场景aclnn 两段式接口含 in-place 变体适合在 acl 编程框架中按算子粒度直接调用实现分层清晰proto 定义算子签名 → infershape 保证同形推导 → tiling 负责 dtype/shape 校验与切分含 hasValue 模板分支→ kernel 以 CompareSelect DAG 完成计算并对 BF16 做 float 中转的精度处理。无论你是在做网络算子迁移、自定义构图还是需要在 acl 层逐算子控制计算流程ThresholdV2 的这套实现与调用范式都可直接复用。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
