Tiling Template Programming

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

If multiple TilingKeys are involved, you rely on these TilingKeys to manage the kernel implementation, which is complex in both management and usage. To simplify this process, template programming instead of traditional TilingKey programming can be used to reduce the dependency on TilingKey values and achieve more intuitive and efficient kernel management. The procedure is as follows:

In the op_kernel directory of a custom operator project, add a header file that contains the template argument declaration and template argument selection. 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 argument declaration ASCENDC_TPL_ARGS_DECL and template argument selection ASCENDC_TPL_ARGS_SEL (available template). For details about the API, see Template Argument Declaration.

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

// Template argument 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 argument declaration for DataType: data type of argument 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 argument declaration for DataType: data type of argument 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 argument declaration for DataType: data type of argument 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 argument declaration of the custom UINT type: The template arguments include the number of tiles, encoding bit width ASCENDC_TPL_8_BW (8 bits), which indicates that the number of template arguments 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 argument declaration of the custom bool type: The template argument is a flag bit indicating whether to split. 1: split; 0: not split.
);

// Template argument selection
// 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 argument 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 argument declaration and template argument selection in step 1.
Call ASCENDC_TPL_SEL_PARAM to automatically generate and configure a TilingKey. The input arguments of ASCENDC_TPL_SEL_PARAM are the specific values of template arguments, and must be in the same sequence as the template arguments in the header file that defines the template argument declaration and template argument selection.

       
        
          
          
            #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(BLOCK_DIM);
    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;
}

           

         

       
      

Implement on the kernel.

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

       
        
          
          
            #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();
        }
    }
}

           

         

       
      

For details, refer to the tiling template programming sample.

Parent topic: Tiling Implementation on the Host