diff --git a/psyneulink/core/components/component.py b/psyneulink/core/components/component.py index a9ba876ee8b..8200f3c40bf 100644 --- a/psyneulink/core/components/component.py +++ b/psyneulink/core/components/component.py @@ -506,6 +506,7 @@ RESET_STATEFUL_FUNCTION_WHEN, VALUE, VARIABLE from psyneulink.core.globals.log import LogCondition from psyneulink.core.scheduling.time import Time, TimeScale +from psyneulink.core.globals.sampleiterator import SampleIterator from psyneulink.core.globals.parameters import \ Defaults, Parameter, ParameterAlias, ParameterError, ParametersBase, copy_parameter_value from psyneulink.core.globals.preferences.basepreferenceset import BasePreferenceSet, VERBOSE_PREF @@ -1297,7 +1298,7 @@ def _get_compilation_params(self): "input_port_variables", "results", "simulation_results", "monitor_for_control", "feature_values", "simulation_ids", "input_labels_dict", "output_labels_dict", - "modulated_mechanisms", "search_space", "grid", + "modulated_mechanisms", "grid", "activation_derivative_fct", "input_specification", # Shape mismatch "costs", "auto", "hetero", @@ -1314,7 +1315,7 @@ def _is_compilation_param(p): #FIXME: this should use defaults val = p.get() # Check if the value type is valid for compilation - return not isinstance(val, (str, dict, ComponentsMeta, + return not isinstance(val, (str, ComponentsMeta, ContentAddressableList, type(max), type(_is_compilation_param), type(self._get_compilation_params))) @@ -1369,6 +1370,11 @@ def _get_param_initializer(self, context): def _convert(x): if isinstance(x, Enum): return x.value + elif isinstance(x, SampleIterator): + if isinstance(x.generator, list): + return (float(v) for v in x.generator) + else: + return (float(x.start), float(x.step), int(x.num)) try: return (_convert(i) for i in x) except TypeError: diff --git a/psyneulink/core/components/functions/optimizationfunctions.py b/psyneulink/core/components/functions/optimizationfunctions.py index 1b2ab53834c..acffba4faa4 100644 --- a/psyneulink/core/components/functions/optimizationfunctions.py +++ b/psyneulink/core/components/functions/optimizationfunctions.py @@ -1352,10 +1352,10 @@ def reset_grid(self): def _gen_llvm_function(self, *, ctx:pnlvm.LLVMBuilderContext, tags:frozenset): if "select_min" in tags: return self._gen_llvm_select_min_function(ctx=ctx, tags=tags) - if self._is_composition_optimize(): - # self.objective_function may be bound method of - # an OptimizationControlMechanism - ocm = self.objective_function.__self__ + ocm = self._get_optimized_composition() + if ocm is not None: + # self.objective_function may be a bound method of + # OptimizationControlMechanism extra_args = [ctx.get_param_struct_type(ocm.agent_rep).as_pointer(), ctx.get_state_struct_type(ocm.agent_rep).as_pointer(), ctx.get_data_struct_type(ocm.agent_rep).as_pointer()] @@ -1382,66 +1382,46 @@ def _get_input_struct_type(self, ctx): return ctx.convert_python_struct_to_llvm_ir(variable) - def _is_composition_optimize(self): - # self.objective_function may be bound method of - # an OptimizationControlMechanism - return hasattr(self.objective_function, '__self__') + def _get_optimized_composition(self): + # self.objective_function may be a bound method of + # OptimizationControlMechanism + return getattr(self.objective_function, '__self__', None) def _get_param_ids(self): - ids = super()._get_param_ids() + ["search_space"] - if self._is_composition_optimize(): + ids = super()._get_param_ids() + if self._get_optimized_composition() is not None: ids.append("objective_function") return ids - def _get_search_dim_type(self, ctx, d): - if isinstance(d.generator, list): - # Make sure we only generate float values - return ctx.convert_python_struct_to_llvm_ir([float(x) for x in d.generator]) - if isinstance(d, SampleIterator): - return pnlvm.ir.LiteralStructType((ctx.float_ty, ctx.float_ty, ctx.int32_ty)) - assert False, "Unsupported dimension type: {}".format(d) - def _get_param_struct_type(self, ctx): param_struct = ctx.get_param_struct_type(super()) - search_space = (self._get_search_dim_type(ctx, d) for d in self.search_space) - search_space_struct = pnlvm.ir.LiteralStructType(search_space) - if self._is_composition_optimize(): + ocm = self._get_optimized_composition() + if ocm is not None: # self.objective_function is a bound method of - # an OptimizationControlMechanism - ocm = self.objective_function.__self__ + # OptimizationControlMechanism obj_func_params = ocm._get_evaluate_param_struct_type(ctx) return pnlvm.ir.LiteralStructType([*param_struct, - search_space_struct, obj_func_params]) - return pnlvm.ir.LiteralStructType([*param_struct, search_space_struct]) - - def _get_search_dim_init(self, context, d): - if isinstance(d.generator, list): - return tuple(d.generator) - if isinstance(d, SampleIterator): - return (d.start, d.step, d.num) - - assert False, "Unsupported dimension type: {}".format(d) + return param_struct def _get_param_initializer(self, context): param_initializer = super()._get_param_initializer(context) - grid_init = (self._get_search_dim_init(context, d) for d in self.search_space) - if self._is_composition_optimize(): + ocm = self._get_optimized_composition() + if ocm is not None: # self.objective_function is a bound method of - # an OptimizationControlMechanism - ocm = self.objective_function.__self__ + # OptimizationControlMechanism obj_func_param_init = ocm._get_evaluate_param_initializer(context) - return (*param_initializer, tuple(grid_init), obj_func_param_init) + return (*param_initializer, obj_func_param_init) - return (*param_initializer, tuple(grid_init)) + return param_initializer def _get_state_ids(self): ids = super()._get_state_ids() - if self._is_composition_optimize(): + if self._get_optimized_composition() is not None: ids.append("objective_function") return ids @@ -1449,10 +1429,10 @@ def _get_state_ids(self): def _get_state_struct_type(self, ctx): state_struct = ctx.get_state_struct_type(super()) - if self._is_composition_optimize(): + ocm = self._get_optimized_composition() + if ocm is not None: # self.objective_function is a bound method of - # an OptimizationControlMechanism - ocm = self.objective_function.__self__ + # OptimizationControlMechanism obj_func_state = ocm._get_evaluate_state_struct_type(ctx) state_struct = pnlvm.ir.LiteralStructType([*state_struct, obj_func_state]) @@ -1462,10 +1442,10 @@ def _get_state_struct_type(self, ctx): def _get_state_initializer(self, context): state_initializer = super()._get_state_initializer(context) - if self._is_composition_optimize(): + ocm = self._get_optimized_composition() + if ocm is not None: # self.objective_function is a bound method of - # an OptimizationControlMechanism - ocm = self.objective_function.__self__ + # OptimizationControlMechanism obj_func_state_init = ocm._get_evaluate_state_initializer(context) state_initializer = (*state_initializer, obj_func_state_init) @@ -1792,7 +1772,7 @@ def _function(self, format(repr(DIRECTION), self.name, direction) - ocm = self.objective_function.__self__ if self._is_composition_optimize() else None + ocm = self._get_optimized_composition() if ocm is not None and \ ocm.parameters.comp_execution_mode._get(context).startswith("PTX"): opt_sample, opt_value, all_samples, all_values = self._run_cuda_grid(ocm, variable, context) diff --git a/psyneulink/core/components/functions/statefulfunctions/memoryfunctions.py b/psyneulink/core/components/functions/statefulfunctions/memoryfunctions.py index 6f96936e60c..d33d15549bb 100644 --- a/psyneulink/core/components/functions/statefulfunctions/memoryfunctions.py +++ b/psyneulink/core/components/functions/statefulfunctions/memoryfunctions.py @@ -732,7 +732,6 @@ def __init__(self, self.parameters.val_size._set(len(self.previous_value[VALS][0]), Context()) self.has_initializers = True - self.stateful_attributes = ["random_state", "previous_value"] def _get_state_ids(self): return super()._get_state_ids() + ["ring_memory"] @@ -772,14 +771,9 @@ def _gen_llvm_function_body(self, ctx, builder, params, state, arg_in, arg_out, var_val_ptr = builder.gep(arg_in, [ctx.int32_ty(0), ctx.int32_ty(1)]) # Zero output + builder.store(arg_out.type.pointee(None), arg_out) out_key_ptr = builder.gep(arg_out, [ctx.int32_ty(0), ctx.int32_ty(0)]) out_val_ptr = builder.gep(arg_out, [ctx.int32_ty(0), ctx.int32_ty(1)]) - with pnlvm.helpers.array_ptr_loop(builder, out_key_ptr, "zero_key") as (b, i): - out_ptr = b.gep(out_key_ptr, [ctx.int32_ty(0), i]) - b.store(out_ptr.type.pointee(0), out_ptr) - with pnlvm.helpers.array_ptr_loop(builder, out_val_ptr, "zero_val") as (b, i): - out_ptr = b.gep(out_val_ptr, [ctx.int32_ty(0), i]) - b.store(out_ptr.type.pointee(0), out_ptr) # Check retrieval probability retr_ptr = builder.alloca(pnlvm.ir.IntType(1)) @@ -813,7 +807,7 @@ def _gen_llvm_function_body(self, ctx, builder, params, state, arg_in, arg_out, builder.gep(distance_arg_in, [ctx.int32_ty(0), ctx.int32_ty(0)])) selection_arg_in = builder.alloca(pnlvm.ir.ArrayType(distance_f.args[3].type.pointee, max_entries)) - with pnlvm.helpers.for_loop_zero_inc(builder, entries, "distance_loop") as (b,idx): + with pnlvm.helpers.for_loop_zero_inc(builder, entries, "distance_loop") as (b, idx): compare_ptr = b.gep(keys_ptr, [ctx.int32_ty(0), idx]) b.store(b.load(compare_ptr), b.gep(distance_arg_in, [ctx.int32_ty(0), ctx.int32_ty(1)])) @@ -833,7 +827,7 @@ def _gen_llvm_function_body(self, ctx, builder, params, state, arg_in, arg_out, builder.store(ctx.int32_ty(0), selected_idx_ptr) with pnlvm.helpers.for_loop_zero_inc(builder, entries, "distance_loop") as (b,idx): selection_val = b.load(b.gep(selection_arg_out, [ctx.int32_ty(0), idx])) - non_zero = b.fcmp_ordered('!=', selection_val, ctx.float_ty(0)) + non_zero = b.fcmp_ordered('!=', selection_val, selection_val.type(0)) with b.if_then(non_zero): b.store(idx, selected_idx_ptr) @@ -842,8 +836,8 @@ def _gen_llvm_function_body(self, ctx, builder, params, state, arg_in, arg_out, selected_idx])) selected_val = builder.load(builder.gep(vals_ptr, [ctx.int32_ty(0), selected_idx])) - builder.store(selected_key, builder.gep(arg_out, [ctx.int32_ty(0), ctx.int32_ty(0)])) - builder.store(selected_val, builder.gep(arg_out, [ctx.int32_ty(0), ctx.int32_ty(1)])) + builder.store(selected_key, out_key_ptr) + builder.store(selected_val, out_val_ptr) # Check storage probability store_ptr = builder.alloca(pnlvm.ir.IntType(1)) diff --git a/psyneulink/core/components/mechanisms/modulatory/control/optimizationcontrolmechanism.py b/psyneulink/core/components/mechanisms/modulatory/control/optimizationcontrolmechanism.py index 4d1bc86a79f..0be581c04a8 100644 --- a/psyneulink/core/components/mechanisms/modulatory/control/optimizationcontrolmechanism.py +++ b/psyneulink/core/components/mechanisms/modulatory/control/optimizationcontrolmechanism.py @@ -1025,8 +1025,7 @@ def _get_evaluate_param_struct_type(self, ctx): def _get_evaluate_param_initializer(self, context): num_estimates = self.parameters.num_estimates.get(context) or 0 - # FIXME: The intensity cost function is not setup with the right execution id - intensity_cost = tuple(op.intensity_cost_function._get_param_initializer(None) for op in self.output_ports) + intensity_cost = tuple(op.intensity_cost_function._get_param_initializer(context) for op in self.output_ports) return (intensity_cost, num_estimates) def _get_evaluate_state_struct_type(self, ctx): diff --git a/psyneulink/core/llvm/builder_context.py b/psyneulink/core/llvm/builder_context.py index 126cef3a27a..a745bfe12a9 100644 --- a/psyneulink/core/llvm/builder_context.py +++ b/psyneulink/core/llvm/builder_context.py @@ -20,6 +20,7 @@ from typing import Set import weakref from psyneulink.core.scheduling.time import Time +from psyneulink.core.globals.sampleiterator import SampleIterator from psyneulink.core import llvm as pnlvm from . import codegen from .debug import debug_env @@ -305,6 +306,11 @@ def convert_python_struct_to_llvm_ir(self, t): return pnlvm.builtins.get_mersenne_twister_state_struct(self) elif isinstance(t, Time): return ir.ArrayType(self.int32_ty, len(Time._time_scale_attr_map)) + elif isinstance(t, SampleIterator): + if isinstance(t.generator, list): + return ir.ArrayType(self.float_ty, len(t.generator)) + # Generic iterator is {start, increment, count} + return ir.LiteralStructType((self.float_ty, self.float_ty, self.int32_ty)) assert False, "Don't know how to convert {}".format(type(t)) diff --git a/psyneulink/core/llvm/execution.py b/psyneulink/core/llvm/execution.py index d3323800d7f..b411cfba67e 100644 --- a/psyneulink/core/llvm/execution.py +++ b/psyneulink/core/llvm/execution.py @@ -241,7 +241,7 @@ def execute(self, variable): class CompExecution(CUDAExecution): - def __init__(self, composition, execution_ids=[None], additional_tags=frozenset()): + def __init__(self, composition, execution_ids=[None], *, additional_tags=frozenset()): super().__init__(buffers=['state_struct', 'param_struct', 'data_struct', 'conditions']) self._composition = composition self._execution_contexts = [