@triton.jit
def store_2d(vals, ptr, sz0: tl.constexpr, sz1: tl.constexpr, n0, n1, max0,
    max1, stride0=None, stride1=1):
    """Store 2d block into (n0,n1)th chunk of matrix (defined by ptr), where each chunk has size (sz0, sz1)"""
    stride0 = stride0 or sz1
    offs0 = get_1d_offest(sz0, n0)
    offs1 = get_1d_offest(sz1, n1)
    offs = get_2d_offset(offs0, offs1, stride0, stride1)
    mask = get_2d_mask(offs0, offs1, max0, max1)
    tl.store(ptr + offs, vals, mask)
