
深入解析 CANN pyasc 的 asc.language.basic.get_hccl_context获取通信域消息区地址的底层实现与实战用法【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc导读asc.language.basic.get_hccl_context是 CANN pyasc 为 Python 开发者提供的算子编程接口之一用于在昇腾 AI 处理器上获取指定 Index 通信域communication domain的 context消息区地址。该接口与 Ascend C 中的GetHcclContext函数一一对应是多核通信与集合通信场景下的基础能力。读完本文你将掌握该接口的调用方式、参数约束、返回值语义理解它从 Python API 到 IR 再到最终 C 代码生成的完整链路并能将其与set_hccl_context配套应用于通算融合算子的通信域初始化流程。接口总览asc.language.basic.get_hccl_context是 pyasc 在python/asc/language/basic/模块中暴露的 basic 层接口其完整签名与语义如下asc.language.basic.get_hccl_context(index: int) → GlobalAddress该接口用于获取指定 Index 通信域的 context消息区地址。所谓消息区是指通信域内用于存放通信上下文、同步状态等消息信息的全局内存区域多核/多卡通信过程中各计算单元需要访问该区域来协同完成集合通信操作。对应的 Ascend C 函数原型pyasc 的 Python 接口与 Ascend C 一一对应get_hccl_context对应的原生 C 原型为template uint32_t index __aicore__ inline __gm__ uint8_t* __gm__ GetHcclContext(void)可以看到index是一个编译期模板参数类型为uint32_t而非运行时函数参数函数返回一个__gm__全局内存空间Global Memory下的uint8_t*指针即全局内存字节地址函数无运行时入参地址完全由模板参数index在编译期确定。这一模板化的设计在 pyasc 的 Python 接口中体现为index作为函数入参传入但底层在生成 IR 时会将其物化为uint32类型的编译期常量最终代码生成阶段仍以模板实参形式展开为AscendC::GetHcclContextindex()。参数说明参数类型含义取值范围indexint要获取 context 的通信域 ID当前只支持 2 个通信域index只能为 0/1返回值说明返回GlobalAddress类型对象表示指定通信域的 context消息区地址。该地址指向全局内存__gm__空间中的字节地址。约束说明当前最多只支持 2 个通信域即index取值仅为0或1超出该范围的行为未定义使用时必须遵守此约束。调用示例ctx asc.get_hccl_context(1)从 Python 源码看接口实现原理get_hccl_context的 Python 实现位于 python/asc/language/basic/common.py完整实现如下require_jit set_common_docstring(api_nameget_hccl_context) def get_hccl_context(index: int) - GlobalAddress: builder global_builder.get_ir_builder() idx_ir _mat(index, KnownTypes.uint32).to_ir() ir_type ir.get_memref_type(builder.get_ui8_type(), [ir.dynshape], ir.AddressSpace.gm) return GlobalAddress(builder.create_asc_GetHcclContextOp(ir_type, idx_ir), KnownTypes.uint8)这段实现清晰地展示了 pyasc 的Python 原生语法 IR 构建架构require_jit装饰器声明该函数必须在 JIT 编译上下文kernel 定义阶段中调用不能在宿主机普通 Python 环境中直接执行。这是 pyasc 所有算子编程接口的统一约束确保代码能在昇腾 AI 处理器的内核kernel中被编译执行。set_common_docstring(api_nameget_hccl_context)自动挂载文档字符串。文档的实际内容定义在 python/asc/language/basic/utils.py 的get_hccl_context_docstring()中这也正是本文所依据的 API 文档docs/python-api/language/generated/asc.language.basic.get_hccl_context.md的生成来源二者内容完全一致。_mat(index, KnownTypes.uint32)将 Python 层的int入参物化materialize为 IR 层的uint32类型值。这一步把 Python 原生值转换为 IR 值与 Ascend C 原型中模板参数uint32_t index的类型对齐。ir.get_memref_type(builder.get_ui8_type(), [ir.dynshape], ir.AddressSpace.gm)构建一个 IR 类型——元素类型为ui8uint8、动态 shapeir.dynshape、位于gmGlobal Memory地址空间的memref类型。这与 C 原型返回__gm__ uint8_t*完全对应全局内存中的字节指针用动态 shape 的 memref 表示。builder.create_asc_GetHcclContextOp(ir_type, idx_ir)调用 IR Builder 创建asc.get_hccl_context操作Op传入目标返回类型与 index 值。返回值包装将 IR 句柄包装为GlobalAddress对象元素数据类型为KnownTypes.uint8。GlobalAddress定义于 python/asc/language/core/ir_value.py是表示全局内存地址的 IR 值封装支持__add__等地址运算例如对该地址做偏移便于从消息区不同位置读取数据。配套接口set_hccl_contextget_hccl_context通常与set_hccl_context成对使用。set_hccl_context用于设置通算融合算子每个通信域对应的 context消息区地址其 Python 实现同样位于 python/asc/language/basic/common.pyrequire_jit set_common_docstring(api_nameset_hccl_context) def set_hccl_context(index: RuntimeInt, context: GlobalAddress) - None: builder global_builder.get_ir_builder() idx_ir _mat(index, KnownTypes.uint32).to_ir() builder.create_asc_SetHcclContextOp(idx_ir, context.to_ir())对应的 Ascend C 原型为template uint32_t index __aicore__ inline void SetHcclContext(__gm__ uint8_t* context)二者的对应关系如下接口作用入参返回get_hccl_context(index)获取指定通信域的 context消息区地址index: 0/1GlobalAddress__gm__ uint8_t*set_hccl_context(index, context)设置指定通信域的 context消息区地址index: 0/1context:GlobalAddress无在通算融合算子的初始化阶段典型的用法是先通过get_hccl_context读取系统为各通信域预分配好的消息区地址必要时对该地址进行解读或搬运再配合set_hccl_context完成通信域的显式配置import asc # 获取通信域 0 / 通信域 1 的 context消息区地址 ctx0 asc.get_hccl_context(0) ctx1 asc.get_hccl_context(1) # 以通信域 0 为例将其 context 地址显式设置回通信域示意 asc.set_hccl_context(0, ctx0)需要注意的是index与 context 的对应关系必须一致读取时使用的通信域 ID 应与后续写入、访问的通信域 ID 匹配。IR 层定义AscendC_GetHcclContextOp在 MLIR 方言层面get_hccl_context对应的操作定义于 include/ascir/Dialect/Asc/IR/Basic/OpHcclContext.tddef AscendC_GetHcclContextOp : APIOpget_hccl_context, GetHcclContext, [AscFunc] { let arguments (ins AnyType:$index); let results (outs AnyType:$context); let paramTypeLists [4]; let assemblyFormat $index attr-dict : type($index) , type($context); }关键信息Op 名称为ascendc.get_hccl_context对应 C 函数名GetHcclContext它属于AscFunc接口族即作为内核函数kernel function内部可调用的功能接口输入为$index输出为$contextparamTypeLists [4]表示该操作对应模板参数列表中的模板参数索引印证了index在底层是以编译期模板参数形式参与代码生成的同文件中还定义了配套的 AscendC_SetHcclContextOp用于set_hccl_context。代码生成验证从 IR 到 Ascend Cpyasc 的完整工具链为Python API → IRMLIR→ Ascend C 代码。get_hccl_context在这一链路中的最终形态可以由仓库中的 Lit 测试用例 test/Target/AscendC/basic/common.mlir 直接验证// CHECK-LABEL:void emit_get_hccl_context(uint32_t v1) { // CHECK-NEXT: __gm__ uint8_t* v2 AscendC::GetHcclContextv1(); // CHECK-NEXT: return; // CHECK-NEXT:} func.func emit_get_hccl_context(%arg0: ui32) { %0 ascendc.get_hccl_context %arg0 : ui32, memref?xui8, 22 return }该测试确认了以下事实IR 输入是ui32对应 Python 层index物化后的uint32类型IR 输出类型为memref?xui8, 22——元素ui8、动态 shape、地址空间 22即AddressSpace.gmGlobal Memory与 Python 实现中构建的 IR 类型完全一致代码生成结果为一行 C__gm__ uint8_t* v2 AscendC::GetHcclContextv1();其中v1是函数入参对应的uint32_t变量被作为编译期模板实参传给AscendC::GetHcclContextindex()。同文件中的 set_hccl_context 测试用例 也验证了配套接口的生成结果AscendC::SetHcclContextv2(v1);即模板实参为 index、运行时参数为 context 指针。典型使用场景与注意事项典型场景通算融合算子的通信域初始化get_hccl_context面向的是**通算融合算子通信与计算融合**场景。在这类算子中多个计算单元需要在同一通信域内交换消息因此必须能够定位每个通信域对应的消息区地址。通过get_hccl_context获取该地址后算子即可在内核内直接访问消息区读取或写入通信上下文信息而无需依赖宿主机侧的额外配置。使用注意事项必须在 JIT 内核上下文中调用get_hccl_context带有require_jit约束只能出现在算子 kernel 定义中由 pyasc 的 JIT 编译框架处理不能在普通 Python 脚本中当作宿主函数调用。通信域数量上限当前最多只支持 2 个通信域index只能是 0 或 1请勿传入其他取值。返回值类型语义返回值是GlobalAddress其元素类型为uint8表示的是全局内存中的字节地址。如需按其他数据类型解读消息区内容可对返回的GlobalAddress做地址偏移或类型转换后使用。与 set_hccl_context 配套读取get与写入set必须针对同一通信域 ID 保持一致避免通信域索引错位导致的消息区访问错误。总结asc.language.basic.get_hccl_context是 pyasc 中访问通信域消息区的基础接口它以 Python 原生函数形式封装了 Ascend C 的模板函数GetHcclContext。从 Python 实现、IR 操作定义 到 代码生成测试整条链路环环相扣index始终以uint32编译期模板参数形式参与编译最终生成访问全局内存消息区的__gm__ uint8_t*指针。理解该接口的语义、约束与底层实现是编写正确的多通信域通算融合算子的前提也为阅读 pyasc 其他 basic 层接口如get_sys_workspace、set_hccl_context提供了可对照的范式。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考