CANN/asc-devkit: SoftMax Tiling使用说明

SoftMax Tiling使用说明【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkitAscend C提供一组SoftMax Tiling API方便用户获取SoftMax kernel计算时所需的Tiling参数。获取Tiling参数主要分为如下两步获取SoftMax接口计算所需最小和最大临时空间大小注意该步骤不是必须的只是作为一个参考供合理分配计算空间。获取输入SoftMax kernel侧接口所需tiling参数需要传入输入shape、剩余的可供softmax接口计算的空间大小和计算的数据类型大小。SoftMax Tiling结构体的定义如下开发者无需关注该tiling结构的具体信息只需要传递到kernel侧传入SoftMax高阶API接口直接进行使用即可。struct SoftMaxTiling { uint32_t srcM 0; uint32_t srcK 0; uint32_t srcSize 0; uint32_t outMaxM 0; uint32_t outMaxK 0; uint32_t outMaxSize 0; uint32_t splitM 0; uint32_t splitK 0; uint32_t splitSize 0; uint32_t reduceM 0; uint32_t reduceK 0; uint32_t reduceSize 0; uint32_t rangeM 0; uint32_t tailM 0; uint32_t tailSplitSize 0; uint32_t tailReduceSize 0; };对于SoftMax/SimpleSoftMax请参考SoftMax/SimpleSoftMax Tiling对于SoftmaxFlash请参考SoftmaxFlash Tiling接口对于SoftmaxGrad请参考SoftmaxGrad Tiling接口对于SoftmaxFlashV2请参考SoftmaxFlashV2 Tiling接口判断SoftMaxTiling是否为基本块Tiling请参考IsBasicBlockInSoftMax。调用示例如下样例介绍了使用SoftMax高阶API时host侧获取Tiling参数的流程以及该参数如何在kernel侧使用。样例中输入Tensor的shape大小为[320,64]输入的数据类型为half。将SoftMaxTiling结构体参数增加至TilingData结构体作为TilingData结构体的一个字段。BEGIN_TILING_DATA_DEF(TilingData) // 注册一个tiling的类以tiling的名字作为入参 TILING_DATA_FIELD_DEF(uint32_t, totalLength); // 添加tiling字段总计算数据量 TILING_DATA_FIELD_DEF(uint32_t, tileNum); // 添加tiling字段每个核上总计算数据分块个数 ... // 添加其他tiling字段 TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxTilingData); // 将SoftMaxTiling结构体参数增加至TilingData结构体 END_TILING_DATA_DEF;Tiling实现函数中首先调用GetSoftMaxMaxTmpSize/GetSoftMaxMinTmpSize接口获取SoftMax接口能完成计算所需最大/最小临时空间大小根据该范围结合实际的内存使用情况设置合适的空间大小然后根据输入shape、剩余的可供计算的空间大小等信息获取SoftMax kernel侧接口所需tiling参数。namespace optiling { const uint32_t NUM_BLOCKS 8; const uint32_t TILE_NUM 8; static ge::graphStatus TilingFunc(gert::TilingContext* context) { TilingData tiling; uint32_t totalLength context-GetInputTensor(0)-GetShapeSize(); context-SetBlockDim(NUM_BLOCKS); tiling.set_totalLength(totalLength); tiling.set_tileNum(TILE_NUM); // 设置其他Tiling参数 ... std::vectorint64_t shapeVec {320,64}; ge::Shape srcShape(shapeVec); // 本样例中仅作为样例说明通过GetSoftMaxMinTmpSize获取最小值并传入来保证功能正确开发者可以根据需要传入合适的空间大小 const uint32_t localWorkSpaceSize AscendC::GetSoftMaxMinTmpSize(srcShape, sizeof(half), false); // 获取SoftMax Tiling参数 AscendC::SoftMaxTilingFunc(srcShape, sizeof(half), localWorkSpaceSize, tiling.softmaxTilingData); ... // 其他逻辑 tiling.SaveToBuffer(context-GetRawTilingData()-GetData(), context-GetRawTilingData()-GetCapacity()); context-GetRawTilingData()-SetDataSize(tiling.GetDataSize()); context-SetTilingKey(1); return ge::GRAPH_SUCCESS; } } // namespace optiling对应的kernel侧通过在核函数中调用GET_TILING_DATA获取TilingData继而将TilingData中的SoftMax Tiling信息传入SoftMax接口参与计算。完整的kernel侧样例请参考调用示例。extern C __global__ __aicore__ void func_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) { GET_TILING_DATA(tilingData, tiling); KernelFunc op; op.Init(x, y, z, tilingData.totalLength, tilingData.tileNum,tilingData.SoftMaxTiling); if (TILING_KEY_IS(1)) { op.Process(); } }【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考