def __init__(self, dtype):
self.dtype = dtype
def float_mapper(x):
if isinstance(x, float):
return "value_type(%s)" % repr(x)
else:
return repr(x)
CCodeMapper.__init__(self, float_mapper, reverse=False)
def __init__(self, dtype):
self.dtype = dtype
def float_mapper(x):
if isinstance(x, float):
return "value_type(%s)" % repr(x)
else:
return repr(x)
CCodeMapper.__init__(self, float_mapper, reverse=False)
def get_segmented_function_preamble(kernel, func_id):
op = func_id.reduction_op
prefix = op.prefix(func_id.scalar_dtype, func_id.segment_flag_dtype)
from pymbolic.mapper.c_code import CCodeMapper
c_code_mapper = CCodeMapper()
return (prefix, """
inline %(scalar_t)s %(prefix)s_init(%(segment_flag_t)s *segment_flag_out)
{
*segment_flag_out = 0;
return %(neutral)s;
}
inline %(scalar_t)s %(prefix)s_update(
%(scalar_t)s op1, %(segment_flag_t)s segment_flag1,
%(scalar_t)s op2, %(segment_flag_t)s segment_flag2,
%(segment_flag_t)s *segment_flag_out)
{
*segment_flag_out = segment_flag1 | segment_flag2;
return segment_flag2 ? op2 : %(combined)s;
}
""" % dict(
scalar_t=kernel.target.dtype_to_typename(func_id.scalar_dtype),
prefix=prefix,
segment_flag_t=kernel.target.dtype_to_typename(
func_id.segment_flag_dtype),
neutral=c_code_mapper(
op.inner_reduction.neutral_element(func_id.scalar_dtype)),
combined=op.op % ("op1", "op2"),
))
def map_power(self, expr, enclosing_prec):
from pymbolic.mapper.stringifier import PREC_NONE
tgt_dtype = self.infer_type(expr)
if 'c' == tgt_dtype.kind:
if expr.exponent in [2, 3, 4]:
value = expr.base
for i in range(expr.exponent - 1):
value = value * expr.base
return self.rec(value, enclosing_prec)
else:
b_complex = 'c' == self.infer_type(expr.base).kind
e_complex = 'c' == self.infer_type(expr.exponent).kind
if b_complex and not e_complex:
return "%s_powr(%s, %s)" % (
self.complex_type_name(tgt_dtype),
self.rec(expr.base, PREC_NONE),
self.rec(expr.exponent, PREC_NONE))
else:
return "%s_pow(%s, %s)" % (
self.complex_type_name(tgt_dtype),
self.rec(expr.base, PREC_NONE),
self.rec(expr.exponent, PREC_NONE))
return CCodeMapperBase.map_power(self, expr, enclosing_prec)
def map_sum(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
is_complex = tgt_dtype.kind == "c"
if not is_complex:
return CCodeMapperBase.map_sum(self, expr, enclosing_prec)
else:
tgt_name = complex_type_name(tgt_dtype)
reals = [child for child in expr.children if "c" != self.infer_type(child).kind]
complexes = [child for child in expr.children if "c" == self.infer_type(child).kind]
from pymbolic.mapper.stringifier import PREC_SUM, PREC_NONE
real_sum = self.join_rec(" + ", reals, PREC_SUM)
if len(complexes) == 1:
myprec = PREC_SUM
else:
myprec = PREC_NONE
complex_sum = self.rec(complexes[0], myprec)
for child in complexes[1:]:
complex_sum = "%s_add(%s, %s)" % (tgt_name, complex_sum, self.rec(child, PREC_NONE))
if real_sum:
result = "%s_add(%s_fromreal(%s), %s)" % (tgt_name, tgt_name, real_sum, complex_sum)
else:
result = complex_sum
return self.parenthesize_if_needed(result, enclosing_prec, PREC_SUM)
def map_sum(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
is_complex = tgt_dtype.kind == 'c'
if not is_complex:
return CCodeMapperBase.map_sum(self, expr, enclosing_prec)
else:
tgt_name = self.complex_type_name(tgt_dtype)
reals = [
child for child in expr.children
if 'c' != self.infer_type(child).kind
]
complexes = [
child for child in expr.children
if 'c' == self.infer_type(child).kind
]
from pymbolic.mapper.stringifier import PREC_SUM
real_sum = self.join_rec(" + ", reals, PREC_SUM)
complex_sum = self.join_rec(" + ", complexes, PREC_SUM)
if real_sum:
result = "%s_fromreal(%s) + %s" % (tgt_name, real_sum,
complex_sum)
else:
result = complex_sum
return self.parenthesize_if_needed(result, enclosing_prec,
PREC_SUM)
def map_quotient(self, expr, enclosing_prec):
from pymbolic.mapper.stringifier import PREC_NONE
n_complex = "c" == self.infer_type(expr.numerator).kind
d_complex = "c" == self.infer_type(expr.denominator).kind
tgt_dtype = self.infer_type(expr)
if not (n_complex or d_complex):
return CCodeMapperBase.map_quotient(self, expr, enclosing_prec)
elif n_complex and not d_complex:
return "%s_divider(%s, %s)" % (
complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator, PREC_NONE),
)
elif not n_complex and d_complex:
return "%s_rdivide(%s, %s)" % (
complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator, PREC_NONE),
)
else:
return "%s_divide(%s, %s)" % (
complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator, PREC_NONE),
)
def map_product(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
is_complex = 'c' == tgt_dtype.kind
if not is_complex:
return CCodeMapperBase.map_product(self, expr, enclosing_prec)
else:
tgt_name = self.complex_type_name(tgt_dtype)
reals = [child for child in expr.children
if 'c' != self.infer_type(child).kind]
complexes = [child for child in expr.children
if 'c' == self.infer_type(child).kind]
from pymbolic.mapper.stringifier import PREC_PRODUCT, PREC_NONE
real_prd = self.join_rec("*", reals, PREC_PRODUCT)
if len(complexes) == 1:
myprec = PREC_PRODUCT
else:
myprec = PREC_NONE
complex_prd = self.rec(complexes[0], myprec)
for child in complexes[1:]:
complex_prd = "%s_mul(%s, %s)" % (
tgt_name, complex_prd,
self.rec(child, PREC_NONE))
if real_prd:
# elementwise semantics are correct
result = "%s * %s" % (real_prd, complex_prd)
else:
result = complex_prd
return self.parenthesize_if_needed(result, enclosing_prec, PREC_PRODUCT)
def map_power(self, expr, enclosing_prec):
from pymbolic.mapper.stringifier import PREC_NONE
tgt_dtype = self.infer_type(expr)
if 'c' == tgt_dtype.kind:
if expr.exponent in [2, 3, 4]:
value = expr.base
for i in range(expr.exponent-1):
value = value * expr.base
return self.rec(value, enclosing_prec)
else:
b_complex = 'c' == self.infer_type(expr.base).kind
e_complex = 'c' == self.infer_type(expr.exponent).kind
if b_complex and not e_complex:
return "%s_powr(%s, %s)" % (
self.complex_type_name(tgt_dtype),
self.rec(expr.base, PREC_NONE),
self.rec(expr.exponent, PREC_NONE))
else:
return "%s_pow(%s, %s)" % (
self.complex_type_name(tgt_dtype),
self.rec(expr.base, PREC_NONE),
self.rec(expr.exponent, PREC_NONE))
return CCodeMapperBase.map_power(self, expr, enclosing_prec)
def get_kernel(self, vector_dtypes, scalar_dtypes):
from pymbolic.mapper.stringifier import PREC_NONE
from pymbolic.mapper.c_code import CCodeMapper
elwise = self.elementwise_mod
result_dtype = self.result_dtype_getter(
dict(zip(self.vector_deps, vector_dtypes)),
dict(zip(self.scalar_deps, scalar_dtypes)),
self.constant_dtypes)
args = [elwise.VectorArg(result_dtype, vei.name)
for vei in self.vec_expr_info_list
if not vei.do_not_return]
def real_const_mapper(num):
# Make sure we do not generate integers or doubles by accident.
# Oh, C and your broken division semantics.
r = repr(num)
if "." not in r or result_dtype == numpy.float32:
from pytools import to_uncomplex_dtype
from cgen import dtype_to_ctype
return "%s(%s)" % (dtype_to_ctype(
to_uncomplex_dtype(result_dtype)), r)
else:
return r
code_mapper = CCodeMapper(constant_mapper=real_const_mapper)
code_lines = []
for vei in self.vec_expr_info_list:
expr_code = code_mapper(vei.expr, PREC_NONE)
if vei.do_not_return:
from cgen import dtype_to_ctype
code_lines.append(
"%s %s = %s;" % (
dtype_to_ctype(result_dtype), vei.name, expr_code))
else:
code_lines.append(
"%s[i] = %s;" % (vei.name, expr_code))
# common subexpressions have been taken care of by the compiler
assert not code_mapper.cse_names
args.extend(
elwise.VectorArg(dtype, name)
for dtype, name in zip(vector_dtypes, self.vector_dep_names))
args.extend(
elwise.ScalarArg(dtype, name)
for dtype, name in zip(scalar_dtypes, self.scalar_dep_names))
return KernelRecord(
kernel=self.make_kernel_internal(args, "\n".join(code_lines)),
result_dtype=result_dtype)
def map_quotient(self, expr, enclosing_prec):
from pymbolic.mapper.stringifier import PREC_NONE
n_complex = 'c' == self.infer_type(expr.numerator).kind
d_complex = 'c' == self.infer_type(expr.denominator).kind
tgt_dtype = self.infer_type(expr)
if not (n_complex or d_complex):
return CCodeMapperBase.map_quotient(self, expr, enclosing_prec)
elif n_complex and not d_complex:
# elementwise semnatics are correct
return CCodeMapperBase.map_quotient(self, expr, enclosing_prec)
elif not n_complex and d_complex:
return "%s_rdivide(%s, %s)" % (self.complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator,
PREC_NONE))
else:
return "%s_divide(%s, %s)" % (self.complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator,
PREC_NONE))
def get_argext_preamble(kernel, func_id):
op = func_id.reduction_op
prefix = op.prefix(func_id.scalar_dtype)
from pymbolic.mapper.c_code import CCodeMapper
c_code_mapper = CCodeMapper()
neutral = get_ge_neutral if op.neutral_sign < 0 else get_le_neutral
return (prefix, """
inline %(scalar_t)s %(prefix)s_init(%(index_t)s *index_out)
{
*index_out = INT_MIN;
return %(neutral)s;
}
inline %(scalar_t)s %(prefix)s_update(
%(scalar_t)s op1, %(index_t)s index1,
%(scalar_t)s op2, %(index_t)s index2,
%(index_t)s *index_out)
{
if (op2 %(comp)s op1)
{
*index_out = index2;
return op2;
}
else
{
*index_out = index1;
return op1;
}
}
""" % dict(
scalar_t=kernel.target.dtype_to_typename(func_id.scalar_dtype),
prefix=prefix,
index_t=kernel.target.dtype_to_typename(kernel.index_dtype),
neutral=c_code_mapper(neutral(func_id.scalar_dtype)),
comp=op.update_comparison,
))
def map_product(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
is_complex = 'c' == tgt_dtype.kind
if not is_complex:
return CCodeMapperBase.map_product(self, expr, enclosing_prec)
else:
tgt_name = self.complex_type_name(tgt_dtype)
reals = [
child for child in expr.children
if 'c' != self.infer_type(child).kind
]
complexes = [
child for child in expr.children
if 'c' == self.infer_type(child).kind
]
from pymbolic.mapper.stringifier import PREC_PRODUCT, PREC_NONE
real_prd = self.join_rec("*", reals, PREC_PRODUCT)
if len(complexes) == 1:
myprec = PREC_PRODUCT
else:
myprec = PREC_NONE
complex_prd = self.rec(complexes[0], myprec)
for child in complexes[1:]:
complex_prd = "%s_mul(%s, %s)" % (tgt_name, complex_prd,
self.rec(child, PREC_NONE))
if real_prd:
# elementwise semantics are correct
result = "%s * %s" % (real_prd, complex_prd)
else:
result = complex_prd
return self.parenthesize_if_needed(result, enclosing_prec,
PREC_PRODUCT)
def map_sum(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
is_complex = tgt_dtype.kind == "c"
if not is_complex:
return CCodeMapperBase.map_sum(self, expr, enclosing_prec)
else:
tgt_name = complex_type_name(tgt_dtype)
reals = [
child for child in expr.children
if "c" != self.infer_type(child).kind
]
complexes = [
child for child in expr.children
if "c" == self.infer_type(child).kind
]
from pymbolic.mapper.stringifier import PREC_SUM, PREC_NONE
real_sum = self.join_rec(" + ", reals, PREC_SUM)
if len(complexes) == 1:
myprec = PREC_SUM
else:
myprec = PREC_NONE
complex_sum = self.rec(complexes[0], myprec)
for child in complexes[1:]:
complex_sum = "{}_add({}, {})".format(
tgt_name, complex_sum, self.rec(child, PREC_NONE))
if real_sum:
result = "{}_add({}_fromreal({}), {})".format(
tgt_name, tgt_name, real_sum, complex_sum)
else:
result = complex_sum
return self.parenthesize_if_needed(result, enclosing_prec,
PREC_SUM)
def map_product(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
is_complex = "c" == tgt_dtype.kind
if not is_complex:
return CCodeMapperBase.map_product(self, expr, enclosing_prec)
else:
tgt_name = complex_type_name(tgt_dtype)
reals = [
child for child in expr.children
if "c" != self.infer_type(child).kind
]
complexes = [
child for child in expr.children
if "c" == self.infer_type(child).kind
]
from pymbolic.mapper.stringifier import PREC_PRODUCT, PREC_NONE
real_prd = self.join_rec("*", reals, PREC_PRODUCT)
if len(complexes) == 1:
myprec = PREC_PRODUCT
else:
myprec = PREC_NONE
complex_prd = self.rec(complexes[0], myprec)
for child in complexes[1:]:
complex_prd = "{}_mul({}, {})".format(
tgt_name, complex_prd, self.rec(child, PREC_NONE))
if real_prd:
result = f"{tgt_name}_rmul({real_prd}, {complex_prd})"
else:
result = complex_prd
return self.parenthesize_if_needed(result, enclosing_prec,
PREC_PRODUCT)
def test_fft():
numpy = pytest.importorskip("numpy")
from pymbolic import var
from pymbolic.algorithm import fft, sym_fft
vars = numpy.array([var(chr(97 + i)) for i in range(16)], dtype=object)
print(vars)
print(fft(vars))
traced_fft = sym_fft(vars)
from pymbolic.mapper.stringifier import PREC_NONE
from pymbolic.mapper.c_code import CCodeMapper
ccm = CCodeMapper()
code = [ccm(tfi, PREC_NONE) for tfi in traced_fft]
for cse_name, cse_str in enumerate(ccm.cse_name_list):
print("%s = %s" % (cse_name, cse_str))
for i, line in enumerate(code):
print("result[%d] = %s" % (i, line))
def map_sum(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
is_complex = tgt_dtype.kind == 'c'
if not is_complex:
return CCodeMapperBase.map_sum(self, expr, enclosing_prec)
else:
tgt_name = self.complex_type_name(tgt_dtype)
reals = [child for child in expr.children
if 'c' != self.infer_type(child).kind]
complexes = [child for child in expr.children
if 'c' == self.infer_type(child).kind]
from pymbolic.mapper.stringifier import PREC_SUM
real_sum = self.join_rec(" + ", reals, PREC_SUM)
complex_sum = self.join_rec(" + ", complexes, PREC_SUM)
if real_sum:
result = "%s_fromreal(%s) + %s" % (tgt_name, real_sum, complex_sum)
else:
result = complex_sum
return self.parenthesize_if_needed(result, enclosing_prec, PREC_SUM)
def map_quotient(self, expr, enclosing_prec):
from pymbolic.mapper.stringifier import PREC_NONE
n_complex = "c" == self.infer_type(expr.numerator).kind
d_complex = "c" == self.infer_type(expr.denominator).kind
tgt_dtype = self.infer_type(expr)
if not (n_complex or d_complex):
return CCodeMapperBase.map_quotient(self, expr, enclosing_prec)
elif n_complex and not d_complex:
return "{}_divider({}, {})".format(
complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator, PREC_NONE))
elif not n_complex and d_complex:
return "{}_rdivide({}, {})".format(
complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator, PREC_NONE))
else:
return "{}_divide({}, {})".format(
complex_type_name(tgt_dtype),
self.rec(expr.numerator, PREC_NONE),
self.rec(expr.denominator, PREC_NONE))
def __init__(self):
CCodeMapper.__init__(self, repr, reverse=False)
def map_remainder(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
if 'c' == tgt_dtype.kind:
raise RuntimeError("complex remainder not defined")
return CCodeMapperBase.map_remainder(self, expr, enclosing_prec)
def __init__(self, infer_type):
CCodeMapperBase.__init__(self)
self.infer_type = infer_type
def map_remainder(self, expr, enclosing_prec):
tgt_dtype = self.infer_type(expr)
if 'c' == tgt_dtype.kind:
raise RuntimeError("complex remainder not defined")
return CCodeMapperBase.map_remainder(self, expr, enclosing_prec)
def __init__(self):
CCodeMapper.__init__(self, repr, reverse=False)
def __init__(self, infer_type):
CCodeMapperBase.__init__(self)
self.infer_type = infer_type
