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 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 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 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))
