class OrderedDictLevel(PriceLevel): def is_empty(self): return len(self.orders) == 0 def __init__(self): PriceLevel.__init__(self) self.orders = OrderedDict() def _add(self, order): self.orders[order[O_ID]] = order def _remove(self, order: tuple) -> list: return self.orders.pop(order[O_ID]) def is_not_empty(self) -> bool: return len(self.orders) > 0 def __len__(self): return len(self.orders) def get_first(self) -> list: for order in self.orders.values(): return order def _remove_first(self): self.orders.popitem(last=False) def get_last(self): for order in self.orders.values().__reversed__(): return order def _remove_last(self): self.orders.popitem(last=True)
class _ProcedureWorker(Worker): def __init__(self, cl_environment, compile_flags, cl_function, kernel_data, double_precision, use_local_reduction): super().__init__(cl_environment) self._cl_function = cl_function self._kernel_data = OrderedDict(sorted(kernel_data.items())) self._double_precision = double_precision self._use_local_reduction = use_local_reduction self._mot_float_dtype = np.float32 if double_precision: self._mot_float_dtype = np.float64 for data in self._kernel_data.values(): data.set_mot_float_dtype(self._mot_float_dtype) self._kernel = self._build_kernel(self._get_kernel_source(), compile_flags) self._workgroup_size = self._kernel.run_procedure.get_work_group_info( cl.kernel_work_group_info.PREFERRED_WORK_GROUP_SIZE_MULTIPLE, self._cl_environment.device) if not self._use_local_reduction: self._workgroup_size = 1 self._kernel_inputs = { name: data.get_kernel_inputs(self._cl_context, self._workgroup_size) for name, data in self._kernel_data.items() } def calculate(self, range_start, range_end): nmr_problems = range_end - range_start func = self._kernel.run_procedure func.set_scalar_arg_dtypes(self.get_scalar_arg_dtypes()) kernel_inputs_list = [] for inputs in [ self._kernel_inputs[name] for name in self._kernel_data ]: kernel_inputs_list.extend(inputs) func(self._cl_queue, (int(nmr_problems * self._workgroup_size), ), (int(self._workgroup_size), ), *kernel_inputs_list, global_offset=(int(range_start * self._workgroup_size), )) for name, data in self._kernel_data.items(): data.enqueue_readouts(self._cl_queue, self._kernel_inputs[name], range_start, range_end) def _build_kernel(self, kernel_source, compile_flags=()): """Convenience function for building the kernel for this worker. Args: kernel_source (str): the kernel source to use for building the kernel Returns: cl.Program: a compiled CL kernel """ from mot import configuration if configuration.should_ignore_kernel_compile_warnings(): warnings.simplefilter("ignore") return cl.Program(self._cl_context, kernel_source).build(' '.join(compile_flags)) def _get_kernel_source(self): assignment = '' if self._cl_function.get_return_type() != 'void': assignment = '__results[gid] = ' variable_inits = [] function_call_inputs = [] post_function_callbacks = [] for parameter in self._cl_function.get_parameters(): data = self._kernel_data[parameter.name] call_args = (parameter.name, '_' + parameter.name, 'gid', parameter.data_type.address_space) variable_inits.append(data.initialize_variable(*call_args)) function_call_inputs.append( data.get_function_call_input(*call_args)) post_function_callbacks.append( data.post_function_callback(*call_args)) kernel_source = '' kernel_source += get_float_type_def(self._double_precision) kernel_source += '\n'.join(data.get_type_definitions() for data in self._kernel_data.values()) kernel_source += self._cl_function.get_cl_code() kernel_source += ''' __kernel void run_procedure(''' + ",\n".join(self._get_kernel_arguments()) + '''){ ulong gid = (ulong)(get_global_id(0) / get_local_size(0)); ''' + '\n'.join(variable_inits) + ''' ''' + assignment + ' ' + self._cl_function.get_cl_function_name() + '(' + \ ', '.join(function_call_inputs) + '''); ''' + '\n'.join(post_function_callbacks) + ''' } ''' return kernel_source def _get_kernel_arguments(self): """Get the list of kernel arguments for loading the kernel data elements into the kernel. This will use the sorted keys for looping through the kernel input items. Returns: list of str: the list of parameter definitions """ declarations = [] for name, data in self._kernel_data.items(): declarations.extend(data.get_kernel_parameters('_' + name)) return declarations def get_scalar_arg_dtypes(self): """Get the location and types of the input scalars. Returns: list: for every kernel input element either None if the data is a buffer or the numpy data type if if is a scalar. """ dtypes = [] for name, data in self._kernel_data.items(): dtypes.extend(data.get_scalar_arg_dtypes()) return dtypes