CANN/ge自动融合代码生成

Codegen【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geCodegen的功能是解析Schedule生成的ImplGraph生成实际可执行的代码包括在Host侧Host Code和Device侧Kernel Code运行的代码其中Host CodeHost侧执行的代码包括如下内容tiling_func负责将计算任务划分为更小的块tiling便于在Device上高效执行。infer_shape推断各个tensor的形状确保后续计算的正确性。get_kernel获取并生成Device侧的kernel代码。Kernel CodeDevice侧执行的kernel代码包括如下内容kernel具体的计算逻辑。tiling_data用于描述如何划分和处理数据的结构。Codegen的输入是Schedule生成的ImplGraph该图包含了shape信息、节点信息、轴信息等生成代码的基本要素。ImplGraph示例如下图 1ImplGraph示例 图上的具体信息如下所示// size描述计算量大小 Sizes: z0z1t_size: VAR z0z1Tb_size: VAR // axis描述轴信息包括大小、类型、对齐方式等 Axis: z0(0) : 200, ORIGINAL, align: -1, allow_oversize_axis: 0, allow_unaligned_tail: 1 z1(1) : 200, ORIGINAL, align: -1, allow_oversize_axis: 0, allow_unaligned_tail: 1 z0z1(2) : 40000, ORIGINAL, align: -1, allow_oversize_axis: 0, allow_unaligned_tail: 1 // ORIGINAL代表未切分的原始轴 z0z1T(3) : Ceiling((40000 / (z0z1t_size))), TILE_OUT, from: {z0z1, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 z0z1t(4) : z0z1t_size, TILE_IN, from: {z0z1, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 // TILE_IN代表UB切分的内轴 z0z1TB(5) : Ceiling((Ceiling((40000 / (z0z1t_size))) / (z0z1Tb_size))), BLOCK_OUT, from: {z0z1T, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 z0z1Tb(6) : z0z1Tb_size, BLOCK_IN, from: {z0z1T, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 // BLOCK_IN代表核间切分的内轴 // node每个节点会生成kernel代码中的api调用 Nodes: ...... abs_test/gather_0: Load (1) ...... abs_test/abs_0: Abs (2) .axis {z0z1TB, z0z1Tb, z0z1t, } .loop_axis z0z1Tb .api: .compute_type elewise .type Compute .unit Vector .x abs_test/gather_0.y .y.dtype float32 .y.axis {z0z1TB, z0z1Tb, z0z1t, } .y.repeats {(40000 / (z0z1Tb_size * z0z1t_size)), z0z1Tb_size, z0z1t_size, } .y.strides {(z0z1Tb_size * z0z1t_size), z0z1t_size, 1, } .y.vectorized_axis {z0z1t, } // 向量化轴代表UB内计算的数据量对应的轴 .y.vectorized_strides {1, } .y.mem: .tensor_id 3 .alloc_type Queue .hardware UB .position TPosition::VECOUT .y.que: .id 1 .depth 2 .buf_num 2 .reuse_id 1 abs_test/store: Store (3) ......根据ImplGraph解析以下示例代码展示了Codegen生成DataCopy API的过程可以看出核心流程包括解析图上的切分策略组装API的入参生成在Device侧执行的代码。ss DataCopyPadExtend( ub , gm [ gm_offset tpipe.tiler.Size(api_attr.offset) ], dma_param.block_count , dma_param.block_len , dma_param.src_stride , dma_param.dst_stride ); std::endl;与传统算子代码类似针对不同的切分策略Schedule会生成多个模板对应不同的ImplGraphCodegen需要依次解析这些模板生成模板函数并在核函数入口处根据tiling_key调用不同的模板函数。示例如下extern C __global__ __aicore__ void abs_test(GM_ADDR abs_test_Data_0, GM_ADDR abs_test_Output_0, GM_ADDR workspace, AutofuseTilingData param) { const AutofuseTilingData t; if (t.tiling_key 0) { abs_test_0_general_0_nil_2_nil(abs_test_Data_0, abs_test_Output_0, workspace, t); } else if (t.tiling_key 1) { abs_test_0_general_0_nil_2_nil_unaligned(abs_test_Data_0, abs_test_Output_0, workspace, t); } }上面AutofuseTilingData结构体示例如下其中的成员变量具体值包括tiling_key等由ATT计算给出BEGIN_TILING_DATA_DEF_T(AutofuseTilingData) const uint32_t block_dim 40; const uint32_t corenum 0; const uint32_t ub_size 196352; const uint32_t hbm_size 0; const uint32_t tiling_key 0; const uint32_t z0z1z2z3t_size 17968; const uint32_t z0z1z2z3Tb_size 57; const uint32_t q0_size 96; const uint32_t q1_size 35936; const uint32_t q2_size 35936; const uint32_t b0_size 35936; const uint32_t tmp_tbuf_size 8192; END_TILING_DATA_DEF_T;不同的模板函数生成的for循环以及每次处理的数据量各不相同示例如下tiling_key0对z1轴进行切分每次在UB中完成计算的向量化轴为z1t。for (int z0z1Tb 0; z0z1Tb z0z1Tb_loop_size; z0z1Tb) { Abs(y_local[0], xlocal[0], z1t_actual_size); }tiling_key1对z0轴进行切分每次在UB中完成计算的向量化轴为z0t以及z1。for (int z0Tb 0; z0Tb z0Tb_loop_size; z0Tb) { Abs(y_local[0], xlocal[0], z0t_actual_size * z1_axis_size); }总体而言Codegen是依据ImplGraph生成kernel代码的但要持续提升kernel代码的性能仍需在Codegen框架优化和API内部优化两方面不断挖掘。【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考