@triton.jit
def _encode_z_kernel(xyz_ptr, distance_ptr, B, N, space_size, x_offset,
    y_offset, z_offset, str_xyz_B, str_xyz_N, str_xyz_C, BLK: tl.constexpr,
    ASSIGN_BATCH_INDEX: tl.constexpr):
    pid_b = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_n = pid_n * BLK + tl.arange(0, BLK)
    mask_n = offs_n < N
    xyz_ptrs = xyz_ptr + pid_b * str_xyz_B + offs_n * str_xyz_N
    fx = tl.load(xyz_ptrs + x_offset * str_xyz_C, mask=mask_n)
    fy = tl.load(xyz_ptrs + y_offset * str_xyz_C, mask=mask_n)
    fz = tl.load(xyz_ptrs + z_offset * str_xyz_C, mask=mask_n)
    ret = _calculate_zorder(fx, fy, fz, space_size)
    if ASSIGN_BATCH_INDEX:
        ret |= pid_b.to(tl.int64) << 48
    tl.store(distance_ptr + pid_b * N + offs_n, ret, mask=mask_n)
