ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

pyasc Matmul 异步迭代同步:asc.language.adv.Matmul.wait_iterate_all 接口深入解析

pyasc Matmul 异步迭代同步:asc.language.adv.Matmul.wait_iterate_all 接口深入解析 pyasc Matmul 异步迭代同步asc.language.adv.Matmul.wait_iterate_all 接口深入解析【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本篇技术指南聚焦 CANN/pyasc 项目中asc.language.adv.Matmul高阶矩阵乘 API 的异步同步接口wait_iterate_all它在iterate_all以异步模式syncFalse将计算结果连续输出到 Global MemoryGM时负责等待异步计算与写回全部完成。读完本文你将掌握wait_iterate_all的签名、配套约束、与iterate_all的完整配合写法、底层 IR 实现原理以及它与wait_iterate_batch、wait_get_tensor_c等其他同步接口的适用场景差异可直接用于在昇腾 AI 处理器上编写异步 Matmul 算子。一、接口概览wait_iterate_all 是什么wait_iterate_all是 pyasc 为 Python 开发者提供的 Ascend C Matmul 高阶 API 同步方法位于asc.language.adv.Matmul类中与 Ascend C 原生接口一一对应并遵循 Python 原生语法见 python/asc/language/adv/matmul.py。函数签名Matmul.wait_iterate_all() → None核心语义等待iterate_all异步接口返回支持连续输出到 Global Memory。也就是说当用户以异步方式发起一次整片 C 矩阵singleCoreM * singleCoreN的计算后通过该接口阻塞等待本次异步迭代彻底结束保证后续对 GM 输出结果的访问是安全、可见的。在 pyasc 的接口清单中官方对该接口的描述为等待 iterate_all 异步接口返回支持连续输出到 Global Memory。见 docs/python-api/language/adv.md 中Matmul.wait_iterate_all条目对应关系说明关联 API 文档中列出的对应的 Ascend C 函数原型为__aicore__ inline void SetSingleShape(int singleM, int singleN, int singleK)见 asc.language.adv.Matmul.wait_iterate_all.md。从源码与 Ascend C 语义看该原型与set_single_shape接口一致属于文档生成时的原型串行wait_iterate_all实际对应 Ascend C 的WaitIterateAll()等待语义在 pyasc 中其行为由 IR 层的等待操作见下文源码级实现剖析体现。二、参数说明wait_iterate_all是一个无参接口参数说明无该方法不接受任何参数直接作用于已构造的Matmul对象由于其无参、无返回值的特性它天然适合作为发射后等待的同步屏障放在异步iterate_all调用之后、对输出 GM 张量执行后续读写如data_copy、Vector 计算或end()释放资源之前。三、配套接口与约束说明wait_iterate_all不能单独使用必须与异步模式的iterate_all配对。以下约束在文档与源码中均有明确依据配套 iterate_all 异步接口使用只有将iterate_all的sync参数置为False并开启wait_iterate_allTrue时wait_iterate_all()的调用才有意义asc.language.adv.Matmul.wait_iterate_all.md。仅支持连续输出至 Global Memoryiterate_all的结果需连续写回 GM此时en_sequential_write应保持关闭False由硬件以连续写模式完成输出。wait_iterate_allTrue 的前提是异步模式源码中对此有强制校验——若wait_iterate_all is True and sync即syncTrue同步模式会直接抛出ValueErrorParam wait_iterate_all can be True only when sync is Falsepython/asc/language/adv/matmul.py。输出到 TSCMLocalTensor/UB时不支持该参数当iterate_all的tensor不是GlobalTensor时传入wait_iterate_all会抛出ValueErrorWhen iterate_all output to TSCM, param wait_iterate_all is not supported.即该同步机制仅在 GM 输出场景生效python/asc/language/adv/matmul.py。四、iterate_all 参数详解wait_iterate_all 的使用上下文要正确使用wait_iterate_all必须先理解iterate_all的完整参数体系。iterate_all一次调用会计算出singleCoreM * singleCoreN大小的 C 矩阵其签名如下asc.language.adv.Matmul.iterate_all.mdMatmul.iterate_all(tensor: BaseTensor, en_atomic: int 0, sync: bool True, en_sequential_write: bool | None None, wait_iterate_all: bool | None None, fake_msg: bool | None None) → None参数说明默认值tensorC 矩阵类型为GlobalTensor或LocalTensor必填en_atomic是否开启 Atomic 操作可选值为 0、1、2、30sync设置同步或异步模式Trueen_sequential_write是否开启连续写模式仅支持输出到 GM 场景NoneGM 场景下实际为Falsewait_iterate_all是否需要通过wait_iterate_all接口等待iterate_all执行结束仅支持异步输出到 GM 场景NoneGM 场景下实际为Falsefake_msg仅在 IBShare 场景和 IntraBlockPartSum 场景使用仅支持输出到 GM 场景NoneGM 场景下实际为False从源码看当tensor为GlobalTensor时en_sequential_write、wait_iterate_all、fake_msg三个可选参数若未显式传入会分别被置为Falseen_atomic的合法取值被check_type(en_atomic, [0, 1, 2, 3])约束tensor.dtype支持int32、intint64、floatfp32、float32、halffp16、float16、int8python/asc/language/adv/matmul.py。此外iterate_all要求传入的 C 矩阵地址空间大小不小于single_core_m * single_core_n个元素否则输出会越界。五、源码级实现剖析wait_iterate_all 的 IR 生成链路pyasc 的算子编程接口本质上是对 Ascend C 的 Python 化封装用户书写的 Python 代码会被编译为 Asc IR再经后续 Pass 与 CodeGen 翻译为昇腾设备代码。wait_iterate_all的实现非常简洁它直接向 IR Builder 请求创建一个MatmulWaitIterateAll操作python/asc/language/adv/matmul.pyrequire_jit set_matmul_docstring(api_namewait_iterate_all) def wait_iterate_all(self) - None: global_builder.get_ir_builder().create_asc_MatmulWaitIterateAllOp(self.to_ir())几点值得注意require_jit装饰器该方法只能在 JIT 编译上下文中调用如asc.jit修饰的算子函数内表明它属于设备端代码生成接口而非宿主侧运行接口。set_matmul_docstring(api_namewait_iterate_all)接口文档即本文所依托的 asc.language.adv.Matmul.wait_iterate_all.md由 docstring 模板自动生成模板定义在 python/asc/language/adv/utils.py 的wait_iterate_all_docstring()其中包含了功能说明、原型、参数与示例并在文件末尾以wait_iterate_all: wait_iterate_all_docstring的形式注册到 API 文档映射表。等待语义落在 IR 层create_asc_MatmulWaitIterateAllOp生成一个以该Matmul对象为操作数的等待操作对应 Ascend C 端WaitIterateAll()的同步屏障语义与之对称的是iterate_all生成的create_asc_MatmulIterateAllOppython/asc/language/adv/matmul.py二者构成发起-等待的异步配对。与iterate_all类似iterate_batch/iterate_n_batch也提供wait_iterate_batch配对接口生成create_asc_MatmulWaitIterateBatchOp而get_tensor_c异步场景则对应wait_get_tensor_c三套等待接口在 IR 层分别建模互不混淆。六、完整调用示例示例一文档给出的最小用法官方 API 文档asc.language.adv.Matmul.wait_iterate_all.md给出的调用示例mm asc.adv.Matmul(a_type, b_type, c_type, bais_type) mm.set_tensor_a(gm_a[offset_a:]) mm.set_tensor_b(gm_b[offset_b:]) if tiling.is_bias: mm.set_bias(gm_bias[offset_bias]) mm.iterate_all(tensorgm_c[offset_c], en_atomic0, syncFalse, en_sequential_writeFalse, wait_iterate_allTrue) mm.wait_iterate_all()要点iterate_all必须显式传入syncFalse否则wait_iterate_allTrue会触发源码中的ValueError校验tensor是带偏移切片gm_c[offset_c]的 GM 张量即连续输出到 GM 的起始位置wait_iterate_all()紧跟iterate_all之后调用作为本次异步迭代的同步屏障。示例二单元测试中的完整算子写法仓库单元测试test_iterate_allpython/test/unit/language/adv/test_matmul.py给出了一个可直接对照的、包含对象构造与 Tiling 初始化的完整内核asc.jit def kernel_iterate_all(a: asc.GlobalAddress, b: asc.GlobalAddress, c: asc.GlobalAddress, workspace: asc.GlobalAddress, quant_vector: asc.GlobalAddress) - None: pipe asc.TPipe() a_type asc.adv.MatmulType(positionasc.TPosition.GM, formatasc.CubeFormat.ND, dtypeasc.float16) b_type asc.adv.MatmulType(positionasc.TPosition.GM, formatasc.CubeFormat.ND, dtypeasc.float16) c_type asc.adv.MatmulType(positionasc.TPosition.GM, formatasc.CubeFormat.ND, dtypeasc.float16) matmul asc.adv.Matmul(a_type, b_type, c_type) asc.adv.register_matmul(pipe, workspace, matmul) m, n, k_a, k_b, k_c 512, 512, 512, 512, 512 single_core_m, single_core_n, single_core_k 256, 256, 256 base_m, base_n, base_k 128, 256, 64 tiling asc.adv.TCubeTiling(used_core_num24, mm, k_ak_a, k_bk_b, nn, base_mbase_m, base_kbase_k, base_nbase_n, single_core_msingle_core_m, single_core_ksingle_core_k, single_core_nsingle_core_n, depth_a11, depth_b11, step_m1, step_n1, share_mode0, share_ub_size0, share_l1_sizeasc.property(asc.TOTAL_L1_SIZE), share_l0c_sizeasc.property(asc.TOTAL_L0C_SIZE)) matmul.init(tiling) a_global asc.GlobalTensor() b_global asc.GlobalTensor() c_global asc.GlobalTensor() a_global.set_global_buffer(a) b_global.set_global_buffer(b) c_global.set_global_buffer(c) matmul.set_tensor_a(a_global) matmul.set_tensor_b(b_global) matmul.set_org_shape(m, n, k_a) matmul.set_single_shape(single_core_m, single_core_n, single_core_k) # 异步发起 GM 连续输出随后等待完成 matmul.iterate_all(c_global, en_atomic0, en_sequential_writeFalse, wait_iterate_allTrue, fake_msgFalse, syncFalse) matmul.wait_iterate_all() # 同一 Matmul 对象可继续向 UB 输出TSCM 场景不支持 wait_iterate_all 参数 ub_matrix asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) matmul.iterate_all(ub_matrix) matmul.end()该用例同时验证了两个关键点iterate_all(..., syncFalse, wait_iterate_allTrue)与wait_iterate_all()的异步配对合法否则会触发源码校验异常同一个Matmul对象在异步 GM 输出完成后仍可切换为向LocalTensorUB输出体现了 Matmul 对象复用能力最后通过end()释放 Matmul 计算资源。七、与其它同步接口的区分pyasc 的 Matmul 高阶 API 提供多个等待类接口适用场景各不相同完整清单见 docs/python-api/language/adv.md接口配套异步接口适用场景wait_iterate_alliterate_all(syncFalse, wait_iterate_allTrue)一次计算整片singleCoreM * singleCoreN的 C 矩阵结果连续输出到 GM 后等待完成wait_iterate_batchiterate_batch/iterate_n_batch异步调用批量多 BatchMatmul 计算结果连续输出到 GM 后等待完成wait_get_tensor_cget_tensor_c(..., syncFalse)将 C 矩阵分片从 GM 拷贝到 UB 且后续需要进行 Vector 计算时进行同步选择建议需要一次调用计算整核输出时使用iterate_allwait_iterate_all需要按baseM * baseN分片迭代获取 C 矩阵时使用iterateget_tensor_cwait_get_tensor_c需要多 Batch 批量计算时使用iterate_batch/iterate_n_batchwait_iterate_batch。八、使用注意事项小结配对使用wait_iterate_all仅当iterate_all处于异步模式syncFalse且wait_iterate_allTrue时使用同步模式下传wait_iterate_allTrue会抛ValueError。仅限 GM 连续输出当iterate_all输出目标为LocalTensorTSCM/UB时不得传入wait_iterate_all参数此时无需也无法通过本接口等待。输出空间充足确保tensor指向的 GM 地址空间不小于single_core_m * single_core_n个元素。对象复用与释放等待完成后可复用同一Matmul对象继续计算多个 Matmul 对象切换时务必在合适的时机调用end()释放计算资源防止资源冲突。适用平台本文所述行为以当前仓库源码与测试为准接口能力受昇腾硬件平台与 CANN 版本约束实际使用时请以目标环境的运行结果为准。通过本文你已完整掌握 pyasc 中wait_iterate_all的签名、约束、底层 IR 生成机制与标准写法可以放心地在昇腾 AI 处理器上编写异步 Matmul 算子实现发起异步计算 → 连续输出 GM → 显式等待完成的高效流水。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表