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