def map_power(self, expr, type_context):
tgt_dtype = self.infer_type(expr)
exponent_dtype = self.infer_type(expr.exponent)
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return 1
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, type_context)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base * expr.base, type_context)
if exponent_dtype.is_integral():
from loopy.codegen import SeenFunction
func_name = (
"loopy_pow_"
f"{tgt_dtype.numpy_dtype}_{exponent_dtype.numpy_dtype}")
self.codegen_state.seen_functions.add(
SeenFunction("int_pow", func_name, (tgt_dtype, exponent_dtype),
(tgt_dtype, )))
return var(func_name)(self.rec(expr.base, type_context),
self.rec(expr.exponent, type_context))
else:
return self.rec(var("pow")(expr.base, expr.exponent), type_context)
def map_power(self, expr, type_context):
tgt_dtype = self.infer_type(expr)
base_dtype = self.infer_type(expr.base)
exponent_dtype = self.infer_type(expr.exponent)
from pymbolic.primitives import is_constant, is_zero
if is_constant(expr.exponent):
if is_zero(expr.exponent):
return 1
elif is_zero(expr.exponent - 1):
return self.rec(expr.base, type_context)
elif is_zero(expr.exponent - 2):
return self.rec(expr.base * expr.base, type_context)
if exponent_dtype.is_integral():
from loopy.codegen import SeenFunction
func_name = (
"loopy_pow_"
f"{tgt_dtype.numpy_dtype}_{exponent_dtype.numpy_dtype}")
self.codegen_state.seen_functions.add(
SeenFunction("int_pow", func_name, (tgt_dtype, exponent_dtype),
(tgt_dtype, )))
# FIXME: This need some more callables to be registered.
return var(func_name)(self.rec(expr.base, type_context),
self.rec(expr.exponent, type_context))
else:
from loopy.codegen import SeenFunction
clbl = self.codegen_state.ast_builder.known_callables["pow"]
clbl = clbl.with_types({
0: tgt_dtype,
1: exponent_dtype
}, self.codegen_state.callables_table)[0]
self.codegen_state.seen_functions.add(
SeenFunction(clbl.name, clbl.name_in_target,
(base_dtype, exponent_dtype), (tgt_dtype, )))
return var(clbl.name_in_target)(self.rec(expr.base, type_context),
self.rec(expr.exponent,
type_context))
def map_call(self, expr, enclosing_prec):
from pymbolic.primitives import Variable
from pymbolic.mapper.stringifier import PREC_NONE
identifier = expr.function
if identifier.name in ["indexof", "indexof_vec"]:
raise LoopyError(
"indexof, indexof_vec not yet supported in Python")
if isinstance(identifier, Variable):
identifier = identifier.name
par_dtypes = tuple(
self.type_inf_mapper(par) for par in expr.parameters)
str_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")
str_parameters = [
self.rec(par, PREC_NONE) for par, par_dtype, tgt_dtype in zip(
expr.parameters, par_dtypes, mangle_result.arg_dtypes)
]
from loopy.codegen import SeenFunction
self.codegen_state.seen_functions.add(
SeenFunction(identifier, mangle_result.target_name,
mangle_result.arg_dtypes or par_dtypes,
mangle_result.result_dtypes))
return "{}({})".format(mangle_result.target_name,
", ".join(str_parameters))
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
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 map_call(self, expr, enclosing_prec, type_context):
from pymbolic.primitives import Variable, Subscript
from pymbolic.mapper.stringifier import PREC_NONE
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)
str_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:
str_parameters = [
self.rec(par, PREC_NONE,
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.
str_parameters = [
self.rec(par,
PREC_NONE,
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 "%s(%s)" % (mangle_result.target_name,
", ".join(str_parameters))
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 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, )))
def seen_func(name):
from loopy.codegen import SeenFunction
self.codegen_state.seen_functions.add(
SeenFunction(name, "%s_%s" % (name, suffix),
(result_dtype, result_dtype)))