Repository: PrincetonUniversity/PsyNeuLink

Issue:
Compiled DDM Integrator function result is off by a factor of 20
The following code:

```Python
import psyneulink as pnl
import numpy as np
import pytest

decisionMaker = pnl.DDM(function=pnl.DriftDiffusionIntegrator(starting_point=0,
                                                              threshold=1,
                                                              noise=0.0),
                        reinitialize_when=pnl.AtTrialStart(),
                        output_ports=[pnl.DECISION_VARIABLE, pnl.RESPONSE_TIME],
                        name='DDM')

comp = pnl.Composition()
comp.add_node(decisionMaker)

print(comp.run([.05], termination_processing={pnl.TimeScale.TRIAL: pnl.WhenFinished(decisionMaker)},
         bin_execute=False))
```

results in an output of 

> [array([[1.0]], dtype=object), array([[20.0]], dtype=object)]

in Python mode, but

> [[0.05], [1.0]]

in compiled mode - off by a factor of exactly 20.



Discussion:
kmantel: retest this please
kmantel: retest this please
jvesely: /retest this please
SamKG: retest this please
SamKG: retest this please
SamKG: retest this please
jvesely: /retest this please

Patch:
diff --git a/psyneulink/core/compositions/composition.py b/psyneulink/core/compositions/composition.py
index 1de7b15250..cb871a7547 100644
--- a/psyneulink/core/compositions/composition.py
+++ b/psyneulink/core/compositions/composition.py
@@ -7960,7 +7960,7 @@ class Composition(Composition_Base, metaclass=ComponentsMeta):
             *comp_param_type_list))
 
     def _get_state_ids(self):
-        return ["nodes", "projections"] + super()._get_param_ids()
+        return ["nodes", "projections"] + super()._get_state_ids()
 
     def _get_state_struct_type(self, ctx):
         node_state_type_list = (ctx.get_state_struct_type(m) for m in self._all_nodes)
diff --git a/psyneulink/core/llvm/codegen.py b/psyneulink/core/llvm/codegen.py
index 4e0889036f..6425169bcf 100644
--- a/psyneulink/core/llvm/codegen.py
+++ b/psyneulink/core/llvm/codegen.py
@@ -292,9 +292,15 @@ def gen_composition_exec(ctx, composition, *, tags:frozenset):
 
         # Generate a while not 'end condition' loop
         builder.position_at_end(loop_condition)
+        is_finished_flags = {}
+        for idx, node in enumerate(composition.nodes):
+            node_state = builder.gep(state, [ctx.int32_ty(0), ctx.int32_ty(0), ctx.int32_ty(idx)])
+            is_finished_flag_ptr = helpers.get_state_ptr(builder, node, node_state, "is_finished_flag")
+            is_finished_flags[node] = builder.load(is_finished_flag_ptr)
+
         run_cond = cond_gen.generate_sched_condition(
             builder, composition.termination_processing[TimeScale.TRIAL],
-            cond, None)
+            cond, None, is_finished_flags)
         run_cond = builder.not_(run_cond, name="not_run_cond")
 
         loop_body = builder.append_basic_block(name="scheduling_loop_body")
@@ -306,6 +312,11 @@ def gen_composition_exec(ctx, composition, *, tags:frozenset):
 
         zero = ctx.int32_ty(0)
         any_cond = ir.IntType(1)(0)
+        is_finished_flags = {}
+        for idx, node in enumerate(composition.nodes):
+            node_state = builder.gep(state, [ctx.int32_ty(0), ctx.int32_ty(0), ctx.int32_ty(idx)])
+            is_finished_flag_ptr = helpers.get_state_ptr(builder, node, node_state, "is_finished_flag")
+            is_finished_flags[node] = builder.load(is_finished_flag_ptr)
 
         # Calculate execution set before running the mechanisms
         for idx, node in enumerate(composition.nodes):
