diff --git a/test/test_loopy.py b/test/test_loopy.py
index 53fe1c1504db3df6f5d5a30f8f3a20a7e997f43c..e33412001573a56ecf008b26e5849ba6c457dbeb 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -2689,97 +2689,6 @@ def test_preamble_with_separate_temporaries(ctx_factory):
from loopy.kernel.data import temp_var_scope as scopes
# create a function mangler
- func_name = 'indirect'
- func_arg_dtypes = (np.int32, np.int32, np.int32)
- func_result_dtypes = (np.int32,)
-
- def __indirectmangler(kernel, name, arg_dtypes):
- """
- A function that will return a :class:`loopy.kernel.data.CallMangleInfo`
- to interface with the calling :class:`loopy.LoopKernel`
- """
- if name != func_name:
- return None
-
- from loopy.types import to_loopy_type
- from loopy.kernel.data import CallMangleInfo
-
- def __compare(d1, d2):
- # compare dtypes ignoring atomic
- return to_loopy_type(d1, for_atomic=True) == \
- to_loopy_type(d2, for_atomic=True)
-
- # check types
- if len(arg_dtypes) != len(arg_dtypes):
- raise Exception('Unexpected number of arguments provided to mangler '
- '{}, expected {}, got {}'.format(
- func_name, len(func_arg_dtypes), len(arg_dtypes)))
-
- for i, (d1, d2) in enumerate(zip(func_arg_dtypes, arg_dtypes)):
- if not __compare(d1, d2):
- raise Exception('Argument at index {} for mangler {} does not '
- 'match expected dtype. Expected {}, got {}'.
- format(i, func_name, str(d1), str(d2)))
-
- # get target for creation
- target = arg_dtypes[0].target
- return CallMangleInfo(
- target_name=func_name,
- result_dtypes=tuple(to_loopy_type(x, target=target) for x in
- func_result_dtypes),
- arg_dtypes=arg_dtypes)
-
- # create the preamble generator
- def create_preamble(arr):
- def __indirectpreamble(preamble_info):
- # find a function matching our name
- func_match = next(
- (x for x in preamble_info.seen_functions
- if x.name == func_name), None)
- desc = 'custom_funcs_indirect'
- if func_match is not None:
- from loopy.types import to_loopy_type
- # check types
- if tuple(to_loopy_type(x) for x in func_arg_dtypes) == \
- func_match.arg_dtypes:
- # if match, create our temporary
- var = lp.TemporaryVariable(
- 'lookup', initializer=arr, dtype=arr.dtype, shape=arr.shape,
- scope=scopes.GLOBAL, read_only=True)
- # and code
- code = """
- int {name}(int start, int end, int match)
- {{
- int result = start;
- for (int i = start + 1; i < end; ++i)
- {{
- if (lookup[i] == match)
- result = i;
- }}
- return result;
- }}
- """.format(name=func_name)
-
- # generate temporary variable code
- from cgen import Initializer
- from loopy.target.c import generate_array_literal
- codegen_state = preamble_info.codegen_state.copy(
- is_generating_device_code=True)
- kernel = preamble_info.kernel
- ast_builder = codegen_state.ast_builder
- target = kernel.target
- decl_info, = var.decl_info(target, index_dtype=kernel.index_dtype)
- decl = ast_builder.wrap_global_constant(
- ast_builder.get_temporary_decl(
- codegen_state, None, var,
- decl_info))
- if var.initializer is not None:
- decl = Initializer(decl, generate_array_literal(
- codegen_state, var, var.initializer))
- # return generated code
- yield (desc, '\n'.join([str(decl), code]))
- return __indirectpreamble
-
# and finally create a test
n = 10
# for each entry come up with a random number of data points
@@ -2809,9 +2718,19 @@ def test_preamble_with_separate_temporaries(ctx_factory):
lp.GlobalArg('data', shape=(data.size,), dtype=np.float64)],
)
# fixt params, and add manglers / preamble
+ from testlib import SeparateTemporariesPreambleTestHelper
+ preamble_with_sep_helper = SeparateTemporariesPreambleTestHelper(
+ func_name='indirect',
+ func_arg_dtypes=(np.int32, np.int32, np.int32),
+ func_result_dtypes=(np.int32,),
+ arr=lookup
+ )
+
kernel = lp.fix_parameters(kernel, **{'n': n})
- kernel = lp.register_preamble_generators(kernel, [create_preamble(lookup)])
- kernel = lp.register_function_manglers(kernel, [__indirectmangler])
+ kernel = lp.register_preamble_generators(
+ kernel, [preamble_with_sep_helper.preamble_gen])
+ kernel = lp.register_function_manglers(
+ kernel, [preamble_with_sep_helper.mangler])
print(lp.generate_code(kernel)[0])
# and call (functionality unimportant, more that it compiles)
diff --git a/test/testlib.py b/test/testlib.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d32e0b512dcb1f7f790b1c211e52987306cd895
--- /dev/null
+++ b/test/testlib.py
@@ -0,0 +1,100 @@
+import loopy as lp
+
+# {{{ test_preamble_with_separate_temporaries
+
+class SeparateTemporariesPreambleTestHelper:
+ def __init__(self, func_name, func_arg_dtypes, func_result_dtypes, arr):
+ self.func_name = func_name
+ self.func_arg_dtypes = func_arg_dtypes
+ self.func_result_dtypes = func_result_dtypes
+ self.arr = arr
+
+ def mangler(self, kernel, name, arg_dtypes):
+ """
+ A function that will return a :class:`loopy.kernel.data.CallMangleInfo`
+ to interface with the calling :class:`loopy.LoopKernel`
+ """
+ if name != self.func_name:
+ return None
+
+ from loopy.types import to_loopy_type
+ from loopy.kernel.data import CallMangleInfo
+
+ def __compare(d1, d2):
+ # compare dtypes ignoring atomic
+ return to_loopy_type(d1, for_atomic=True) == \
+ to_loopy_type(d2, for_atomic=True)
+
+ # check types
+ if len(arg_dtypes) != len(arg_dtypes):
+ raise Exception('Unexpected number of arguments provided to mangler '
+ '{}, expected {}, got {}'.format(
+ self.func_name, len(self.func_arg_dtypes),
+ len(arg_dtypes)))
+
+ for i, (d1, d2) in enumerate(zip(self.func_arg_dtypes, arg_dtypes)):
+ if not __compare(d1, d2):
+ raise Exception('Argument at index {} for mangler {} does not '
+ 'match expected dtype. Expected {}, got {}'.
+ format(i, self.func_name, str(d1), str(d2)))
+
+ # get target for creation
+ target = arg_dtypes[0].target
+ return CallMangleInfo(
+ target_name=self.func_name,
+ result_dtypes=tuple(to_loopy_type(x, target=target) for x in
+ self.func_result_dtypes),
+ arg_dtypes=arg_dtypes)
+
+ def preamble_gen(self, preamble_info):
+ from loopy.kernel.data import temp_var_scope as scopes
+
+ # find a function matching our name
+ func_match = next(
+ (x for x in preamble_info.seen_functions
+ if x.name == self.func_name), None)
+ desc = 'custom_funcs_indirect'
+ if func_match is not None:
+ from loopy.types import to_loopy_type
+ # check types
+ if tuple(to_loopy_type(x) for x in self.func_arg_dtypes) == \
+ func_match.arg_dtypes:
+ # if match, create our temporary
+ var = lp.TemporaryVariable(
+ 'lookup', initializer=self.arr, dtype=self.arr.dtype,
+ shape=self.arr.shape,
+ scope=scopes.GLOBAL, read_only=True)
+ # and code
+ code = """
+ int {name}(int start, int end, int match)
+ {{
+ int result = start;
+ for (int i = start + 1; i < end; ++i)
+ {{
+ if (lookup[i] == match)
+ result = i;
+ }}
+ return result;
+ }}
+ """.format(name=self.func_name)
+
+ # generate temporary variable code
+ from cgen import Initializer
+ from loopy.target.c import generate_array_literal
+ codegen_state = preamble_info.codegen_state.copy(
+ is_generating_device_code=True)
+ kernel = preamble_info.kernel
+ ast_builder = codegen_state.ast_builder
+ target = kernel.target
+ decl_info, = var.decl_info(target, index_dtype=kernel.index_dtype)
+ decl = ast_builder.wrap_global_constant(
+ ast_builder.get_temporary_decl(
+ codegen_state, None, var,
+ decl_info))
+ if var.initializer is not None:
+ decl = Initializer(decl, generate_array_literal(
+ codegen_state, var, var.initializer))
+ # return generated code
+ yield (desc, '\n'.join([str(decl), code]))
+
+# }}}