Codegen
Codegen is used to parse the ImplGraph generated by Schedule and generate executable code, including the code to be run on the host and device (kernel code).
- Host code: is executed on the host and includes the following content:
- tiling_func: divides a compute task into smaller blocks (tiling) for efficient execution on the device.
- infer_shape: infers the shape of each tensor to ensure the correctness of subsequent computation.
- get_kernel: obtains and generates the kernel code on the device.
- Kernel code: is executed on the device and includes the following content:
- kernel: defines the specific computation logic.
- tiling_data: describes the structure of how to divide and process data.
The input of Codegen is the ImplGraph generated by Schedule. The graph contains the basic elements for generating code, including shape information, node information, and axis information. The following figure shows an example of ImplGraph.
The detailed information in the figure is as follows:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | // Size describes the computation amount. Sizes: z0z1t_size: VAR z0z1Tb_size: VAR // Axis describes the axis information, including the size, type, and alignment mode. 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 indicates the original axis that has not been tiled. 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 indicates the inner axis of the UB tiling. 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 indicates the inner axis of the inter-kernel tiling. // Nodes: Each node generates the API call in the kernel code. 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, } // Vectorized axis: represents the axis corresponding to the amount of data computed in the 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) ...... |
The following sample code shows the process of generating the DataCopy API by Codegen based on the parsing result of ImplGraph. The core process covers parsing the tiling strategy on the graph, assembling the input parameters of the API, and generating the code to be executed on the device.
1 2 3 | 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; |
Similar to the traditional operator code, the Schedule module generates multiple templates for different tiling strategies, corresponding to different ImplGraphs. The Codegen needs to parse these templates in sequence, generate template functions, and call different template functions based on the tiling key at the kernel function entry. The following is an example:
1 2 3 4 5 6 7 8 | 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); } } |
Below is an example of the AutofuseTilingData structure. The values of member variables, including tiling_key, are computed by the ATT.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 | 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; |
The generated for loop and the amount of data processed each time vary depending on the template function. The following is an example:
- tiling_key=0
The z1 axis is tiled, and the vectorized axis for each computation in the UB is z1t.
1 2 3
for (int z0z1Tb = 0; z0z1Tb < z0z1Tb_loop_size; z0z1Tb++) { Abs(y_local[0], xlocal[0], z1t_actual_size); }
- tiling_key=1
The z0 axis is tiled, and the vectorized axes for each computation in the UB are z0t and z1.
1 2 3
for (int z0Tb = 0; z0Tb < z0Tb_loop_size; z0Tb++) { Abs(y_local[0], xlocal[0], z0t_actual_size * z1_axis_size); }
In general, Codegen generates kernel code based on the ImplGraph. However, to continuously improve the performance of kernel code, Codegen framework optimization and internal API optimization need to be continuously explored.