@@ -314,7 +325,7 @@ def gen_composition_exec(ctx, composition, *, tags:frozenset):
                                            name="run_cond_ptr_" + node.name)
             node_cond = cond_gen.generate_sched_condition(
                 builder, composition._get_processing_condition_set(node),
-                cond, node)
+                cond, node, is_finished_flags)
             ran = cond_gen.generate_ran_this_pass(builder, cond, node)
             node_cond = builder.and_(node_cond, builder.not_(ran),
                                      name="run_cond_" + node.name)
diff --git a/psyneulink/core/llvm/helpers.py b/psyneulink/core/llvm/helpers.py
index 26ebf40943..b2807555e3 100644
--- a/psyneulink/core/llvm/helpers.py
+++ b/psyneulink/core/llvm/helpers.py
@@ -260,12 +260,22 @@ class ConditionGenerator:
         return tuple(data)
 
     def bump_ts(self, builder, cond_ptr, count=(0, 0, 1)):
+        """
+        Increments the time structure of the composition.
+        Count should be a tuple where there is a number in only one spot, and zeroes elsewhere.
+        Indices greater than that of the one are zeroed.
+        """
+        
+        # Validate count tuple
+        assert count.count(0) == len(count) - 1
+
+        # Get timestruct pointer
         ts_ptr = builder.gep(cond_ptr, [self._zero, self._zero, self._zero])
         ts = builder.load(ts_ptr)
 
-        # run, pass, step
+        # Update run, pass, step of ts
         for idx in range(3):
-            if idx == 0 or not all(v == 0 for v in count[:idx]):
+            if all(v == 0 for v in count[:idx]):
                 el = builder.extract_value(ts, idx)
                 el = builder.add(el, self.ctx.int32_ty(count[idx]))
             else:
@@ -342,9 +352,10 @@ class ConditionGenerator:
 
         return builder.icmp_signed("==", node_run, global_run)
 
-    def generate_sched_condition(self, builder, condition, cond_ptr, node):
+    def generate_sched_condition(self, builder, condition, cond_ptr, node, is_finished_flags=None):
+
+        from psyneulink.core.scheduling.condition import All, AllHaveRun, Always, AtPass, AtTrial, EveryNCalls, Never, WhenFinished
 
-        from psyneulink.core.scheduling.condition import All, AllHaveRun, Always, AtPass, AtTrial, EveryNCalls, Never
         if isinstance(condition, Always):
             return ir.IntType(1)(1)
         if isinstance(condition, Never):
@@ -352,13 +363,18 @@ class ConditionGenerator:
         elif isinstance(condition, All):
             agg_cond = ir.IntType(1)(1)
             for cond in condition.args:
-                cond_res = self.generate_sched_condition(builder, cond, cond_ptr, node)
+                cond_res = self.generate_sched_condition(builder, cond, cond_ptr, node, is_finished_flags)
                 agg_cond = builder.and_(agg_cond, cond_res)
             return agg_cond
         elif isinstance(condition, AllHaveRun):
+            # Extract dependencies
+            dependencies = self.composition.nodes
+            if len(condition.args) > 0:
+                dependencies = condition.args
+
             run_cond = ir.IntType(1)(1)
             array_ptr = builder.gep(cond_ptr, [self._zero, self._zero, self.ctx.int32_ty(1)])
-            for node in self.composition.nodes:
+            for node in dependencies:
                 node_ran = self.generate_ran_this_trial(builder, cond_ptr, node)
                 run_cond = builder.and_(run_cond, node_ran)
             return run_cond
@@ -397,5 +413,12 @@ class ConditionGenerator:
 
             # Return: target.calls % N == 0 AND me.last_time < target.last_time
             return builder.and_(completedNruns, ran_after_me)
