cann/cannbot-skills HIVM三元向量运算

HIVM 三元向量运算【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills关键词HIVM, Ternary, vsel, select, conditional概述HIVM 三元向量运算继承自HIVM_ElementwiseTernaryOp对三个输入操作数执行逐元条件选择运算。目前仅有一个操作hir.vsel它根据条件向量的值从两个数据源向量中选择元素。Python API 对应tl.where(condition, x, y)— 条件选择操作。IR 操作定义基类HIVM_ElementwiseTernaryOpclass HIVM_ElementwiseTernaryOpstring mnemonic, listTrait traits [] : HIVM_ElementwiseNaryOpmnemonic, !listconcat([ElementwiseNaryOpTrait3], traits);源码参考HIVMVectorOps.td#L1016-L1018hir.vsel — 逐元条件选择TableGen 定义def VSelOp : HIVM_ElementwiseTernaryOpvsel, [StaticMaxRankTrait1, OperElemTypeConstraints[/*condition*/0], [I1,I8], OperElemTypeConstraints[/*src0*/1, /*src0*/2], [I1, AnyI8, AnyI16, F16, BF16, AnyI32, F32, I64], DeclareOpInterfaceMethodsExtraBufferOpInterface, [getExtraBufferSize], BroadcastableOTF] { let summary Elementwise Vector Selection Op; let description baseClassDescription # [{ Select elements from two source vector according to the binary condition vector. If the corresponding bit of the indicator is 1, select src0. Otherwise, select src1. Additional constraints: 1. The input vectors and output vector must have the same ranks. 2. The element type of indicator vector must be bool. }]; let arguments (ins VariadicAnyType:$src, VariadicAnyShaped:$dst, OptionalAnyMemRef:$temp_buffer, DefaultValuedAttrDenseI64ArrayAttr, {}:$transpose, DefaultValuedAttrDenseI64ArrayAttr, {}:$broadcast ); let assemblyFormat [{ attr-dict ins ( $src : type($src) ) outs ( $dst : type($dst) ) (temp_buffer ( $temp_buffer^ : type($temp_buffer) ))? (- type($result)^)? }]; }源码参考HIVMVectorOps.td#L1020-L1050参数说明参数类型必选说明约束$src[0]AnyType是条件向量condition/indicator元素类型 I1 或 I8$src[1]AnyType是数据源 src0条件为真时选择元素类型 I1/AnyI8/AnyI16/F16/BF16/AnyI32/F32/I64$src[2]AnyType是数据源 src1条件为假时选择与 src0 相同元素类型$dstVariadicAnyShaped是输出向量与 src0/src1 相同元素类型$temp_bufferOptionalAnyMemRef否临时缓冲区ExtraBufferOpInterface$transposeDenseI64ArrayAttr (默认 {})否OTF 转置维度-$broadcastDenseI64ArrayAttr (默认 {})否OTF 广播维度-数据类型约束操作数位置语义支持的元素类型约束来源$src[0]条件I1, I8OperElemTypeConstraints[0], [I1, I8]$src[1]数据源0I1, AnyI8, AnyI16, F16, BF16, AnyI32, F32, I64OperElemTypeConstraints[1, 2], [I1, AnyI8, AnyI16, F16, BF16, AnyI32, F32, I64]$src[2]数据源1I1, AnyI8, AnyI16, F16, BF16, AnyI32, F32, I64OperElemTypeConstraints[1, 2], ...语义说明dst[i] condition[i] ? src0[i] : src1[i]当条件向量的对应位为 1 时选择 src0 的元素否则选择 src1 的元素。IR 示例%result hivm.hir.vsel ins(%cond, %src0, %src1 : tensor23x77xi1, f32, tensor23x77xf32) outs(%dst : tensor23x77xf32) - tensor23x77xf32 %result hivm.hir.vsel ins(%cond, %a, %b : i1, tensor5x?x10xf32, tensor5x?x10xf32) outs(%dst : tensor5x?x10xf32) - tensor5x?x10xf32完整使用示例来自测试文件%cond hivm.hir.vcmp ins(%a, %b : tensor23x77xf32, f32) outs(%init : tensor23x77xi1) compare_mode ne - tensor23x77xi1 %result hivm.hir.vsel ins(%cond, %val_true, %val_false : tensor23x77xi1, f32, tensor23x77xf32) outs(%dst : tensor23x77xf32) - tensor23x77xf32IR 层约束与验证条件类型约束条件向量$src[0]的元素类型必须为 I1 或 I8数据源类型一致性src0 和 src1 必须具有相同的元素类型Rank 一致性所有输入向量和输出向量必须具有相同的 rank最大 Rank 限制StaticMaxRankTrait1仅支持 1 维BroadcastableOTF支持 OTF 广播允许条件或数据源在指定维度上进行广播与其他 IR 操作的关系HIVM 操作上游降级HFusion 降级说明vsellinalg.select-语义等价于 linalg.selectvsel 通常与 vcmp 配合使用形成条件选择模式vcmp → vsel (比较后选择)降级示例%cond hivm.hir.vcmp ... compare_mode ne - tensorNxi1 %result hivm.hir.vsel ins(%cond, %a, %b : ...) - tensorNxf32 %result linalg.select ins(%cond, %a, %b : tensorNxi1, tensorNxf32, tensorNxf32) - tensorNxf32常见问题Q: vsel 的条件向量为什么支持 I8 而不仅仅是 I1A: 硬件实现中条件判断可以基于 I8 类型的非零值不仅仅是布尔值。这提供了更大的灵活性允许直接使用比较结果或掩码向量。Q: vsel 支持标量条件吗A: 是的。从 IR 示例可以看到$src 支持 AnyType而非仅 AnyShaped因此条件可以是标量i1值。当条件为标量时所有元素使用相同的条件值。Q: vsel 的最大 Rank 为什么只有 1A: 这是当前硬件实现的限制。对于多维条件选择需要先展平为 1 维操作或在编译阶段通过循环分解处理。相关文档Python APIdocs_triton_ascend 中的tl.where()文档源码参考HIVMVectorOps.td - VSelOpconvert-hivm-to-upstream.mlir【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考