@triton.jit
def test(at, bt, ct, k):
    midx = tl.arange(0, 32)
    kidx = tl.arange(0, 32)
    nidx = tl.arange(0, 32)
    aidx = midx[:, None] * 32 + kidx[None, :]
    bidx = kidx[:, None] * 32 + nidx[None, :]
    cidx = midx[:, None] * 32 + nidx[None, :]
    a_ptrs = at + aidx
    b_ptrs = bt + bidx
    c_ptrs = ct + cidx
    for i in range(k):
        a = tl.load(a_ptrs)
        b = tl.load(b_ptrs)
        x = tl.dot(a, b)
        tl.atomic_add(c_ptrs, x)
        a_ptrs += 32
        b_ptrs += 32
        c_ptrs += 32
