pyasc Matmul.set_tensor_b 详解设置矩阵乘右矩阵 B 的三种方式与底层实现【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc导读本文围绕 CANN pyasc 项目面向昇腾 AI 处理器的 Python 算子编程接口中的高阶矩阵乘 APIasc.language.adv.Matmul.set_tensor_b展开完整讲解其两种函数重载、与 Ascend C 原生SetTensorB的对应关系、参数语义、约束条件及调用流程。读完本文你将掌握在 pyasc 中为 Matmul 对象设置右矩阵 B 的正确姿势理解标量/GlobalTensor/LocalTensor 三种形态的适用场景并能结合 matmul_mix.py 示例完成一个可运行的 Matmul 算子编写。接口定位Matmul 高阶 API 中的右矩阵设置入口pyasc 为 Python 用户提供与 Ascend C 一一对应的算子编程接口。在 asc.language.adv.Matmul 这一组高阶 API 中矩阵乘的计算公式为C A * B Bias其中右矩阵 B 的传入正是由set_tensor_b负责。它属于 Matmul 对象在正式迭代计算iterate_all/iterate/iterate_batch等之前必须完成的输入装配环节与左矩阵设置接口set_tensor_a对称存在。该接口的完整函数原型定义在 matmul.py 中通过overload声明了两种调用形态def set_tensor_b(self, scalar: int) - None: ... def set_tensor_b(self, tensor: BaseTensor, transpose: bool False) - None: ...标量形态当 B 矩阵是单值标量时使用将一个整数常量直接作为右矩阵参与计算张量形态当 B 矩阵是真实数据时使用可传入GlobalTensor全局内存或LocalTensor本地内存类型并通过transpose控制是否转置。对应的 Ascend C 函数原型set_tensor_b与 Ascend C 原生接口SetTensorB一一对应包含三个重载版本分别覆盖全局内存张量、本地内存张量和标量三种场景__aicore__ inline void SetTensorB(const GlobalTensorSrcBT gm, bool isTransposeB false)__aicore__ inline void SetTensorB(const LocalTensorSrcBT leftMatrix, bool isTransposeB false)__aicore__ inline void SetTensorB(SrcBT bScalar)对照可见Python 层的两个重载正是对这三个 C 重载的封装tensor参数在运行时根据实际传入的是GlobalTensor还是LocalTensor分别映射到前两个原型scalar参数映射到第三个原型transpose则对应 C 侧的isTransposeB默认均为false。参数说明参数类型必填说明scalarint二选一B 矩阵中设置的值为标量。仅用于右矩阵为常量的场景tensorBaseTensor二选一B 矩阵类型为GlobalTensor或LocalTensortransposebool否B 矩阵是否需要转置默认False两种重载二选一使用要么传scalar要么传tensor可搭配transpose不能同时传入。约束说明传入的 TensorB 地址空间大小需要保证不小于single_k * single_n以元素个数计。其中single_k、single_n是 Matmul 单核计算的 K、N 方向分片大小B 矩阵需要为每个核的分片计算提供完整的数据。数据类型受源码级校验约束。查看 matmul.py 的实现可知张量形态下check_type仅允许以下类型half、float即 float32、int8以及对应的float16/float32标量形态下则仅支持half、float、float16、float32matmul.py传入其他类型会抛出ValueError(Tensor type is not supported in set_tensor_b)。这也意味着 int8 量化场景下 B 矩阵必须走张量形态无法用标量形态表达。底层实现重载分发与 IR 构建从源码结构看set_tensor_b的 Python 实现采用OverloadDispatcher机制完成运行时重载分发matmul.pyrequire_jit set_matmul_docstring(api_nameset_tensor_b) def set_tensor_b(self, *args, **kwargs) - None: dispatcher OverloadDispatcher(__name__) builder global_builder.get_ir_builder() dispatcher.register(scalarRuntimeInt) def _(scalar: RuntimeInt): check_type(self.b_dtype, [KT.half, KT.float_, KT.float16, KT.float32], Scalar type is not supported in set_tensor_b) builder.create_asc_MatmulSetTensorBScalarOp(self.to_ir(), _mat(scalar, self.b_dtype).to_ir()) dispatcher.register(tensorBaseTensor, transposeDefaultValued(RuntimeBool, False)) def _(tensor: BaseTensor, transpose: RuntimeBool False): check_type(tensor.dtype, [KT.half, KT.float_, KT.int8, KT.float16, KT.float32], Tensor type is not supported in set_tensor_b) transpose _mat(transpose, KT.bit) builder.create_asc_MatmulSetTensorBOp(self.to_ir(), tensor.to_ir(), transpose.to_ir()) dispatcher(*args, **kwargs)关键点解读require_jit装饰器确保该方法仅在 JIT 编译上下文中被调用由编译器在运行时将 Python 调用降级为 Asc IR 指令。重载匹配OverloadDispatcher根据实参形态scalar或tensor自动选择对应的内部函数transpose参数通过DefaultValued(RuntimeBool, False)提供默认值。类型检查前置在构建 IR 之前先做check_type校验失败立即抛异常将错误拦截在编译期之前。IR 生成标量形态生成asc_MatmulSetTensorBScalarOp张量形态生成asc_MatmulSetTensorBOptranspose被转换为 bit 类型的 IR 值这些 Op 最终由后端发射为对应的 Ascend C 内核代码。调用示例与完整使用流程原文档给出的最小调用示例asc.adv.register_matmul(pipe, workspace, mm, tiling) mm.set_tensor_a(gm_a) mm.set_tensor_b(gm_b) # 设置右矩阵B mm.set_bias(gm_bias) mm.iterate_all(gm_c)在真实算子中set_tensor_b通常与多核切分逻辑配合。仓库示例 matmul_mix.py 给出了一个完整的 Matmul 内核写法其中 B 矩阵的设置流程如下asc.jit(always_compileTrue) def matmul_kernel(a: asc.GlobalAddress, b: asc.GlobalAddress, c: asc.GlobalAddress, tiling: asc.adv.TCubeTiling, workspace: asc.GlobalAddress): offset_a, offset_b, offset_c, tail_m, tail_n calc_offsets(tiling, IS_TRANS_A, IS_TRANS_B) a_global asc.GlobalTensor() b_global asc.GlobalTensor() c_global asc.GlobalTensor() a_global.set_global_buffer(a offset_a) b_global.set_global_buffer(b offset_b) # 为 B 矩阵绑定全局内存基址 片内偏移 c_global.set_global_buffer(c offset_c) pipe asc.TPipe() matmul asc.adv.Matmul( aasc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, a_global.dtype, IS_TRANS_A), basc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, b_global.dtype, IS_TRANS_B), casc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, c_global.dtype), ) asc.adv.register_matmul(pipe, workspace, matmul, tiling) if asc.get_block_idx() tiling.used_core_num: matmul.set_tensor_a(a_global, IS_TRANS_A) matmul.set_tensor_b(b_global, IS_TRANS_B) # 传入转置标志 matmul.set_tail(tail_m, tail_n) # 尾核调整分片 matmul.iterate_all(c_global) matmul.end() asc.pipe_barrier(asc.PipeID.PIPE_ALL)其中IS_TRANS_B即为set_tensor_b的transpose参数。若 B 需要转置还需在 Tiling 侧同步设置见 generate_tiling 中matmul_tiling.set_b_type(host.TPosition.GM, host.CubeFormat.ND, host.DataType.DT_FLOAT16, False)的最后一个布尔参数。与相邻接口的配合关系set_tensor_b不是一个孤立接口它在 Matmul 生命周期中的位置如下初始化先通过 register_matmul 完成 Matmul 对象初始化分离模式下需在init_buffer之前调用最多支持 4 个 Matmul 对象。装配输入依次调用set_tensor_a左矩阵、set_tensor_b右矩阵、set_bias可选偏置。调整分片如果当前是尾核还需调用 set_tail 重新设置single_core_m/single_core_n/single_core_k——注意 B 矩阵的地址空间约束正是基于single_k * single_n计算的因此set_tail调整 N/K 分片后需要确保 B 的内存范围仍然满足该约束。迭代计算调用 iterate_all一次计算single_core_m * single_core_n的 C 矩阵或iterate/iterate_batch等接口。资源释放多个 Matmul 对象切换时调用end()释放计算资源。常见错误与避坑提示类型不匹配向set_tensor_b传入int8以外的非支持类型或标量形态传入不支持的类型会直接抛出ValueError。请先确认MatmulType中 b 的 dtype 声明与实际传入张量一致。内存越界B 张量基址加偏移后的可读范围若小于single_k * single_n个元素将产生越界访问。多核场景下务必像 calc_offsets 那样按n_index * single_core_n转置时按n_index * k_b * single_core_n计算 B 的片内偏移。转置标志不一致set_tensor_b的transpose必须与 Tiling 侧set_b_type中的转置配置、以及 Matmul 构造时MatmulType(..., IS_TRANS_B)保持一致三者不一致会导致计算错误。总结asc.language.adv.Matmul.set_tensor_b是 pyasc 矩阵乘编程中装配右矩阵 B 的唯一入口支持标量、GlobalTensor、LocalTensor三种数据形态与可选的转置开关底层通过OverloadDispatcher分发并经类型校验后生成对应的 Asc IR 指令最终发射为 Ascend C 的SetTensorB内核代码。结合register_matmul→set_tensor_a→set_tensor_b→set_tail→iterate_all的标准调用链即可在昇腾 AI 处理器上完成一个完整、可验证的 Matmul 算子。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考