Conditions and Branches

Conditions and branches are used to implement conditional judgment in a program, so that different code logic is executed based on different conditions. The programming framework supports two types of conditions and branches:

• Static conditional branch: Configure conditional branches at compile time and generate fixed instructions for execution. Different kernels can be generated through multiple JIT compilations.
• Dynamic conditional branch: Determine conditions and branches at runtime and execute the corresponding functions.

Static Conditional Branch

# Generate kernel with input parameter add1_flag=False.
@pypto.frontend.jit
def add_kernel_false(
    input0: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    input1: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    output: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    val: int
):
    add_core(input0, input1, output, val, False)

# Generate kernel with input parameter add1_flag=True.
@pypto.frontend.jit
def add_kernel_true(
    input0: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    input1: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    output: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    val: int
):
    add_core(input0, input1, output, val, True)

Example code:

def add_core(input0: pypto.Tensor, input1: pypto.Tensor, output: pypto.Tensor, val: int, add1_flag: bool = False):
    # Tiling configuration and loop logic.
    pypto.set_vec_tile_shapes(1, 4, 1, 64)

    # Calculate the loop parameters.
    b = input0.shape[0]
    tile_b = 1
    b_loop = b // tile_b

    for idx in pypto.loop(b_loop):
        b_offset = idx * tile_b
        b_offset_end = (idx + 1) * tile_b
        t0_sub = input0[b_offset:b_offset_end, ...]
        t1_sub = input1[b_offset:b_offset_end, ...]
        t3_sub = t0_sub + t1_sub
        if add1_flag:
            output[b_offset:b_offset_end, ...] = t3_sub + val
        else:
            output[b_offset:b_offset_end, ...] = t3_sub

This use case adds an optional parameter add1_flag to the add_kernel function and uses this parameter for different processing. If add1_flag is True, the parameter val is added to the output result; otherwise, the result of the previous processing is directly output.

For the complete example, see condition.py.

Dynamic Conditional Branch

Evaluate conditions and branches at runtime to execute the corresponding functionality. Core APIs include:

• pypto.cond(condition): Evaluates a condition at runtime.
• pypto.is_loop_begin(idx): Determines whether the current iteration is the first iteration of a loop.
• pypto.is_loop_end(idx): Determines whether the current iteration is the last iteration of a loop.

@pypto.frontend.jit
def add_kernel(
    input0: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    input1: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    output: pypto.Tensor([pypto.DYNAMIC, 4, 1, 64], pypto.DT_FP32),
    val: int
):
    ...
    for idx in pypto.loop(b_loop):
        t3_sub = t0_sub + t1_sub
        if idx < 2:  # Dynamic conditional check.
            output[b_offset:b_offset_end, ...] = t3_sub + val
        else:
            output[b_offset:b_offset_end, ...] = t3_sub

        # Or a condition based on the loop position.
        if pypto.is_loop_begin(idx):
            output[b_offset:b_offset_end, ...] = t3_sub + val
        elif pypto.is_loop_end(idx):
            output[b_offset:b_offset_end, ...] = t3_sub + val + 1
        else:
            output[b_offset:b_offset_end, ...] = t3_sub

For the complete example, see:

condition.py