Tiling Template Programming

As described in TilingKey programming, TilingKey is hard to remember and understand because it is usually a long number without clear meaning.

If multiple TilingKeys are involved, you rely on the TilingKeys to manage the kernel implementation, which is complex in both management and usage. To simplify this process, you can use template programming instead of traditional TilingKey programming to reduce the dependency on TilingKey values and achieve more intuitive and efficient kernel management. The procedure is as follows. For details about the complete sample, see tiling template programming sample.

In the op_kernel folder of the custom operator project, add a header file that contains the template parameter definition and template parameter combination. In this example, the header file is named tiling_key_add_custom.h.

Include the template header file ascendc/host_api/tiling/template_argument.h.
Define the template parameter declaration ASCENDC_TPL_ARGS_DECL and template parameter combination ASCENDC_TPL_ARGS_SEL (that is, the template that can be used). For details about the APIs, see Template Argument Declaration.

       
        
          
          
            #include "ascendc/host_api/tiling/template_argument.h"

// Template parameter declaration
ASCENDC_TPL_ARGS_DECL(AddCustomTemplate, // Operator type
ASCENDC_TPL_DATATYPE_DECL(D_T_X, C_DT_FLOAT, C_DT_FLOAT16, ASCENDC_TPL_INPUT(0)),  // Template parameter definition for DataType: data type of parameter x. The value can be float16 or float32. ASCENDC_TPL_INPUT(0) indicates the 0th input on the kernel.
ASCENDC_TPL_DATATYPE_DECL(D_T_Y, C_DT_FLOAT, C_DT_FLOAT16, ASCENDC_TPL_INPUT(1)),  // Template parameter definition for DataType: data type of parameter y. The value can be float16 or float32. ASCENDC_TPL_INPUT(1) indicates the first input on the kernel.
ASCENDC_TPL_DATATYPE_DECL(D_T_Z, C_DT_FLOAT, C_DT_FLOAT16, ASCENDC_TPL_OUTPUT(0)), // Template parameter declaration for DataType: data type of parameter z. The value can be float16 or float32. ASCENDC_TPL_OUTPUT(0) indicates the 0th output on the kernel.
ASCENDC_TPL_UINT_DECL(TILE_NUM, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_MIX, 2, 0, 2, 3, 5, 10, 12, 13, 9, 8),// Template parameter definition of the custom UINT type (unsigned integer): The template parameters include the number of tiles, encoding bit width ASCENDC_TPL_8_BW (8 bits), which indicates that the number of template parameters does not exceed the 8-bit range, and ASCENDC_TPL_UI_MIX, which indicates that the value range is expressed in mixed mode. There are two groups of data {0–2} and {3–5} and exhaustive values 10, 12, 13, 9, and 8. The final result is {0, 1, 2, 3, 4, 5, 10, 12, 13, 9, 8}.
ASCENDC_TPL_BOOL_DECL(IS_SPLIT, 0, 1), // Template parameter definition of the custom bool type: The template parameter is a flag bit indicating whether to split. 1: split; 0: not split.
);

// Template parameter combination
// Used to check whether the TilingKey is valid when GET_TPL_TILING_KEY is called to obtain the TilingKey.
ASCENDC_TPL_SEL(
    ASCENDC_TPL_ARGS_SEL(
    ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY), // Kernel type selection, which does not need to be defined in the template parameter declaration and is directly configured in SEL. All ASCENDC_TPL_ARGS_SEL configurations must be consistent. If it is not configured, the automatic inference process will be triggered.
    ASCENDC_TPL_DATATYPE_SEL(D_T_X, C_DT_FLOAT16),
    ASCENDC_TPL_DATATYPE_SEL(D_T_Y, C_DT_FLOAT16),
    ASCENDC_TPL_DATATYPE_SEL(D_T_Z, C_DT_FLOAT16),
    ASCENDC_TPL_UINT_SEL(TILE_NUM, ASCENDC_TPL_UI_LIST, 1, 8),
    ASCENDC_TPL_BOOL_SEL(IS_SPLIT, 0, 1)
    ),
    ASCENDC_TPL_ARGS_SEL(
    ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
    ASCENDC_TPL_DATATYPE_SEL(D_T_X, C_DT_FLOAT),
    ASCENDC_TPL_DATATYPE_SEL(D_T_Y, C_DT_FLOAT),
    ASCENDC_TPL_DATATYPE_SEL(D_T_Z, C_DT_FLOAT),
    ASCENDC_TPL_UINT_SEL(TILE_NUM, ASCENDC_TPL_UI_LIST, 1, 8),
    ASCENDC_TPL_BOOL_SEL(IS_SPLIT, 0, 1)
    ),
);

           

         

       
      