+        
+        elif isinstance(condition, WhenFinished):
+            target = condition.args[0]
+            assert target in is_finished_flags
+            target_is_finished = is_finished_flags[target]
+            
+            return builder.fcmp_ordered("==", target_is_finished, target_is_finished.type(1))
 
         assert False, "Unsupported scheduling condition: {}".format(condition)
diff --git a/psyneulink/library/components/mechanisms/processing/integrator/ddm.py b/psyneulink/library/components/mechanisms/processing/integrator/ddm.py
index 2e786cd642..c1829ee779 100644
--- a/psyneulink/library/components/mechanisms/processing/integrator/ddm.py
+++ b/psyneulink/library/components/mechanisms/processing/integrator/ddm.py
@@ -1195,3 +1195,36 @@ class DDM(ProcessingMechanism):
             )
             return True
         return False
+    
+    def _gen_llvm_is_finished_cond(self, ctx, builder, params, state, current):
+        # Setup pointers to internal function
+        func_state_ptr = pnlvm.helpers.get_state_ptr(builder, self, state, 'function')
+        func_param_ptr = pnlvm.helpers.get_state_ptr(builder, self, params, 'function')
+
+        # Find the single numeric entry in previous_value
+        try:
+            prev_val_ptr = pnlvm.helpers.get_state_ptr(builder, self.function, func_state_ptr, "previous_value")
+        except ValueError:
+            return pnlvm.ir.IntType(1)(1)
+
+        # Extract scalar value from ptr
+        prev_val_ptr = builder.gep(prev_val_ptr, [ctx.int32_ty(0), ctx.int32_ty(0), ctx.int32_ty(0)])
+        prev_val = builder.load(prev_val_ptr)
+
+        # Take abs of previous val
+        llvm_fabs = ctx.get_builtin("fabs", [ctx.float_ty])
+        prev_val = builder.call(llvm_fabs, [prev_val])
+
+
+        # obtain threshold value
+        threshold_ptr = pnlvm.helpers.get_param_ptr(builder,
+                                                    self.function,
+                                                    func_param_ptr,
+                                                    "threshold")
+
+        threshold_ptr = builder.gep(threshold_ptr, [ctx.int32_ty(0), ctx.int32_ty(0)])
+        threshold = builder.load(threshold_ptr)
+        is_prev_greater_or_equal = builder.fcmp_ordered('>=', prev_val, threshold)
+
+        return is_prev_greater_or_equal
+


Test patch:
diff --git a/tests/composition/test_composition.py b/tests/composition/test_composition.py
index 61633f5bcf..13f046afb1 100644
--- a/tests/composition/test_composition.py
+++ b/tests/composition/test_composition.py
@@ -3436,6 +3436,62 @@ class TestSystemComposition:
 #         assert 16 == output[0][0]
 #
 
