diff --git a/loopy/codegen/expression.py b/loopy/codegen/expression.py index 5a387a68e64bdfca4fb65005605fc2606c94e779..56ed006c3eae4ee51c1c398ac4715049b2c78cc3 100644 --- a/loopy/codegen/expression.py +++ b/loopy/codegen/expression.py @@ -59,11 +59,10 @@ class TypeInferenceMapper(CombineMapper): arg_dtypes = tuple(self.rec(par) for par in expr.parameters) - for mangler in self.kernel.function_manglers: - mangle_result = mangler(identifier, arg_dtypes) - if mangle_result is not None: - result_dtype, _ = mangle_result - return result_dtype + mangle_result = self.kernel.mangle_function(identifier, arg_dtypes) + if mangle_result is not None: + result_dtype, c_name = mangle_result + return result_dtype raise RuntimeError("no type inference information on " "function '%s'" % identifier) @@ -305,10 +304,9 @@ class LoopyCCodeMapper(CCodeMapper): arg_dtypes = tuple(self.infer_type(par) for par in expr.parameters) - for mangler in self.kernel.function_manglers: - mangle_result = mangler(identifier, arg_dtypes) - if mangle_result is not None: - result_dtype, c_name = mangle_result + mangle_result = self.kernel.mangle_function(identifier, arg_dtypes) + if mangle_result is not None: + result_dtype, c_name = mangle_result self.seen_functions.add((identifier, c_name, arg_dtypes)) diff --git a/loopy/kernel.py b/loopy/kernel.py index de67d95722336a0c4e292dfe5c81778e225cd2e1..e5e15f995b6bdc8e3e9aa4f0eceba3b636ab3b47 100644 --- a/loopy/kernel.py +++ b/loopy/kernel.py @@ -495,6 +495,15 @@ def default_preamble_generator(seen_dtypes, seen_functions): # {{{ loop kernel object +def _generate_unique_possibilities(prefix): + yield prefix + + try_num = 0 + while True: + yield "%s_%d" % (prefix, try_num) + try_num += 1 + + class LoopKernel(Record): """ :ivar device: :class:`pyopencl.Device` @@ -779,14 +788,29 @@ class LoopKernel(Record): function_manglers=function_manglers, symbol_manglers=symbol_manglers) + # {{{ function mangling + + def register_function_mangler(self, mangler): + return self.copy( + function_manglers=[mangler]+self.function_manglers) + + def mangle_function(self, identifier, arg_dtypes): + for mangler in self.function_manglers: + mangle_result = mangler(identifier, arg_dtypes) + if mangle_result is not None: + return mangle_result + + return None + + # }}} + def make_unique_instruction_id(self, insns=None, based_on="insn", extra_used_ids=set()): if insns is None: insns = self.instructions used_ids = set(insn.id for insn in insns) | extra_used_ids - from loopy.tools import generate_unique_possibilities - for id_str in generate_unique_possibilities(based_on): + for id_str in _generate_unique_possibilities(based_on): if id_str not in used_ids: return id_str @@ -916,8 +940,7 @@ class LoopKernel(Record): def make_unique_var_name(self, based_on="var", extra_used_vars=set()): used_vars = self.all_variable_names() | extra_used_vars - from loopy.tools import generate_unique_possibilities - for var_name in generate_unique_possibilities(based_on): + for var_name in _generate_unique_possibilities(based_on): if var_name not in used_vars: return var_name diff --git a/loopy/tools.py b/loopy/tools.py deleted file mode 100644 index 95d4e4e413ae88171f2a1ede26351c41bdbe40d2..0000000000000000000000000000000000000000 --- a/loopy/tools.py +++ /dev/null @@ -1,13 +0,0 @@ -from __future__ import division - - - - -def generate_unique_possibilities(prefix): - yield prefix - - try_num = 0 - while True: - yield "%s_%d" % (prefix, try_num) - try_num += 1 -