def get_reduction_kernel_and_types(out_type,
block_size,
neutral,
reduce_expr,
map_expr=None,
arguments=None,
name="reduce_kernel",
keep=False,
options=[],
preamble=""):
if map_expr is None:
map_expr = "in[i]"
if arguments is None:
arguments = "const %s *in" % out_type
mod = get_reduction_module(out_type, block_size, neutral, reduce_expr,
map_expr, arguments, name, keep, options,
preamble)
from pycuda.tools import get_arg_type
func = mod.get_function(name)
arg_types = [get_arg_type(arg) for arg in arguments.split(",")]
func.prepare("P%sII" % "".join(arg_types), (block_size, 1, 1))
return func, arg_types
def get_reduction_kernel_and_types(stage, out_type, block_size,
neutral, reduce_expr, map_expr=None, arguments=None,
name="reduce_kernel", keep=False, options=None, preamble=""):
if stage == 1:
if map_expr is None:
map_expr = "in[i]"
elif stage == 2:
if map_expr is None:
map_expr = "pycuda_reduction_inp[i]"
in_arg = "const %s *pycuda_reduction_inp" % out_type
if arguments:
arguments = in_arg + ", " + arguments
else:
arguments = in_arg
else:
assert False
mod = get_reduction_module(out_type, block_size,
neutral, reduce_expr, map_expr, arguments,
name, keep, options, preamble)
from pycuda.tools import get_arg_type
func = mod.get_function(name)
arg_types = [get_arg_type(arg) for arg in arguments.split(",")]
func.prepare("P%sII" % "".join(arg_types))
return func, arg_types
def get_reduction_kernel_and_types(
out_type,
block_size,
neutral,
reduce_expr,
map_expr=None,
arguments=None,
name="reduce_kernel",
keep=False,
options=[],
preamble="",
):
if map_expr is None:
map_expr = "in[i]"
if arguments is None:
arguments = "const %s *in" % out_type
mod = get_reduction_module(
out_type, block_size, neutral, reduce_expr, map_expr, arguments, name, keep, options, preamble
)
from pycuda.tools import get_arg_type
func = mod.get_function(name)
arg_types = [get_arg_type(arg) for arg in arguments.split(",")]
func.prepare("P%sII" % "".join(arg_types), (block_size, 1, 1))
return func, arg_types
