def emit_multiple_assignment(self, codegen_state, insn):
ecm = codegen_state.expression_to_code_mapper
from pymbolic.primitives import Variable
from pymbolic.mapper.stringifier import PREC_NONE
func_id = insn.expression.function
parameters = insn.expression.parameters
if isinstance(func_id, Variable):
func_id = func_id.name
assignee_var_descriptors = [
codegen_state.kernel.get_var_descriptor(a)
for a in insn.assignee_var_names()]
par_dtypes = tuple(ecm.infer_type(par) for par in parameters)
mangle_result = codegen_state.kernel.mangle_function(func_id, par_dtypes)
if mangle_result is None:
raise RuntimeError("function '%s' unknown--"
"maybe you need to register a function mangler?"
% func_id)
assert mangle_result.arg_dtypes is not None
if mangle_result.target_name == "loopy_make_tuple":
# This shorcut avoids actually having to emit a 'make_tuple' function.
return self.emit_tuple_assignment(codegen_state, insn)
from loopy.expression import dtype_to_type_context
c_parameters = [
ecm(par, PREC_NONE,
dtype_to_type_context(self.target, tgt_dtype),
tgt_dtype).expr
for par, par_dtype, tgt_dtype in zip(
parameters, par_dtypes, mangle_result.arg_dtypes)]
from loopy.codegen import SeenFunction
codegen_state.seen_functions.add(
SeenFunction(func_id,
mangle_result.target_name,
mangle_result.arg_dtypes))
from pymbolic import var
for i, (a, tgt_dtype) in enumerate(
zip(insn.assignees[1:], mangle_result.result_dtypes[1:])):
if tgt_dtype != ecm.infer_type(a):
raise LoopyError("type mismatch in %d'th (1-based) left-hand "
"side of instruction '%s'" % (i+1, insn.id))
c_parameters.append(
# TODO Yuck: The "where-at function": &(...)
var("&")(
ecm(a, PREC_NONE,
dtype_to_type_context(self.target, tgt_dtype),
tgt_dtype).expr))
from pymbolic import var
result = var(mangle_result.target_name)(*c_parameters)
# In case of no assignees, we are done
if len(mangle_result.result_dtypes) == 0:
from cgen import ExpressionStatement
return ExpressionStatement(
CExpression(self.get_c_expression_to_code_mapper(), result))
result = ecm.wrap_in_typecast(
mangle_result.result_dtypes[0],
assignee_var_descriptors[0].dtype,
result)
lhs_code = ecm(insn.assignees[0], prec=PREC_NONE, type_context=None)
from cgen import Assign
return Assign(
lhs_code,
CExpression(self.get_c_expression_to_code_mapper(), result))
def seen_func(name):
idt = self.kernel.index_dtype
from loopy.codegen import SeenFunction
self.codegen_state.seen_functions.add(
SeenFunction(name, name, (idt, idt)))
def map_call(self, expr, type_context):
from pymbolic.primitives import Variable, Subscript
identifier = expr.function
# {{{ implement indexof, indexof_vec
if identifier.name in ["indexof", "indexof_vec"]:
if len(expr.parameters) != 1:
raise LoopyError("%s takes exactly one argument" % identifier.name)
arg, = expr.parameters
if not isinstance(arg, Subscript):
raise LoopyError(
"argument to %s must be a subscript" % identifier.name)
ary = self.find_array(arg)
from loopy.kernel.array import get_access_info
from pymbolic import evaluate
access_info = get_access_info(self.kernel.target, ary, arg.index,
lambda expr: evaluate(expr, self.codegen_state.var_subst_map),
self.codegen_state.vectorization_info)
from loopy.kernel.data import ImageArg
if isinstance(ary, ImageArg):
raise LoopyError("%s does not support images" % identifier.name)
if identifier.name == "indexof":
return access_info.subscripts[0]
elif identifier.name == "indexof_vec":
from loopy.kernel.array import VectorArrayDimTag
ivec = None
for iaxis, dim_tag in enumerate(ary.dim_tags):
if isinstance(dim_tag, VectorArrayDimTag):
ivec = iaxis
if ivec is None:
return access_info.subscripts[0]
else:
return (
access_info.subscripts[0]*ary.shape[ivec]
+ access_info.vector_index)
else:
raise RuntimeError("should not get here")
# }}}
if isinstance(identifier, Variable):
identifier = identifier.name
par_dtypes = tuple(self.infer_type(par) for par in expr.parameters)
processed_parameters = None
mangle_result = self.kernel.mangle_function(
identifier, par_dtypes,
ast_builder=self.codegen_state.ast_builder)
if mangle_result is None:
raise RuntimeError("function '%s' unknown--"
"maybe you need to register a function mangler?"
% identifier)
if len(mangle_result.result_dtypes) != 1:
raise LoopyError("functions with more or fewer than one return value "
"may not be used in an expression")
if mangle_result.arg_dtypes is not None:
processed_parameters = tuple(
self.rec(par,
dtype_to_type_context(self.kernel.target, tgt_dtype),
tgt_dtype)
for par, par_dtype, tgt_dtype in zip(
expr.parameters, par_dtypes, mangle_result.arg_dtypes))
else:
# /!\ FIXME For some functions (e.g. 'sin'), it makes sense to
# propagate the type context here. But for many others, it does
# not. Using the inferred type as a stopgap for now.
processed_parameters = tuple(
self.rec(par,
type_context=dtype_to_type_context(
self.kernel.target, par_dtype))
for par, par_dtype in zip(expr.parameters, par_dtypes))
from warnings import warn
warn("Calling function '%s' with unknown C signature--"
"return CallMangleInfo.arg_dtypes"
% identifier, LoopyWarning)
from loopy.codegen import SeenFunction
self.codegen_state.seen_functions.add(
SeenFunction(identifier,
mangle_result.target_name,
mangle_result.arg_dtypes or par_dtypes))
return var(mangle_result.target_name)(*processed_parameters)
def seen_func(name):
from loopy.codegen import SeenFunction
self.codegen_state.seen_functions.add(
SeenFunction(name, f"{name}_{suffix}",
(result_dtype, result_dtype), (result_dtype, )))