+
+class TestSchedulerConditions:
+    @pytest.mark.composition
+    @pytest.mark.parametrize("mode",['Python',
+                            #  pytest.param('LLVM', marks=pytest.mark.llvm), #FIXME: Fails for `LLVM` and `LLVMExec` modes?
+                            #  pytest.param('LLVMExec', marks=pytest.mark.llvm),
+                             pytest.param('LLVMRun', marks=pytest.mark.llvm),
+                             pytest.param('PTXExec', marks=[pytest.mark.llvm, pytest.mark.cuda]),
+                             pytest.param('PTXRun', marks=[pytest.mark.llvm, pytest.mark.cuda])
+                             ])
+    @pytest.mark.parametrize(["condition", "expected_result"],[
+                             (pnl.EveryNCalls, [[.25, .25]]),
+                             (pnl.WhenFinished, [[1.0, 1.0]]),
+                             (pnl.All, [[1.0, 1.0]]),
+                             (pnl.AllHaveRun, [[.05, .05]]),
+                             (pnl.Always, [[0.05, 0.05]]),
+                            #  (pnl.AtPass, [[.3, .3]]), #FIXME: Differing result between llvm and python
+                             (pnl.AtTrial,[[0.05, 0.05]]),
+                            #  (pnl.Never), #TODO: Find a good test case for this!
+                            ])
+    def test_scheduler_conditions(self, mode, condition, expected_result):
+        print(mode, condition)
+        decisionMaker = pnl.DDM(
+                        function=pnl.DriftDiffusionIntegrator(starting_point=0,
+                                                              threshold=1,
+                                                              noise=0.0),
+                        reinitialize_when=pnl.AtTrialStart(),
+                        output_ports=[pnl.DECISION_VARIABLE, pnl.RESPONSE_TIME],
+                        name='DDM')
+
+        response = pnl.ProcessingMechanism(size=2, name="GATE")
+
+        comp = pnl.Composition()
+        comp.add_node(decisionMaker)
+        comp.add_node(response)
+        comp.add_projection(pnl.MappingProjection(), sender=decisionMaker, receiver=response)
+
+        if condition is pnl.EveryNCalls:
+            comp.scheduler.add_condition(response, condition(decisionMaker, 5))
+        elif condition is pnl.WhenFinished:
+            comp.scheduler.add_condition(response, condition(decisionMaker))
+        elif condition is pnl.All:
+            comp.scheduler.add_condition(response, condition(pnl.WhenFinished(decisionMaker)))
+        elif condition is pnl.AllHaveRun:
+            comp.scheduler.add_condition(response, condition(decisionMaker))
+        elif condition is pnl.Always:
+            comp.scheduler.add_condition(response, condition())
+        elif condition is pnl.AtPass:
+            comp.scheduler.add_condition(response, condition(5))
+        elif condition is pnl.AtTrial:
+            comp.scheduler.add_condition(response, condition(0))
+        result = comp.run([0.05], bin_execute=mode)
+        result = [x for x in np.array(result).flatten()] #HACK: The result is an object dtype in Python mode for some reason?
+        assert np.allclose(result, np.array(expected_result).flatten())
+
+
 class TestNestedCompositions:
 
     @pytest.mark.composition
diff --git a/tests/mechanisms/test_ddm_mechanism.py b/tests/mechanisms/test_ddm_mechanism.py
index c54ee3ab49..5dac944757 100644
--- a/tests/mechanisms/test_ddm_mechanism.py
+++ b/tests/mechanisms/test_ddm_mechanism.py
@@ -713,3 +713,27 @@ def test_DDM_threshold_modulation(mode):
     val = np.asfarray(val)
     assert np.allclose(val[0], [60.0])
     assert np.allclose(val[1], [60.2])
+
+@pytest.mark.parametrize("mode", ['Python',
+                                    pytest.param('LLVM', marks=pytest.mark.llvm),
+                                    pytest.param('LLVMExec', marks=pytest.mark.llvm),
+                                    pytest.param('LLVMRun', marks=pytest.mark.llvm),
+                                    pytest.param('PTXExec', marks=[pytest.mark.llvm, pytest.mark.cuda]),
+                                    pytest.param('PTXRun', marks=[pytest.mark.llvm, pytest.mark.cuda])
+                                    ])
+@pytest.mark.parametrize(["noise", "threshold", "expected_results"],[
+                            (1.0, 0.0, (0.0, 1.0)),
+                            (1.5, 2, (-2.0, 1.0)),
+                            (10.0, 10.0, (10.0, 29.0)),
+                            (100.0, 100.0, (100.0, 76.0)),
+                        ])
+def test_ddm_is_finished(mode, noise, threshold, expected_results):
+    comp = Composition()
+    ddm = DDM(execute_until_finished=True,
+                function=DriftDiffusionIntegrator(threshold=threshold, noise=noise))
+    comp.add_node(ddm)
+
+    results = comp.run([0], bin_execute=mode)
+
+    results = [x for x in np.array(results).flatten()] #HACK: The result is an object dtype in Python mode for some reason?
+    assert np.allclose(results, np.array(expected_results).flatten())