def sort_by_keys(ary_list, out_list=None, key_bits=None, backend=None):
# first arg of ary_list is the key
if backend is None:
backend = ary_list[0].backend
if backend == 'opencl':
from .jit import get_ctype_from_arg
from compyle.opencl import get_queue
arg_types = [get_ctype_from_arg(arg) for arg in ary_list]
sort_knl = get_cl_sort_kernel(arg_types, ary_list)
allocator = get_allocator(get_queue())
arg_list = [ary.dev for ary in ary_list]
out_list, event = sort_knl(*arg_list,
key_bits=key_bits,
allocator=allocator)
return out_list
else:
order = argsort(ary_list[0], backend=backend)
out_list = align(ary_list[1:],
order,
out_list=out_list,
backend=backend)
return [ary_list[0]] + out_list
def sort_by_keys(ary_list,
out_list=None,
key_bits=None,
backend=None,
use_radix_sort=False):
# FIXME: Need to use returned values, cuda backend uses
# thrust that will internally allocate a new array for storing
# the sorted data so out_list will not have the sorted arrays
# first arg of ary_list is the key
if backend is None:
backend = ary_list[0].backend
if backend == 'opencl':
from .jit import get_ctype_from_arg
from compyle.opencl import get_queue
if not out_list:
out_list = [
Array(ary.dtype, allocate=False, backend=backend)
for ary in ary_list
]
arg_types = [
get_ctype_from_arg(arg, backend=backend) for arg in ary_list
]
sort_knl = get_cl_sort_kernel(arg_types, ary_list)
allocator = get_allocator(get_queue())
arg_list = [ary.dev for ary in ary_list]
out_arrays, event = sort_knl(*arg_list,
key_bits=key_bits,
allocator=allocator)
for i, out in enumerate(out_list):
out.set_data(out_arrays[i])
return out_list
elif backend == 'cython' and use_radix_sort:
out_list, order = radix_sort(ary_list,
out_list=out_list,
max_key_bits=key_bits,
backend=backend)
return out_list
elif backend == 'cython':
order = wrap(np.argsort(ary_list[0].dev), backend=backend)
out_list = align(ary_list, order, out_list=out_list, backend=backend)
return out_list
else:
order = argsort(ary_list[0], backend=backend)
modified_out_list = None
if out_list:
modified_out_list = out_list[1:]
out_list = align(ary_list[1:],
order,
out_list=modified_out_list,
backend=backend)
return [ary_list[0]] + out_list
def argsort(ary, backend=None):
# FIXME: Implement an OpenCL backend and add tests
# NOTE: argsort also sorts the array
if backend is None:
backend = ary.backend
if backend == 'cython':
result = np.argsort(ary.dev)
ary.dev = np.take(ary.dev, result)
return wrap_array(result, backend=backend)
elif backend == 'cuda':
from compyle.cuda import argsort
return argsort(ary)
else:
raise ValueError("Only cython and cuda backends supported")
