hir.custom — 自定义单 Pipe 操作【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills关键词custom, CustomOp, Builtin, __builtin_gather_load, __builtin_index_select, SinglePipe概述hir.custom是 HIVM 方言中的自定义单 Pipe 操作允许用户在现有操作无法满足需求时编写自定义实现。该操作在单个 Pipeline 上执行通过name属性标识操作类型。Custom 操作适用于以下场景现有操作无法实现所需功能现有操作可以实现但性能不佳需要私有操作内置操作name 以__builtin开头由编译器自带实现用户无需指定属性即可使用。Python API 对应Triton 的自定义操作可以通过tl.extern等机制映射为 hir.custom。IR 操作定义从 HIVMOps.td 提取def CustomOp : HIVM_CustomOpcustom, [SinglePipeOpTrait] { let arguments (ins StrAttr:$name, VariadicAnyType:$inputs, VariadicAnyType:$outputs); let results (outs VariadicAnyType:$results); }基类HIVM_CustomOp定义HIVMOps.tdclass HIVM_CustomOpstring mnemonic, listTrait traits [] : HIVM_StructuredOpmnemonic, !listconcat([AttrSizedOperandSegments, MemoryEffects[MemRead, MemWrite], HIVMInferCoreTypeInterface], traits) { dag args (ins StrAttr:$name, VariadicAnyType:$inputs, VariadicAnyType:$outputs); dag res (outs VariadicAnyType:$results); }参数说明输入操作数ins参数类型必选说明$nameStrAttr是操作名称__builtin_前缀为内置操作$inputsVariadicAnyType是输入参数可变数量$outputsVariadicAnyType是输出/初始化参数DestinationStyle输出结果outs参数类型说明$resultsVariadicAnyType结果值必需属性属性类型说明hivm.tcore_typeTCoreTypeAttr执行的 Core 类型hivm.pipePipeAttr执行的 Pipelinehivm.vf_modeVFModeAttrVector 运行模式Cube Core 时忽略可选属性属性类型说明gm_addr_args_indicesDenseI32ArrayAttr指示哪些参数包含 GM 纯地址额外类方法方法返回类型说明getCoreType()optionalTCoreType获取 Core 类型setCoreType(TCoreType)void设置 Core 类型getPipe()PIPE获取 PipesetPipe(PIPE)void设置 PipegetVFMode()optionalVFMode获取 VF 模式setVFMode(VFMode)void设置 VF 模式getGMAddrArgsIndices()optionalSmallVectorsize_t获取 GM 地址参数索引setGMAddrArgsIndices(SmallVectorsize_t)void设置 GM 地址参数索引isBuiltin()bool判断是否为内置操作getDpsInitsMutable()MutableOperandRangeDestinationStyleOpInterface内置操作常量static constexpr StringLiteral kBuiltinGatherLoadName __builtin_gather_load; static constexpr StringLiteral kBuiltinIndexSelectName __builtin_index_select;IR 示例自定义操作%empty tensor.empty() : tensor3x3xf32 %0 hivm.hir.custom { hivm.tcore_type #hivm.tcore_typeVECTOR, hivm.pipe #hivm.pipePIPE_V, hivm.vf_mode #hivm.vf_modeSIMD } my_custom_op ins(%arg0, %arg1, %c4_i64, %c0_i32, %c2_i64, %c1_i64, %c2_i32, %c2_i32, %c0_i32, %c0_i32 : memref?xf32, tensor3x3xi64, i64, i32, i64, i64, i32, i32, i32, i32) outs(%empty : tensor3x3xf32) - tensor3x3xf32内置操作无需指定属性%empty tensor.empty() : tensor3x3xf32 %0 hivm.hir.custom __builtin_gather_load ins(%arg0, %arg1, %c4_i64, %c0_i32, %c2_i64, %c1_i64, %c2_i32, %c2_i32, %c0_i32, %c0_i32 : memref?xf32, tensor3x3xi64, i64, i32, i64, i64, i32, i32, i32, i32) outs(%empty : tensor3x3xf32) - tensor3x3xf32__builtin_index_select%0 tensor.empty() : tensor1x4x32xf32 %1 hivm.hir.custom {extra_attr srcStrideLength3, hivm.vf_mode #hivm.vf_modeSIMT} __builtin_index_select ins(%arg0, %arg1, %c0_i32, %c9000_i64, %c0_i32, %c0_i32, %c0_i32, %c0_i32, %c0_i32, %c0_i32, %c32_i32, %c4000_i32, %c32_i32 : memref?xf32, tensor16x400xi32, i32, i64, i32, i32, i32, i32, i32, i32, i32, i32, i32) outs(%0 : tensor1x4x32xf32) - tensor1x4x32xf32IR 层约束与验证SinglePipeOpTrait操作在单个 Pipeline 上执行。AttrSizedOperandSegmentsinputs 和 outputs 的数量由属性决定。DestinationStyleOpInterfaceoutputs 作为 DPS init 操作数。Core Type必须通过hivm.tcore_type属性指定内置操作可省略编译器自动推断。Pipe必须通过hivm.pipe属性指定内置操作可省略。VFMode必须通过hivm.vf_mode属性指定内置操作可省略Cube Core 时忽略。Builtin 验证内置操作编译器会检查参数正确性并规范化属性。常见问题Q: 自定义操作如何提供实现A: 当前自定义操作的用户实现机制尚未完全开放标记为 TODO。内置操作由编译器自动链接到模板库。Q: __builtin_gather_load 和 hir.gather_load 的区别A:__builtin_gather_load是通过 custom 操作机制实现的 gather load而hir.gather_load是独立的 HIVM 操作。两者功能类似但 custom 版本更灵活可以通过属性配置。Q: gm_addr_args_indices 的用途A: 指示哪些输入参数包含 GM 的纯地址值。Lowering Pass 在处理这些参数时会保留地址值而不是尝试从 memref 中提取数据。相关文档源码参考HIVMOps.td测试用例custom-op.mlirCustomMacroOp02-custom-macro-op.md【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
