CANN/asc-devkit矩阵切K轴API

📅 2026/7/5 18:59:02 👁️ 阅读次数 📝 编程学习
CANN/asc-devkit矩阵切K轴API

SetSplitK

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit

功能说明

EnableMultiCoreSplitK接口功能与该接口相同,建议使用EnableMultiCoreSplitK。

多核场景,通过该接口开启切K轴。不调用该接口的情况下,默认不切K轴。在GetTiling接口调用前使用。

函数原型

void SetSplitK(bool flag)

参数说明

表1参数说明

参数名输入/输出描述
flag输入是否开启切K轴。

true:开启切K轴

false:不开启切K轴

返回值说明

约束说明

  • 如果在算子中使用该接口,获取C矩阵结果时仅支持输出到Global Memory。
  • 如果在算子中使用该接口,需在Kernel侧代码中首次将C矩阵分片的结果写入Global Memory之前,先清零Global Memory,随后在获取C矩阵分片的结果时,再开启AtomicAdd累加。如果不预先清零Global Memory,可能会因为累加Global Memory中原始的无效数据而产生精度问题。

调用示例

auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()) matmul_tiling::MultiCoreMatmulTiling tiling(ascendcPlatform); tiling.SetDim(useCoreNums); tiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); tiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); tiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT); tiling.SetBiasType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT); tiling.SetShape(M, N, K); tiling.SetOrgShape(M, N, K); tiling.SetBias(true); tiling.SetBufferSpace(-1, -1, -1); tiling.SetSplitK(true); optiling::TCubeTiling tilingData; int ret = tiling.GetTiling(tilingData);

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考