def __call__(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]))
def __call__(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]))
def generate_code(self, preamble_info): 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, = self.array.decl_info(target, index_dtype=kernel.index_dtype) decl = ast_builder.wrap_global_constant( ast_builder.get_temporary_decl(codegen_state, 1, self.array, decl_info)) if self.array.initializer is not None: decl = Initializer( decl, generate_array_literal(codegen_state, self.array, self.array.initializer)) return '\n'.join([str(decl), self.code])