Documentation

TILING_KEY_IS

Function Usage

Checks whether the tilingKey in the current kernel function execution is equal to a specific key, so as to identify a kernel branch with tilingKey==key.

Prototype

1
TILING_KEY_IS(key)

Parameters

Parameter

Input/Output

Description

key

Input

The parameter value is a positive number, indicating a branch of a kernel function.

Restrictions

  • The key value must be a non-negative integer.
  • TILING_KEY_IS applies only to the if and else if branches, and does not support the else branch. That is, using the TILING_KEY_IS function to represent N branches requires N TILING_KEY_IS (keys).
  • Currently, the kernel launch project is not supported.

Example

 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
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;
    op.Init(x, y, z, tilingData.blockDim, tilingData.totalLength, tilingData.tileNum);
    // When TilingKey is 1, Process1 is executed. When TilingKey is 2, Process2 is executed. When TilingKey is 3, Process3 is executed.
    if (TILING_KEY_IS(1)) {
        op.Process1();
    } else if (TILING_KEY_IS(2)) {
        op.Process2();
    } else if (TILING_KEY_IS(3)) {
        op.Process3();
    }
    // Other code logic
    ...
    // In this example, when TilingKey is 3, ProcessOther is executed.
    if (TILING_KEY_IS(3)) {
        op.ProcessOther();
    }
}
// Example of a matching tiling function on the host:
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);
    }
}
Parent topic: Kernel Tiling