# Get 1d mask for 
    for _ in range(0, k, bk):
        mask_a = get_2d_mask(rm, rn, m, k) # add mask
        mask_b = get_2d_mask(rm, rn, k, n)  # add mask
            
        a = tl.load(offs_a, mask_a)
        b = tl.load(offs_b, mask_b)
        acc += tl.dot(a, b, allow_tf32=False) # matmul in block ; Weirdness: allow_tf32 must be set to False for older GPUs, otherwise won't compile
        
        # print_if(f"pid: {pid_m, pid_n} |\na: {a} | \nb: {b}", "")
        
        # increase offets, so next iteration loads next chunks
        offs_a += bk * stride_ak
        offs_b += bk * stride_bk