CANN/pypto设置立方体切片形状

pypto.set_cube_tile_shapes【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas A3 训练系列产品/Atlas A3 推理系列产品√功能说明在调用pypto.matmul或pypto.scaled_mm前必须调用本接口设置矩阵运算的切分大小具体切分配置可参考Matmul高性能编程。函数原型set_cube_tile_shapes(m: List[int], k: List[int], n: List[int], enable_split_k: bool False) - None参数说明参数名输入/输出说明m输入m维度在L0和L1上的TileShape切片形状的切分大小分别对应mL0和mL1的切分大小k输入k维度在L0和L1上的TileShape切片形状的切分大小分别对应kL0和kL1的切分大小n输入n维度在L0和L1上的TileShape切片形状的切分大小分别对应nL0和nL1的切分大小enable_split_k输入设置True表示使能matmul的多核切K功能False表示未使能多核切K默认为False返回值说明void约束说明对齐约束通用对齐约束要求mL0、mL1、kL0、kL1、nL0、nL1均满足32字节对齐DT_FP32输入场景要求满足16元素对齐。例如输入矩阵的数据类型为DT_FP16时kL0 * sizeof(DT_FP16) % 32 0。基础关系约束约束项要求mL0 与 mL1mL0 0且mL0 ≤ mL1且mL1 % mL0 0kL0 与 kL1kL0 0且kL0 ≤ kL1且kL1 % kL0 0nL0 与 nL1nL0 0且nL0 ≤ nL1且nL1 % nL0 0ND格式特有约束当A矩阵在format为ND且转置场景时即数据排布为[K, M]要求mL0满足32字节对齐。NZ格式特有约束A、B矩阵在format为NZ场景时要求外轴切分大小满足16元素对齐内轴切分大小满足32字节对齐。例如在A矩阵非转置场景外轴为M、内轴为K要求mL0、mL1满足16元素对齐kL0、kL1满足32字节对齐。空间约束输入dtype为DT_FP16或DT_BF16或DT_FP32CeilAlign(mL0, 16) × CeilAlign(kL0, 16) × sizeof(aDtype) ≤ L0A_size CeilAlign(nL0, 16) × CeilAlign(kL0, 16) × sizeof(bDtype) ≤ L0B_size CeilAlign(mL0, 16) × CeilAlign(nL0, 16) × sizeof(cDtype) ≤ L0C_size CeilAlign(mL1, 16) × CeilAlign(kL1, 16) × sizeof(aDtype) CeilAlign(nL1, 16) × CeilAlign(kL1, 16) × sizeof(bDtype) ≤ L1_size输入dtype为DT_INT8或DT_FP8E5M2或DT_FP8E4M3或DT_HF8CeilAlign(mL0, 32) × CeilAlign(kL0, 32) × sizeof(aDtype) ≤ L0A_size CeilAlign(nL0, 32) × CeilAlign(kL0, 32) × sizeof(bDtype) ≤ L0B_size CeilAlign(mL0, 32) × CeilAlign(nL0, 32) × sizeof(cDtype) ≤ L0C_size CeilAlign(mL1, 32) × CeilAlign(kL1, 32) × sizeof(aDtype) CeilAlign(nL1, 32) × CeilAlign(kL1, 32) × sizeof(bDtype) ≤ L1_sizeBias空间约束bias数据到达BTBuffer全部转为fp32需满足以下约束nL0 × 4 ≤ BTBuffer_sizeFixPipe空间约束scaleTensor数据为uint64_t需满足以下约束nL0 × 8 ≤ FixBuffer_size其中aDtype、bDtype为输入矩阵数据类型cDtype为输出矩阵数据类型当输入为DT_INT8时cDtype为DT_INT32其余场景cDtype为DT_FP32CeilAlign(value, align)元素对齐实现为(value align - 1) / align * align多核切K约束仅支持2维矩阵3维/4维矩阵不支持多核切K多核切K场景只支持out_dtype数据类型为DT_FP32或DT_INT32。Bias、FixPipe反量化场景不支持叠加多核切K功能。调用示例# 基本配置 pypto.set_cube_tile_shapes([128, 128], [128, 128], [128, 128]) # 启用多核切K pypto.set_cube_tile_shapes([128, 128], [64, 256], [128, 128], enable_split_kTrue)【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考