CANN/pyasc矩阵乘加API文档

asc.language.basic.mmad

【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc

asc.language.basic.mmad(dst: LocalTensor, fm: LocalTensor, filter: LocalTensor, params: MmadParams) → None

asc.language.basic.mmad(dst: LocalTensor, fm: LocalTensor, filter: LocalTensor, bias: LocalTensor, params: MmadParams) → None

完成矩阵乘加（C += A * B）操作。矩阵ABC分别为A2/B2/CO1中的数据。 ABC矩阵的数据排布格式分别为ZZ，ZN，NZ。

对应的 Ascend C 函数原型

template <typename T, typename U, typename S> __aicore__ inline void Mmad(const LocalTensor<T>& dst, const LocalTensor<U>& fm, const LocalTensor<S>& filter, const MmadParams& mmadParams)

template <typename T, typename U, typename S, typename V> __aicore__ inline void Mmad(const LocalTensor<T>& dst, const LocalTensor<U>& fm, const LocalTensor<S>& filter, const LocalTensor<V>& bias, const MmadParams& mmadParams)

参数说明

• dst：结果输出 Tensor，类型为 LocalTensor。
  • 用于存放矩阵乘累加的结果。
  • 必须位于 CO1 存储位置（TPosition.CO1）。
  • 元素数据类型需与累加结果类型匹配。
• fm：左矩阵（A 矩阵）输入，类型为 LocalTensor。
  • 表示矩阵乘法中的左操作数。
  • 必须位于 A2 存储位置（TPosition.A2）。
  • 需要按照满足 Mmad 格式要求的 A2 布局存储。
• filter：右矩阵（B 矩阵）输入，类型为 LocalTensor。
  • 表示矩阵乘法中的右操作数。
  • 必须位于 B2 存储位置（TPosition.B2）。
  • 需要按照符合指令格式的 B2 分块布局排布。
• bias（可选）：偏置项，类型为 LocalTensor。
  • 用于执行 dst += fm × filter + bias 的计算。
  • 当提供 bias 时，将使用带偏置版本的指令。
• params：MmadParams 类型的矩阵乘参数。
  • m：左矩阵Height，取值范围：m∈[0, 4095] 。默认值为0。
  • n：右矩阵Width，取值范围：n∈[0, 4095] 。默认值为0。
  • k：左矩阵Width、右矩阵Height，取值范围：k∈[0, 4095] 。默认值为0。
  • cmatrixInitVal：配置C矩阵初始值是否为0。默认值true。
  • cmatrixSource：配置C矩阵初始值是否来源于C2（存放Bias的硬件缓存区）。默认值为false。
  • isBias：该参数废弃，新开发内容不要使用该参数。
  • fmOffset：预留参数。
  • enSsparse：预留参数。
  • enWinogradA：预留参数。
  • enWinogradB：预留参数。
  • unitFlag：预留参数。
  • kDirectionAlign：预留参数。

约束说明

• dst只支持位于CO1，fm只支持位于A2，filter只支持位于B2。
• 当M、K、N中的任意一个值为0时，该指令不会被执行。
• 当M = 1时，会默认开启GEMV（General Matrix-Vector Multiplication）功能。在这种情况下，Mmad API从L0A Buffer读取数据时，会以ND格式进行读取，而不会将其视为ZZ格式。所以此时左矩阵需要直接按照ND格式进行排布。
• 操作数地址对齐要求请参见 《Ascend C算子开发接口》 中的“通用说明和约束-通用地址对齐约束”。

调用示例

• 基本 Mmad（无 bias）

  @asc.jit def kernel_mmad_basic(): dst = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.CO1, addr=0, tile_size=1024) fm = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.A2, addr=0, tile_size=1024) filter = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.B2, addr=0, tile_size=1024) params = asc.MmadParams(4, 4, 4) asc.mmad(dst, fm, filter, params)

• Mmad 带 bias

  @asc.jit def kernel_mmad_bias(): dst = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.CO1, addr=0, tile_size=1024) fm = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.A2, addr=0, tile_size=1024) filter = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.B2, addr=0, tile_size=1024) bias = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECIN, addr=0, tile_size=1024) params = asc.MmadParams(4, 4, 4) asc.mmad(dst, fm, filter, bias, params)

【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc