How many times the Triton kernel function will be called.

[XinDongol]

I am playing with https://triton-lang.org/main/getting-started/tutorials/01-vector-add.html.

I changed the block_size=4 and the len(x)=len(y)=10.

In the kernel function, I try to print out the pid

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(
    x_ptr,  # *Pointer* to first input vector.
    y_ptr,  # *Pointer* to second input vector.
    output_ptr,  # *Pointer* to output vector.
    n_elements,  # Size of the vector.
    BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
                 # NOTE: `constexpr` so it can be used as a shape value.
):

    pid = tl.program_id(axis=0)  
    tl.printf("pid:", pid).                            # print the program id
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

After runing one add op,

def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and output.is_cuda
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=4)
    return output

size = 10
x = torch.rand(size, device='cuda')
y = torch.rand(size, device='cuda')
output_torch = x + y
output_triton = add(x, y)

I got output like this

pid:2
pid:2
pid:2
pid:2
pid:1
pid:2
pid:2
pid:2
pid:2
pid:1
pid:2
pid:1
pid:2
pid:2
pid:2
pid:2
pid:2
pid:2
pid:2
pid:0
pid:2
pid:2
pid:0
pid:2
pid:2
...

After counting, I found that

Counter({'pid:2': 128, 'pid:0': 128, 'pid:1': 128})

I can understand there are 3 different program id since 10(i.e., inp len)//4(i.e., block_size) + 1 = 3
But why each pid is repeated by 128 times?