def _generate_enqueue_and_set_args_module(function_name,
num_passed_args, num_cl_args,
scalar_arg_dtypes,
work_around_arg_count_bug, warn_about_arg_count_bug):
from pytools.py_codegen import PythonCodeGenerator, Indentation
arg_names = ["arg%d" % i for i in range(num_passed_args)]
if scalar_arg_dtypes is None:
body = generate_generic_arg_handling_body(num_passed_args)
else:
body = generate_specific_arg_handling_body(
function_name, num_cl_args, scalar_arg_dtypes,
warn_about_arg_count_bug=warn_about_arg_count_bug,
work_around_arg_count_bug=work_around_arg_count_bug)
err_handler = wrap_in_error_handler(body, arg_names)
gen = PythonCodeGenerator()
gen("from struct import pack")
gen("from pyopencl import status_code")
gen("")
# {{{ generate _enqueue
enqueue_name = "enqueue_knl_%s" % function_name
gen("def %s(%s):"
% (enqueue_name,
", ".join(
["self", "queue", "global_size", "local_size"]
+ arg_names
+ ["global_offset=None", "g_times_l=None",
"wait_for=None"])))
with Indentation(gen):
add_local_imports(gen)
gen.extend(err_handler)
gen("""
return _cl.enqueue_nd_range_kernel(queue, self, global_size, local_size,
global_offset, wait_for, g_times_l=g_times_l)
""")
# }}}
# {{{ generate set_args
gen("")
gen("def set_args(%s):"
% (", ".join(["self"] + arg_names)))
with Indentation(gen):
add_local_imports(gen)
gen.extend(err_handler)
# }}}
return gen.get_picklable_module(), enqueue_name
def _generate_enqueue_and_set_args_module(function_name, num_passed_args,
num_cl_args, scalar_arg_dtypes,
work_around_arg_count_bug,
warn_about_arg_count_bug):
from pytools.py_codegen import PythonCodeGenerator, Indentation
arg_names = ["arg%d" % i for i in range(num_passed_args)]
if scalar_arg_dtypes is None:
body = generate_generic_arg_handling_body(num_passed_args)
else:
body = generate_specific_arg_handling_body(
function_name,
num_cl_args,
scalar_arg_dtypes,
warn_about_arg_count_bug=warn_about_arg_count_bug,
work_around_arg_count_bug=work_around_arg_count_bug)
err_handler = wrap_in_error_handler(body, arg_names)
gen = PythonCodeGenerator()
gen("from struct import pack")
gen("from pyopencl import status_code")
gen("")
# {{{ generate _enqueue
enqueue_name = "enqueue_knl_%s" % function_name
gen("def %s(%s):" %
(enqueue_name,
", ".join(["self", "queue", "global_size", "local_size"] + arg_names +
["global_offset=None", "g_times_l=None", "wait_for=None"])))
with Indentation(gen):
add_local_imports(gen)
gen.extend(err_handler)
gen("""
return _cl.enqueue_nd_range_kernel(queue, self, global_size, local_size,
global_offset, wait_for, g_times_l=g_times_l)
""")
# }}}
# {{{ generate set_args
gen("")
gen("def set_args(%s):" % (", ".join(["self"] + arg_names)))
with Indentation(gen):
add_local_imports(gen)
gen.extend(err_handler)
# }}}
return gen.get_picklable_module(), enqueue_name
def _generate_enqueue_and_set_args_module(function_name, num_passed_args,
num_cl_args, arg_types,
work_around_arg_count_bug,
warn_about_arg_count_bug):
arg_names = ["arg%d" % i for i in range(num_passed_args)]
def gen_arg_setting(in_enqueue):
if arg_types is None:
return generate_generic_arg_handling_body(num_passed_args)
else:
return generate_specific_arg_handling_body(
function_name,
num_cl_args,
arg_types,
warn_about_arg_count_bug=warn_about_arg_count_bug,
work_around_arg_count_bug=work_around_arg_count_bug,
in_enqueue=in_enqueue)
gen = PythonCodeGenerator()
gen("from struct import pack")
gen("from pyopencl import status_code")
gen("import numpy as np")
gen("import pyopencl._cl as _cl")
gen("")
# {{{ generate _enqueue
enqueue_name = "enqueue_knl_%s" % function_name
gen("def %s(%s):" %
(enqueue_name,
", ".join(["self", "queue", "global_size", "local_size"] + arg_names +
[
"global_offset=None", "g_times_l=None",
"allow_empty_ndrange=False", "wait_for=None"
])))
with Indentation(gen):
gen.extend(gen_arg_setting(in_enqueue=True))
# Using positional args here because pybind is slow with keyword args
gen("""
return _cl.enqueue_nd_range_kernel(queue, self, global_size, local_size,
global_offset, wait_for, g_times_l,
allow_empty_ndrange)
""")
# }}}
# {{{ generate set_args
gen("")
gen("def set_args(%s):" % (", ".join(["self"] + arg_names)))
with Indentation(gen):
gen.extend(gen_arg_setting(in_enqueue=False))
# }}}
return gen.get_picklable_module(), enqueue_name
def _generate_enqueue_and_set_args_module(function_name, num_passed_args,
num_cl_args, arg_types,
include_debug_code,
work_around_arg_count_bug,
warn_about_arg_count_bug):
arg_names = ["arg%d" % i for i in range(num_passed_args)]
def gen_arg_setting(in_enqueue):
if arg_types is None:
result = generate_generic_arg_handling_body(num_passed_args)
if in_enqueue:
return result, []
else:
return result
else:
return generate_specific_arg_handling_body(
function_name,
num_cl_args,
arg_types,
warn_about_arg_count_bug=warn_about_arg_count_bug,
work_around_arg_count_bug=work_around_arg_count_bug,
in_enqueue=in_enqueue,
include_debug_code=include_debug_code)
gen = PythonCodeGenerator()
gen("from struct import pack")
gen("from pyopencl import status_code")
gen("import numpy as np")
gen("import pyopencl._cl as _cl")
gen("")
# {{{ generate _enqueue
enqueue_name = "enqueue_knl_%s" % function_name
gen("def %s(%s):" %
(enqueue_name,
", ".join(["self", "queue", "global_size", "local_size"] + arg_names +
[
"global_offset=None", "g_times_l=None",
"allow_empty_ndrange=False", "wait_for=None"
])))
with Indentation(gen):
subgen, wait_for_parts = gen_arg_setting(in_enqueue=True)
gen.extend(subgen)
if wait_for_parts:
wait_for_expr = ("[*(() if wait_for is None else wait_for), " +
", ".join("*" + wfp
for wfp in wait_for_parts) + "]")
else:
wait_for_expr = "wait_for"
# Using positional args here because pybind is slow with keyword args
gen(f"""
return _cl.enqueue_nd_range_kernel(queue, self,
global_size, local_size, global_offset,
{wait_for_expr},
g_times_l, allow_empty_ndrange)
""")
# }}}
# {{{ generate set_args
gen("")
gen("def set_args(%s):" % (", ".join(["self"] + arg_names)))
with Indentation(gen):
gen.extend(gen_arg_setting(in_enqueue=False))
# }}}
return (gen.get_picklable_module(
name=f"<pyopencl invoker for '{function_name}'>"), enqueue_name)