Why using tl.multiple_of(a_ptr, [16, 16]) makes it very fast?

[foreverpiano]

code copy from grouped GEMM

def grouped_matmul_kernel(
    group_a_ptrs,
    group_b_ptrs,
    group_c_ptrs,
    g_sizes,
    g_lds,
    group_size,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    NUM_SM: tl.constexpr,
):
    tile_id = tl.program_id(0)
    last = 0
    for g in range(group_size):
        gm = tl.load(g_sizes + g*3)
        gn = tl.load(g_sizes + g*3 + 1)
        gk = tl.load(g_sizes + g*3 + 2)
        
        num_m_tiles = tl.cdiv(gm, BLOCK_SIZE_M)
        num_n_tiles = tl.cdiv(gn, BLOCK_SIZE_N)
        num_tiles = num_m_tiles * num_n_tiles
        
        while tile_id >= last and tile_id < last + num_tiles:
            kk=gk
            stride_ma = tl.load(g_lds + g*3)
            stride_mb = tl.load(g_lds + g*3 + 1)
            stride_mc = tl.load(g_lds + g*3 + 2)
            # solve tile_idx
            relative_id = tile_id - last
            m_tile_idx = relative_id // num_n_tiles
            n_tile_idx = relative_id % num_n_tiles
            a_ptrs = tl.load(group_a_ptrs + g).to(tl.pointer_type(tl.float16))
            b_ptrs = tl.load(group_b_ptrs + g).to(tl.pointer_type(tl.float16))
            c_ptrs = tl.load(group_c_ptrs + g).to(tl.pointer_type(tl.float16))
            
            offset_m = tl.arange(0, BLOCK_SIZE_M) + m_tile_idx * BLOCK_SIZE_M
            offset_n = tl.arange(0, BLOCK_SIZE_N) + n_tile_idx * BLOCK_SIZE_N
            offset_k = tl.arange(0, BLOCK_SIZE_K)
            
            a_ptr = a_ptrs + offset_m[:, None] * stride_ma + offset_k[None, :] # [BLOCK_SIZE_M, BLOCK_SIZE_K]
            b_ptr = b_ptrs + offset_k[:, None] * stride_mb + offset_n[None, :] # [BLOCK_SIZE_K, BLOCK_SIZE_N]
            
            acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
            
            for k in range(0, tl.cdiv(kk, BLOCK_SIZE_K)):
            # without mask
            #for k in range(0, gk, BLOCK_SIZE_K):
                # hint to Triton compiler to do proper loop pipelining
                tl.multiple_of(a_ptr, [16, 16]) 
                tl.multiple_of(b_ptr, [16, 16])
                a_k = tl.load(a_ptr)
                b_k = tl.load(b_ptr)
                acc += tl.dot(a_k, b_k)
                a_ptr += BLOCK_SIZE_K
                b_ptr += BLOCK_SIZE_K * stride_mb
                
                
            # 算 c 用 c 自己的 stride 会快很多
            c_ptr = c_ptrs + offset_m[:, None] * stride_mc + offset_n[None, :]
            tl.store(c_ptr, acc.to(tl.float16))
            
            tile_id += NUM_SM
            
        last += num_tiles

no tl.multiple_of(a_ptr, [16, 16]) 0.100s
with tl.multiple_of(a_ptr, [16, 16]) 0.030s
3.3x faster