TILING_KEY_LIST

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	x
Atlas inference product AI Core	x
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

Determines within the kernel function whether the currently executed TilingKey matches the specified TilingKey configured on the host, thereby identifying the branch logic that meets the condition of TilingKey == key1 or TilingKey == key2.

Prototype

TILING_KEY_LIST(key1,key2)

Parameters

Parameter	Input/Output	Description
key	Input	Branch of a kernel function. The value must be a non-negative integer.

Restrictions

TILING_KEY_LIST applies only to the if and else if branches and does not support the else branch. That is, using the TILING_KEY_LIST function to represent N branches requires N separate instances of TILING_KEY_LIST(key1,key2).
Two TilingKey values can be passed. Each TilingKey must be unique.
When this API is used, the default kernel type must be set. You can also configure the kernel type for a specific TilingKey, which will overwrite the default kernel type. The kernel type can only be set to KERNEL_TYPE_MIX_AIC_1_1 or KERNEL_TYPE_MIX_AIC_1_2.
Currently, the kernel launch project is not supported.

Example

extern "C" __global__ __aicore__ void add_custom(__gm__ uint8_t *x, __gm__ uint8_t *y, __gm__ uint8_t *z, __gm__ uint8_t *workspace, __gm__ uint8_t *tiling)
{
    GET_TILING_DATA(tilingData, tiling);
    if (workspace == nullptr) {
        return;
    }
    KernelAdd op;
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_1);
    op.Init(x, y, z, tilingData.numBlocks, tilingData.totalLength, tilingData.tileNum);
    // When TilingKey is 1 or 2, execute Process1. When TilingKey is 3 or 4, execute Process2.
    if (TILING_KEY_LIST(1,2)) {
        op.Process1();
    } else if (TILING_KEY_LIST(3,4)) {
        KERNEL_TASK_TYPE(4, KERNEL_TYPE_MIX_AIC_1_2);
        op.Process2();
    } 
    // Other code logic
    ...
    // In this example, when TilingKey is 3 or 4, ProcessOther is executed.
    if (TILING_KEY_LIST(3,4)) {
        op.ProcessOther();
    }
}

Example of a matching tiling function on the host (pseudocode):

ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    // Other code logic
    ...
    if (context->GetInputShape(0) > 10) {
        context->SetTilingKey(1);
    } else if (some condition) {
        context->SetTilingKey(2);
    } else if (some condition) {
        context->SetTilingKey(3);
    } else if (some condition) {
        context->SetTilingKey(4);
    }
}

Parent topic: Kernel Tiling