def get_kernel(self, kernel_name, **kwargs):
data = kernel_name, tuple(kwargs.items())
if data in self.cache:
return profile_kernel(self.cache[data], kernel_name)
else:
args, src = self._get_code(kernel_name, **kwargs)
knl = ElementwiseKernel(self.ctx,
args,
src,
kernel_name,
preamble=self.preamble)
self.cache[data] = knl
return profile_kernel(knl, kernel_name)
def _setup_calls(self):
calls = []
prg = self.program
array_index = self._array_index
for item in self.data:
type = item.get('type')
if type == 'kernel':
kernel = item.get('kernel')
method = getattr(prg, kernel)
method = profile_kernel(method, method.function_name)
dest = item['dest']
src = item.get('source', dest)
args = [self._queue, None, None]
for arg in item['args']:
args.append(self._get_argument(arg, dest, src))
loop = item['loop']
args.append(self._get_argument('kern', dest, src))
info = dict(method=method,
dest=self._array_map[dest],
src=self._array_map[src],
args=args,
loop=loop,
src_idx=array_index[src],
dst_idx=array_index[dest],
type='kernel')
elif type == 'method':
info = dict(item)
if info.get('method') == 'do_reduce':
args = info.get('args')
grp = args[0]
args[0] = [
x for x in grp.equations if hasattr(x, 'reduce')
]
args[1] = self._array_map[args[1]]
elif 'iteration' in type:
group = item['group']
equations = get_equations_with_converged(group._orig_group)
info = dict(type=type, equations=equations, group=group)
else:
raise RuntimeError('Unknown type %s' % type)
calls.append(info)
return calls
def _setup_call_data(self):
array_map = self.acceleration_eval_helper._array_map
q = self.acceleration_eval_helper._queue
calls = self.calls
for method, info in self.data.items():
for dest_name, (kernel, args) in info.items():
dest = array_map[dest_name]
# Note: This is done to do some late binding. Instead of just
# directly storing the dest.gpu.x, we compute it on the fly
# as the number of particles and the actual buffer may change.
def _getter(dest_gpu, x):
return getattr(dest_gpu, x).data
_args = [
functools.partial(_getter, dest.gpu, x[2:]) for x in args
]
all_args = [q, None, None] + _args
call = getattr(self.program, kernel)
call = profile_kernel(call, call.function_name)
calls[method][dest] = (call, all_args, dest)