On the host, call ASCENDC_TPL_SEL_PARAM to automatically generate and configure a TilingKey.

In the host implementation file, include the header file that defines the template parameter definition and template parameter combination in step 1.
Call ASCENDC_TPL_SEL_PARAM to automatically generate and configure a TilingKey. The input parameters of ASCENDC_TPL_SEL_PARAM are the specific values of template parameters, and must be in the same sequence as the template parameters in the header file that defines the template parameter definition and template parameter combination.

       
        
          
          
            #include "tiling_key_add_custom.h"
static ge::graphStatus TilingFunc(gert::TilingContext *context)
{
    TilingData tiling;
    uint32_t totalLength = context->GetInputShape(0)->GetOriginShape().GetShapeSize();
    ge::DataType dtype_x = context->GetInputDesc(0)->GetDataType();
    ge::DataType dtype_y = context->GetInputDesc(1)->GetDataType();
    ge::DataType dtype_z = context->GetOutputDesc(1)->GetDataType();
    uint32_t D_T_X = static_cast<int>(dtype_x), D_T_Y = static_cast<int>(dtype_y), D_T_Z = static_cast<int>(dtype_z), TILE_NUM = 1, IS_SPLIT = 0;
    if(totalLength< MIN_LENGTH_FOR_SPLIT){
        IS_SPLIT = 0;
        TILE_NUM = 1;
    }else{
        IS_SPLIT = 1;
        TILE_NUM = DEFAULT_TILE_NUM;
    }
    context->SetBlockDim(NUM_BLOCKS);
    tiling.set_totalLength(totalLength);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    ASCENDC_TPL_SEL_PARAM(context, D_T_X, D_T_Y, D_T_Z, TILE_NUM, IS_SPLIT);
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}

           

         

       
      

Perform implementation on the kernel.

In the kernel implementation file, include the header file that defines the template parameter definition and template parameter combination in step 1.
Add a template to the kernel function to support the input of template parameters. These parameters must be in the same sequence as those in the header file that defines the template parameter definition and template parameter combination.
Select different kernels based on the branch judgment of template parameters.

       
        
          
          
            #include "tiling_key_add_custom.h"
...
...
template<typename D_T_X, typename D_T_Y, typename D_T_Z, int TILE_NUM, int IS_SPLIT>
 __global__ __aicore__ void add_custom_template(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    KernelAdd<D_T_X, D_T_Y, D_T_Z> op;
    op.Init(x, y, z, tiling_data.totalLength, TILE_NUM);
    if constexpr (std::is_same_v<D_T_X, float> && std::is_same_v<D_T_Y, float> && std::is_same_v<D_T_Z, float>) {
        op.Process1();
    } else if constexpr (std::is_same_v<D_T_X, half> && std::is_same_v<D_T_Y, half> && std::is_same_v<D_T_Z, half>){
        if (IS_SPLIT == 0) {
            op.Process1();
        } else if(IS_SPLIT == 1) {
            op.Process2();
        }
    }
}

           

         

       
      

In the tiling template programming scenario, you can use the --kernel-template-input compilation option to specify that only the kernel code related to the specified template parameter combination, accelerating the compilation process.

Parent topic: Tiling Implementation on the Host