def generate_body(kernel):
if kernel.schedule is None:
from loopy.schedule import get_one_scheduled_kernel
kernel = get_one_scheduled_kernel(kernel)
from loopy.kernel import kernel_state
if kernel.state != kernel_state.SCHEDULED:
raise LoopyError("cannot generate code for a kernel that has not been " "scheduled")
from loopy.preprocess import infer_unknown_types
kernel = infer_unknown_types(kernel, expect_completion=True)
from loopy.check import pre_codegen_checks
pre_codegen_checks(kernel)
logger.info("%s: generate code: start" % kernel.name)
allow_complex = False
for var in kernel.args + list(six.itervalues(kernel.temporary_variables)):
if var.dtype.kind == "c":
allow_complex = True
seen_dtypes = set()
seen_functions = set()
initial_implemented_domain = isl.BasicSet.from_params(kernel.assumptions)
codegen_state = CodeGenerationState(
kernel=kernel,
implemented_domain=initial_implemented_domain,
implemented_predicates=frozenset(),
seen_dtypes=seen_dtypes,
seen_functions=seen_functions,
var_subst_map={},
allow_complex=allow_complex,
)
code_str, implemented_domains = kernel.target.generate_body(kernel, codegen_state)
from loopy.check import check_implemented_domains
assert check_implemented_domains(kernel, implemented_domains, code_str)
logger.info("%s: generate code: done" % kernel.name)
return code_str
def generate_code_v2(kernel):
"""
:returns: a :class:`CodeGenerationResult`
"""
from loopy.kernel import kernel_state
if kernel.state == kernel_state.INITIAL:
from loopy.preprocess import preprocess_kernel
kernel = preprocess_kernel(kernel)
if kernel.schedule is None:
from loopy.schedule import get_one_scheduled_kernel
kernel = get_one_scheduled_kernel(kernel)
if kernel.state != kernel_state.SCHEDULED:
raise LoopyError("cannot generate code for a kernel that has not been "
"scheduled")
# {{{ cache retrieval
from loopy import CACHING_ENABLED
if CACHING_ENABLED:
input_kernel = kernel
try:
result = code_gen_cache[input_kernel]
logger.debug("%s: code generation cache hit" % kernel.name)
return result
except KeyError:
pass
# }}}
from loopy.type_inference import infer_unknown_types
kernel = infer_unknown_types(kernel, expect_completion=True)
from loopy.check import pre_codegen_checks
pre_codegen_checks(kernel)
logger.info("%s: generate code: start" % kernel.name)
# {{{ examine arg list
from loopy.kernel.data import ValueArg
from loopy.kernel.array import ArrayBase
implemented_data_info = []
for arg in kernel.args:
is_written = arg.name in kernel.get_written_variables()
if isinstance(arg, ArrayBase):
implemented_data_info.extend(
arg.decl_info(
kernel.target,
is_written=is_written,
index_dtype=kernel.index_dtype))
elif isinstance(arg, ValueArg):
implemented_data_info.append(ImplementedDataInfo(
target=kernel.target,
name=arg.name,
dtype=arg.dtype,
arg_class=ValueArg,
is_written=is_written))
else:
raise ValueError("argument type not understood: '%s'" % type(arg))
allow_complex = False
for var in kernel.args + list(six.itervalues(kernel.temporary_variables)):
if var.dtype.involves_complex():
allow_complex = True
# }}}
seen_dtypes = set()
seen_functions = set()
seen_atomic_dtypes = set()
initial_implemented_domain = isl.BasicSet.from_params(kernel.assumptions)
codegen_state = CodeGenerationState(
kernel=kernel,
implemented_data_info=implemented_data_info,
implemented_domain=initial_implemented_domain,
implemented_predicates=frozenset(),
seen_dtypes=seen_dtypes,
seen_functions=seen_functions,
seen_atomic_dtypes=seen_atomic_dtypes,
var_subst_map={},
allow_complex=allow_complex,
var_name_generator=kernel.get_var_name_generator(),
is_generating_device_code=False,
gen_program_name=(
kernel.target.host_program_name_prefix
+ kernel.name
+ kernel.target.host_program_name_suffix),
schedule_index_end=len(kernel.schedule))
from loopy.codegen.result import generate_host_or_device_program
codegen_result = generate_host_or_device_program(
codegen_state,
schedule_index=0)
device_code_str = codegen_result.device_code()
from loopy.check import check_implemented_domains
assert check_implemented_domains(kernel, codegen_result.implemented_domains,
device_code_str)
# {{{ handle preambles
for arg in kernel.args:
seen_dtypes.add(arg.dtype)
for tv in six.itervalues(kernel.temporary_variables):
seen_dtypes.add(tv.dtype)
preambles = kernel.preambles[:]
preamble_info = PreambleInfo(
kernel=kernel,
seen_dtypes=seen_dtypes,
seen_functions=seen_functions,
# a set of LoopyTypes (!)
seen_atomic_dtypes=seen_atomic_dtypes)
preamble_generators = (kernel.preamble_generators
+ kernel.target.get_device_ast_builder().preamble_generators())
for prea_gen in preamble_generators:
preambles.extend(prea_gen(preamble_info))
codegen_result = codegen_result.copy(device_preambles=preambles)
# }}}
logger.info("%s: generate code: done" % kernel.name)
if CACHING_ENABLED:
code_gen_cache[input_kernel] = codegen_result
return codegen_result
def generate_code_v2(kernel):
"""
:returns: a :class:`CodeGenerationResult`
"""
from loopy.kernel import KernelState
if kernel.state == KernelState.INITIAL:
from loopy.preprocess import preprocess_kernel
kernel = preprocess_kernel(kernel)
if kernel.schedule is None:
from loopy.schedule import get_one_scheduled_kernel
kernel = get_one_scheduled_kernel(kernel)
if kernel.state != KernelState.SCHEDULED:
raise LoopyError("cannot generate code for a kernel that has not been "
"scheduled")
# {{{ cache retrieval
from loopy import CACHING_ENABLED
if CACHING_ENABLED:
input_kernel = kernel
try:
result = code_gen_cache[input_kernel]
logger.debug("%s: code generation cache hit" % kernel.name)
return result
except KeyError:
pass
# }}}
from loopy.type_inference import infer_unknown_types
kernel = infer_unknown_types(kernel, expect_completion=True)
from loopy.check import pre_codegen_checks
pre_codegen_checks(kernel)
logger.info("%s: generate code: start" % kernel.name)
# {{{ examine arg list
from loopy.kernel.data import ValueArg
from loopy.kernel.array import ArrayBase
implemented_data_info = []
for arg in kernel.args:
is_written = arg.name in kernel.get_written_variables()
if isinstance(arg, ArrayBase):
implemented_data_info.extend(
arg.decl_info(
kernel.target,
is_written=is_written,
index_dtype=kernel.index_dtype))
elif isinstance(arg, ValueArg):
implemented_data_info.append(ImplementedDataInfo(
target=kernel.target,
name=arg.name,
dtype=arg.dtype,
arg_class=ValueArg,
is_written=is_written))
else:
raise ValueError("argument type not understood: '%s'" % type(arg))
allow_complex = False
for var in kernel.args + list(six.itervalues(kernel.temporary_variables)):
if var.dtype.involves_complex():
allow_complex = True
# }}}
seen_dtypes = set()
seen_functions = set()
seen_atomic_dtypes = set()
initial_implemented_domain = isl.BasicSet.from_params(kernel.assumptions)
codegen_state = CodeGenerationState(
kernel=kernel,
implemented_data_info=implemented_data_info,
implemented_domain=initial_implemented_domain,
implemented_predicates=frozenset(),
seen_dtypes=seen_dtypes,
seen_functions=seen_functions,
seen_atomic_dtypes=seen_atomic_dtypes,
var_subst_map={},
allow_complex=allow_complex,
var_name_generator=kernel.get_var_name_generator(),
is_generating_device_code=False,
gen_program_name=(
kernel.target.host_program_name_prefix
+ kernel.name
+ kernel.target.host_program_name_suffix),
schedule_index_end=len(kernel.schedule))
from loopy.codegen.result import generate_host_or_device_program
codegen_result = generate_host_or_device_program(
codegen_state,
schedule_index=0)
device_code_str = codegen_result.device_code()
from loopy.check import check_implemented_domains
assert check_implemented_domains(kernel, codegen_result.implemented_domains,
device_code_str)
# {{{ handle preambles
for arg in kernel.args:
seen_dtypes.add(arg.dtype)
for tv in six.itervalues(kernel.temporary_variables):
seen_dtypes.add(tv.dtype)
preambles = kernel.preambles[:]
preamble_info = PreambleInfo(
kernel=kernel,
seen_dtypes=seen_dtypes,
seen_functions=seen_functions,
# a set of LoopyTypes (!)
seen_atomic_dtypes=seen_atomic_dtypes,
codegen_state=codegen_state
)
preamble_generators = (kernel.preamble_generators
+ kernel.target.get_device_ast_builder().preamble_generators())
for prea_gen in preamble_generators:
preambles.extend(prea_gen(preamble_info))
codegen_result = codegen_result.copy(device_preambles=preambles)
# }}}
# For faster unpickling in the common case when implemented_domains isn't needed.
from loopy.tools import LazilyUnpicklingDict
codegen_result = codegen_result.copy(
implemented_domains=LazilyUnpicklingDict(
codegen_result.implemented_domains))
logger.info("%s: generate code: done" % kernel.name)
if CACHING_ENABLED:
code_gen_cache.store_if_not_present(input_kernel, codegen_result)
return codegen_result
def generate_code(kernel, device=None):
if device is not None:
from warnings import warn
warn("passing 'device' to generate_code() is deprecated",
DeprecationWarning, stacklevel=2)
if kernel.schedule is None:
from loopy.schedule import get_one_scheduled_kernel
kernel = get_one_scheduled_kernel(kernel)
from loopy.kernel import kernel_state
if kernel.state != kernel_state.SCHEDULED:
raise LoopyError("cannot generate code for a kernel that has not been "
"scheduled")
# {{{ cache retrieval
from loopy import CACHING_ENABLED
if CACHING_ENABLED:
input_kernel = kernel
try:
result = code_gen_cache[input_kernel]
logger.info("%s: code generation cache hit" % kernel.name)
return result
except KeyError:
pass
# }}}
from loopy.preprocess import infer_unknown_types
kernel = infer_unknown_types(kernel, expect_completion=True)
from loopy.check import pre_codegen_checks
pre_codegen_checks(kernel)
logger.info("%s: generate code: start" % kernel.name)
# {{{ examine arg list
from loopy.kernel.data import ValueArg
from loopy.kernel.array import ArrayBase
impl_arg_info = []
for arg in kernel.args:
if isinstance(arg, ArrayBase):
impl_arg_info.extend(
arg.decl_info(
kernel.target,
is_written=arg.name in kernel.get_written_variables(),
index_dtype=kernel.index_dtype))
elif isinstance(arg, ValueArg):
impl_arg_info.append(ImplementedDataInfo(
target=kernel.target,
name=arg.name,
dtype=arg.dtype,
cgen_declarator=arg.get_arg_decl(kernel.target),
arg_class=ValueArg))
else:
raise ValueError("argument type not understood: '%s'" % type(arg))
allow_complex = False
for var in kernel.args + list(six.itervalues(kernel.temporary_variables)):
if var.dtype.kind == "c":
allow_complex = True
# }}}
seen_dtypes = set()
seen_functions = set()
initial_implemented_domain = isl.BasicSet.from_params(kernel.assumptions)
codegen_state = CodeGenerationState(
kernel=kernel,
implemented_domain=initial_implemented_domain,
implemented_predicates=frozenset(),
seen_dtypes=seen_dtypes,
seen_functions=seen_functions,
var_subst_map={},
allow_complex=allow_complex)
code_str, implemented_domains = kernel.target.generate_code(
kernel, codegen_state, impl_arg_info)
from loopy.check import check_implemented_domains
assert check_implemented_domains(kernel, implemented_domains,
code_str)
# {{{ handle preambles
for arg in kernel.args:
seen_dtypes.add(arg.dtype)
for tv in six.itervalues(kernel.temporary_variables):
seen_dtypes.add(tv.dtype)
preambles = kernel.preambles[:]
preamble_generators = (kernel.preamble_generators
+ kernel.target.preamble_generators())
for prea_gen in preamble_generators:
preambles.extend(prea_gen(kernel, seen_dtypes, seen_functions))
seen_preamble_tags = set()
dedup_preambles = []
for tag, preamble in sorted(preambles, key=lambda tag_code: tag_code[0]):
if tag in seen_preamble_tags:
continue
seen_preamble_tags.add(tag)
dedup_preambles.append(preamble)
from loopy.tools import remove_common_indentation
preamble_codes = [
remove_common_indentation(lines) + "\n"
for lines in dedup_preambles]
code_str = "".join(preamble_codes) + code_str
# }}}
logger.info("%s: generate code: done" % kernel.name)
result = code_str, impl_arg_info
if CACHING_ENABLED:
code_gen_cache[input_kernel] = result
return result