pyopencl copy_if(): возможно ли минимизировать размер возвращаемого буфера?

Question

Вот простой пример pyopencl copy_if().

Во-первых, давайте создадим большой набор (2^25) случайных целых чисел и запросим их ниже порога 500000:

import pyopencl as cl
import numpy as np
import my_pyopencl_algorithm
import time

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

from pyopencl.clrandom import rand as clrand
random_gpu = clrand(queue, (2^25,), dtype=np.int32, a=0, b=10**6)

start = time.time()
final_gpu, count_gpu, evt = my_pyopencl_algorithm.copy_if(random_gpu, "ary[i] < 500000", queue = queue)
final = final_gpu.get()
count = int(count_gpu.get())
print '\ncopy_if():\nresults=',final[:count], '\nfound=', count, '\ntime=', (time.time()-start), '\n========\n'

Вы, возможно, заметили, что я не вызываю pyopencl copy_if, а его ответвление (my_pyopencl_algorithm.copy_if). Форк pyopencl.algorithm.py можно найти здесь.

Прелесть copy_if в том, что у вас есть готовый счетчик желаемого результата и порядка от gid=0 до gid=count. Что не кажется оптимальным, так это то, что он распределяет и возвращает (из графического процессора) весь буфер, причем только первые записи имеют значение. Итак, в моей ветке pyopencl.algorithm.py я пытаюсь оптимизировать размер буфера возврата, и у меня есть это:

def sparse_copy_if(ary, predicate, extra_args=[], preamble="", queue=None, wait_for=None):
    """Copy the elements of *ary* satisfying *predicate* to an output array.

:arg predicate: a C expression evaluating to a `bool`, represented as a string.
    The value to test is available as `ary[i]`, and if the expression evaluates
    to `true`, then this value ends up in the output.
:arg extra_args: |scan_extra_args|
:arg preamble: |preamble|
:arg wait_for: |explain-waitfor|
:returns: a tuple *(out, count, event)* where *out* is the output array, *count*
    is an on-device scalar (fetch to host with `count.get()`) indicating
    how many elements satisfied *predicate*, and *event* is a
    :class:`pyopencl.Event` for dependency management. *out* is allocated
    to the same length as *ary*, but only the first *count* entries carry
    meaning.

.. versionadded:: 2013.1
"""
if len(ary) > np.iinfo(np.int32).max:
    scan_dtype = np.int64
else:
    scan_dtype = np.int32

extra_args_types, extra_args_values = extract_extra_args_types_values(extra_args)


knl = _copy_if_template.build(ary.context,
        type_aliases=(("scan_t", scan_dtype), ("item_t", ary.dtype)),
        var_values=(("predicate", predicate),),
        more_preamble=preamble, more_arguments=extra_args_types)
out = cl.array.empty_like(ary)
count = ary._new_with_changes(data=None, offset=0,
        shape=(), strides=(), dtype=scan_dtype)

# **dict is a Py2.5 workaround
evt = knl(ary, out, count, *extra_args_values,
        **dict(queue=queue, wait_for=wait_for))

'''
Now I need to copy the first num_results values from out to final_gpu (in which buffer size is minimized)
'''

prg = cl.Program(ary.context, """ 
__kernel void copy_final_results(__global int *final_gpu, __global int *out_gpu) 
{ 
__private uint gid; 
gid = get_global_id(0); 
final_gpu [gid] = out_gpu [gid]; 
} 
""").build() 

num_results= int(count.get())

final_gpu = pyopencl.array.zeros(queue, (num_results,), dtype=scan_dtype)

prg.copy_final_results (queue, (num_results,), None, final_gpu.data, out.data).wait()  

return final_gpu, evt 
#return out, count, evt 

То есть я создаю буфер final_gpu точно в соответствии с размером вывода, затем копирую в него значимые записи и возвращаю его.

Если я сейчас бегу:

start = time.time()
final_gpu, evt = my_pyopencl_algorithm.sparse_copy_if(random_gpu, "ary[i] < 500000", queue = queue)
final = final_gpu.get()
print '\ncopy_if_2():\nresults=',final, '\nfound=', count, '\ntime=', (time.time()-start) here

... кажется, это приводит к увеличению скорости на порядки. Чем более разреженные результаты, тем быстрее это происходит, так как размер буфера, подлежащего передаче (с высокой задержкой), минимизируется.

Мой вопрос: есть ли причина, по которой мы возвращаем полноразмерный буфер? Другими словами, я представляю какие-либо ошибки или я должен просто отправить патч?