CANN/pyasc矩阵乘法N批处理迭代

asc.language.adv.Matmul.iterate_n_batch【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyascMatmul.iterate_n_batch(batch_loop: PlainValue | int, batch_a: PlainValue | int, batch_b: PlainValue | int, en_sequential_write: PlainValue | bool, matrix_stride_a: PlainValue | int 0, matrix_stride_b: PlainValue | int 0, matrix_stride_c: PlainValue | int 0, sync: PlainValue | bool True, wait_iterate_batch: PlainValue | bool False) → None调用一次IterateNBatch会进行N次IterateBatch计算计算出N个多Batch的singleCoreM * singleCoreN大小的C矩阵。 在调用该接口前需将MatmulConfig中的isNBatch参数设为true使能多Batch输入多Batch输出功能并调用SetWorkspace接口申请临时空间 用于缓存计算结果即IterateNBatch的结果输出至SetWorkspace指定的Global Memory内存中。 对于BSNGD、SBNGD、BNGS1S2的Layout格式 调用该接口之前需要在tiling中使用SetALayout/SetBLayout/SetCLayout/SetBatchNum设置A/B/C的Layout轴信息和最大BatchNum数 对于Normal数据格式则需使用SetBatchInfoForNormal设置A/B/C的M/N/K轴信息和A/B矩阵的BatchNum数。 实例化Matmul时通过MatmulType设置Layout类型当前支持3种Layout类型BSNGD、SBNGD、BNGS1S2。对应的Ascend C函数原型template bool sync true, bool waitIterateBatch false __aicore__ inline void IterateNBatch(const uint32_t batchLoop, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0)参数说明sync设置同步或者异步模式。wait_iterate_batch是否需要通过WaitIterateBatch接口等待IterateNBatch执行结束仅在异步场景下使用。batch_loop当前计算的BMM个数。batch_a当前单次BMM调用计算左矩阵的batch数。batch_b当前单次BMM调用计算右矩阵的batch数brc场景batchA/B不相同。en_sequential_write输出是否连续存放数据。matrix_stride_aA矩阵源操作数相邻nd矩阵起始地址间的偏移默认值是0。matrix_stride_bB矩阵源操作数相邻nd矩阵起始地址间的偏移默认值是0。matrix_stride_c该参数预留开发者无需关注。约束说明单BMM内计算遵循之前的约束条件。对于BSNGD、SBNGD、BNGS1S2 Layout格式输入A、B矩阵多Batch数据总和应小于L1 Buffer的大小。当使能MixDualMaster双主模式场景时即模板参数enableMixDualMaster设置为true不支持使用该接口。调用示例asc.jit def kernel_matmul_rpc_batch(a_gm: asc.GlobalAddress, b_gm: asc.GlobalAddress, c_gm: asc.GlobalAddress, bias_gm: asc.GlobalAddress, tiling: asc.adv.TCubeTiling, workspace_gm: asc.GlobalAddress, is_transpose_a_in: int, is_transpose_b_in: int, batch_a: int, batch_b: int): # 定义matmul type a_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, False, asc.LayoutMode.BSNGD) b_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, True, asc.LayoutMode.BSNGD) c_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float, False, asc.LayoutMode.BNGS1S2) bias_type asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float) a_global asc.GlobalTensor() size_a tiling.a_layout_info_b * tiling.a_layout_info_s * tiling.a_layout_info_n * tiling.a_layout_info_g * tiling.a_layout_info_d * 4 size_a tiling.b_layout_info_b * tiling.b_layout_info_s * tiling.b_layout_info_n * tiling.b_layout_info_g * tiling.b_layout_info_d * 4 size_bias tiling.c_layout_info_b * tiling.c_layout_info_n * tiling.c_layout_info_g * tiling.c_layout_info_s2 * 8 a_global set_global_buffer(a_gm, size_a) b_global set_global_buffer(b_gm, size_b) bias_global set_global_buffer(bias_gm, size_bias) tiling.share_mode 0 tiling.share_l1_size 512 * 1024 tiling.share_l0c_size 128 * 1024 tiling.share_ub_size 0 offset_a 0 offset_b 0 offset_c 0 offset_bias 0 a_global a_global[offset_a] b_global b_global[offset_b] bias_global bias_global[offset_bias] # 创建Matmul实例 mm asc.adv.Matmul(a_type, b_type, c_type, bias_type) pipe asc.Tpipe() asc.adv.register_matmul(pipe, mm) mm.init(tiling) mm.set_tensor_a(a_global, is_transpose_a_in) mm.set_tensor_b(b_global, is_transpose_b_in) g_lay tiling.a_layout_info_g if tiling.b_layout_info g_lay: g_lay tiling.b_layout_info_g for_extent tiling.a_layout_info_b * tiling.a_layout_info_n * g_lay / tiling.batch_num mm.set_workspace(c_global) mm.iterate_n_batch(for_extent, batch_a, batch_b, False)【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考