def _prep_all_data(self):
     # -- replace the numpy-allocated RaggedArray with OpenCL one
     self.all_data = CLRaggedArray(self.queue, self.all_data)

 def RaggedArray(self, listofarrays, **kwargs):
     return CLRaggedArray.from_arrays(self.queue, listofarrays, **kwargs)

#.........这里部分代码省略.........
        if (i < ${shape0} && j == 0)
            ybuf[i] = ${float_alpha} * (sums[i][0] + sums[i][1]);
    }
    """

    text = as_ascii(Template(text, output_encoding='ascii').render(**textconf))
    kernel = cl.Program(p.queue.context, text).build().fn
    kernel.set_args(*[arr.data for arr in full_args])

    plan = Plan(p.queue, kernel, gsize, lsize,
                name='clra_gemv.block_impl',
                tag=p.tag,
                bw_per_call=bw_from_geometry(p.geometry, items),
                flops_per_call=flops_from_geometry(p.geometry, items),
                )
    plan.full_args = full_args  # prevent GC the args
    plan.description = p.geometry_summary(items)
    plan.Ybuf = clYbuf

    # --- Reduce kernel
    align = False

    Nreduce = len(Yshape0s_reduce)
    clYshape0s_reduce = to_device(
        p.queue, np.array(Yshape0s_reduce, dtype=np.int32))
    clYinstride0s_reduce = to_device(
        p.queue, np.array(Yinstride0s_reduce, dtype=np.int32))
    clYinstarts_reduce = to_device(
        p.queue, np.array(Yinstarts_reduce, dtype=np.int32))
    clYstride0s_reduce = to_device(
        p.queue, np.array(Ystride0s_reduce, dtype=np.int32))
    clYstarts_reduce = to_device(
        p.queue, np.array(Ystarts_reduce, dtype=np.int32))
    clYbufinds_reduce = CLRaggedArray.from_arrays(
        p.queue, Ybufinds_reduce, dtype=np.int32, align=align)
    assert len(clYbufinds_reduce) == Nreduce
    assert (clYbufinds_reduce.shape1s == 1).all()

    textconf_reduce = dict(
        Ybuf=clYbuf,
        Yin=p.Y_in,
        Y=p.Y,
        float_beta=p.float_beta,
        float_gamma=p.float_gamma,
    )

    full_args_reduce = (
        clYshape0s_reduce,
        clYbufinds_reduce.cl_shape0s,
        clYbufinds_reduce.cl_starts,
        clYbufinds_reduce.cl_buf,
        clYbuf,
        clYinstride0s_reduce,
        clYinstarts_reduce,
        p.Y_in.cl_buf,
        clYstride0s_reduce,
        clYstarts_reduce,
        p.Y.cl_buf,
    )

    lsize_reduce = None
    gsize_reduce = (block_y, Nreduce)

    text_reduce = """
    __kernel void reduce(
        __global const int *shape0s,

#.........这里部分代码省略.........
            cache = get_default_decoder_cache()
        else:
            cache = model.decoder_cache

        with cache, Timer() as nengo_timer:
            if model is None:
                self.model = Model(dt=float(dt),
                                   label="%s, dt=%f" % (network, dt),
                                   decoder_cache=cache)
            else:
                self.model = model

            if network is not None:
                # Build the network into the model
                self.model.build(network)

            cache.shrink()

        logger.info("Nengo build in %0.3f s" % nengo_timer.duration)

        # --- operators
        with Timer() as planner_timer:
            operators = list(self.model.operators)

            # convert DotInc and Copy to MultiDotInc
            operators = list(map(MultiDotInc.convert_to, operators))
            operators = MultiDotInc.compress(operators)

            # plan the order of operations, combining where appropriate
            op_groups = planner(operators)
            assert len([typ for typ, _ in op_groups if typ is Reset]) < 2, (
                "All resets not planned together")

            self.operators = operators
            self.op_groups = op_groups

        logger.info("Planning in %0.3f s" % planner_timer.duration)

        with Timer() as signals_timer:
            # Initialize signals
            all_signals = stable_unique(
                sig for op in operators for sig in op.all_signals)
            all_bases = stable_unique(sig.base for sig in all_signals)

            sigdict = SignalDict()  # map from Signal.base -> ndarray
            for op in operators:
                op.init_signals(sigdict)

            # Add built states to the probe dictionary
            self._probe_outputs = dict(self.model.params)

            # Provide a nicer interface to probe outputs
            self.data = ProbeDict(self._probe_outputs)

            # Create data on host and add views
            self.all_data = RaggedArray(
                [sigdict[sb] for sb in all_bases],
                names=[getattr(sb, 'name', '') for sb in all_bases],
                dtype=np.float32)

            view_builder = ViewBuilder(all_bases, self.all_data)
            view_builder.setup_views(operators)
            for probe in self.model.probes:
                view_builder.append_view(self.model.sig[probe]['in'])
            view_builder.add_views_to(self.all_data)

            self.all_bases = all_bases
            self.sidx = {
                k: np.int32(v) for k, v in iteritems(view_builder.sidx)}
            self._A_views = view_builder._A_views
            self._X_views = view_builder._X_views
            self._YYB_views = view_builder._YYB_views
            del view_builder

            # Copy data to device
            self.all_data = CLRaggedArray(self.queue, self.all_data)

        logger.info("Signals in %0.3f s" % signals_timer.duration)

        # --- set seed
        self.seed = np.random.randint(npext.maxint) if seed is None else seed
        self._reset_rng()

        # --- create list of plans
        self._raggedarrays_to_reset = {}
        self._cl_rngs = {}

        with Timer() as plans_timer:
            self._plan = []
            for op_type, op_list in op_groups:
                self._plan.extend(self.plan_op_group(op_type, op_list))
            self._plan.extend(self.plan_probes())

        logger.info("Plans in %0.3f s" % plans_timer.duration)

        # -- create object to execute list of plans
        self._plans = Plans(self._plan, self.profiling)

        self._reset_cl_rngs()
        self._probe_step_time()

class Simulator(nengo.Simulator):

    unsupported = []

    def Array(self, val, dtype=np.float32):
        return to_device(self.queue, np.asarray(val, dtype=dtype))

    def RaggedArray(self, listofarrays, **kwargs):
        return CLRaggedArray.from_arrays(self.queue, listofarrays, **kwargs)

    def __init__(self, network, dt=0.001, seed=None, model=None, context=None,
                 n_prealloc_probes=32, profiling=None, ocl_only=False,
                 planner=greedy_planner):
        # --- create these first since they are used in __del__
        self.closed = False
        self.model = None

        # --- check version
        if nengo.version.version_info[:2] != latest_nengo_version_info[:2]:
            raise ValueError(
                "This simulator only supports Nengo %s.x (got %s)" %
                ('.'.join(str(i) for i in latest_nengo_version_info[:2]),
                 nengo.__version__))
        elif nengo.version.version_info > latest_nengo_version_info:
            warnings.warn("This version of `nengo_ocl` has not been tested "
                          "with your `nengo` version (%s). The latest fully "
                          "supported version is %s" % (
                              nengo.__version__, latest_nengo_version))

        # --- arguments/attributes
        if context is None:
            print('No context argument was provided to nengo_ocl.Simulator')
            print("Calling pyopencl.create_some_context() for you now:")
            context = cl.create_some_context()
        if profiling is None:
            profiling = int(os.getenv("NENGO_OCL_PROFILING", 0))
        self.context = context
        self.profiling = profiling
        if self.profiling:
            self.queue = cl.CommandQueue(context, properties=PROFILING_ENABLE)
        else:
            self.queue = cl.CommandQueue(context)

        self.n_prealloc_probes = n_prealloc_probes
        self.ocl_only = ocl_only

        # --- Nengo build
        if model is None or model.decoder_cache is None:
            cache = get_default_decoder_cache()
        else:
            cache = model.decoder_cache

        with cache, Timer() as nengo_timer:
            if model is None:
                self.model = Model(dt=float(dt),
                                   label="%s, dt=%f" % (network, dt),
                                   decoder_cache=cache)
            else:
                self.model = model

            if network is not None:
                # Build the network into the model
                self.model.build(network)

            cache.shrink()

        logger.info("Nengo build in %0.3f s" % nengo_timer.duration)

        # --- operators
        with Timer() as planner_timer:
            operators = list(self.model.operators)

            # convert DotInc and Copy to MultiDotInc
            operators = list(map(MultiDotInc.convert_to, operators))
            operators = MultiDotInc.compress(operators)

            # plan the order of operations, combining where appropriate
            op_groups = planner(operators)
            assert len([typ for typ, _ in op_groups if typ is Reset]) < 2, (
                "All resets not planned together")

            self.operators = operators
            self.op_groups = op_groups

        logger.info("Planning in %0.3f s" % planner_timer.duration)

        with Timer() as signals_timer:
            # Initialize signals
            all_signals = stable_unique(
                sig for op in operators for sig in op.all_signals)
            all_bases = stable_unique(sig.base for sig in all_signals)

            sigdict = SignalDict()  # map from Signal.base -> ndarray
            for op in operators:
                op.init_signals(sigdict)

            # Add built states to the probe dictionary
            self._probe_outputs = dict(self.model.params)

            # Provide a nicer interface to probe outputs
#.........这里部分代码省略.........

#.........这里部分代码省略.........
        self.if_python_code = if_python_code
        self.n_prealloc_probes = n_prealloc_probes
        self.progress_bar = progress_bar

        # --- Nengo build
        with Timer() as nengo_timer:
            if model is None:
                self.model = Model(dt=float(dt),
                                   label="%s, dt=%f" % (network, dt),
                                   decoder_cache=get_default_decoder_cache())
            else:
                self.model = model

            if network is not None:
                # Build the network into the model
                self.model.build(network)

        logger.info("Nengo build in %0.3f s" % nengo_timer.duration)

        # --- operators
        with Timer() as planner_timer:
            operators = list(self.model.operators)

            # convert DotInc and Copy to MultiDotInc
            operators = list(map(MultiDotInc.convert_to, operators))
            operators = MultiDotInc.compress(operators)

            # plan the order of operations, combining where appropriate
            op_groups = planner(operators)
            assert len([typ for typ, _ in op_groups if typ is Reset]) < 2, (
                "All resets not planned together")

            self.operators = operators
            self.op_groups = op_groups

        logger.info("Planning in %0.3f s" % planner_timer.duration)

        with Timer() as signals_timer:
            # Initialize signals
            all_signals = stable_unique(
                sig for op in operators for sig in op.all_signals)
            all_bases = stable_unique(sig.base for sig in all_signals)

            sigdict = SignalDict()  # map from Signal.base -> ndarray
            for op in operators:
                op.init_signals(sigdict)

            # Add built states to the probe dictionary
            self._probe_outputs = dict(self.model.params)

            # Provide a nicer interface to probe outputs
            self.data = ProbeDict(self._probe_outputs)

            # Create data on host and add views
            self.all_data = RaggedArray(
                [sigdict[sb] for sb in all_bases],
                names=[getattr(sb, 'name', '') for sb in all_bases],
                dtype=np.float32)

            view_builder = ViewBuilder(all_bases, self.all_data)
            view_builder.setup_views(operators)
            for probe in self.model.probes:
                view_builder.append_view(self.model.sig[probe]['in'])
            view_builder.add_views_to(self.all_data)

            self.all_bases = all_bases
            self.sidx = {
                k: np.int32(v) for k, v in iteritems(view_builder.sidx)}
            self._A_views = view_builder._A_views
            self._X_views = view_builder._X_views
            self._YYB_views = view_builder._YYB_views
            del view_builder

            # Copy data to device
            self.all_data = CLRaggedArray(self.queue, self.all_data)

        logger.info("Signals in %0.3f s" % signals_timer.duration)

        # --- set seed
        self.seed = np.random.randint(npext.maxint) if seed is None else seed
        self.rng = np.random.RandomState(self.seed)

        # --- create list of plans
        self._raggedarrays_to_reset = {}
        self._cl_rngs = {}
        self._python_rngs = {}

        plans = []
        with Timer() as plans_timer:
            for op_type, op_list in op_groups:
                plans.extend(self.plan_op_group(op_type, op_list))
            plans.extend(self.plan_probes())

        logger.info("Plans in %0.3f s" % plans_timer.duration)

        # -- create object to execute list of plans
        self._plans = Plans(plans, self.profiling)

        self.rng = None  # all randomness set, should no longer be used
        self._reset_probes()  # clears probes from previous model builds

class Simulator(object):
    """Simulator for running Nengo models in OpenCL.

    Parameters
    ----------
    network, dt, seed, model
        These parameters are the same as in `nengo.Simulator`.
    context : `pyopencl.Context` (optional)
        OpenCL context specifying which device(s) to run on. By default, we
        will create a context by calling `pyopencl.create_some_context`
        and use this context as the default for all subsequent instances.
    n_prealloc_probes : int (optional)
        Number of timesteps to buffer when probing. Larger numbers mean less
        data transfer with the device (faster), but use more device memory.
    profiling : boolean (optional)
        If ``True``, ``print_profiling()`` will show profiling information.
        By default, will check the environment variable ``NENGO_OCL_PROFILING``
    if_python_code : 'none' | 'warn' | 'error'
        How the simulator should react if a Python function cannot be converted
        to OpenCL code.
    planner : callable
        A function to plan operator order. See ``nengo_ocl.planners``.
    """

    # --- Store the result of create_some_context so we don't recreate it
    some_context = None

    # 'unsupported' defines features unsupported by a simulator.
    # The format is a list of tuples of the form `(test, reason)` with `test`
    # being a string with wildcards (*, ?, [abc], [!abc]) matched against Nengo
    # test paths and names, and `reason` is a string describing why the feature
    # is not supported by the backend. For example:
    #     unsupported = [('test_pes*', 'PES rule not implemented')]
    # would skip all test whose names start with 'test_pes'.
    unsupported = [
        # advanced indexing
        ('nengo/tests/test_connection.py:test_list_indexing*',
         "Advanced indexing with repeated indices not implemented"),

        # neuron types
        ('nengo/tests/test_neurons.py:test_izhikevich',
         "Izhikevich neurons not implemented"),
        ('nengo/tests/test_neurons.py:test_lif_min_voltage*',
         "Min voltage not implemented"),

        # nodes
        ('nengo/tests/test_node.py:test_none',
         "No error if nodes output None"),
        ('nengo/tests/test_node.py:test_invalid_values*',
         "No error for invalid node values"),
        ('nengo/tests/test_neurons.py:test_direct_mode_nonfinite_value',
         "No error for non-finite values"),

        # processes
        ('nengo/tests/test_processes.py:test_brownnoise',
         "Filtered noise processes not yet implemented"),
        ('nengo/tests/test_ensemble.py:test_noise_copies_ok*',
         "Filtered noise processes not yet implemented"),
        ('nengo/tests/test_simulator.py:test_noise_copies_ok',
         "Filtered noise processes not yet implemented"),
        ('nengo/tests/test_processes.py:TestPiecewise.test_interpolation_?d',
         "float32 rounding issues"),

        # synapses
        ('nengo/tests/test_synapses.py:test_triangle',
         "Only linear filters implemented"),

        # learning rules
        ('nengo/tests/test_learning_rules.py:test_custom_type',
         "Copying 2-D arrays not implemented"),

        # simulator features
        ('nengo/tests/test_simulator.py:test_probe_cache',
         "Changing simulator seed not implemented"),

        # specific to nengo.Simulator (functionality does not need testing)
        ('nengo/tests/test_builder.py:test_commonsig_readonly',
         "Specific to nengo.Simulator"),
    ]

    def Array(self, val, dtype=np.float32):
        return to_device(self.queue, np.asarray(val, dtype=dtype))

    def RaggedArray(self, listofarrays, **kwargs):
        return CLRaggedArray.from_arrays(self.queue, listofarrays, **kwargs)

    def __init__(self, network, dt=0.001, seed=None, model=None, context=None,
                 n_prealloc_probes=32, profiling=None, if_python_code='none',
                 planner=greedy_planner, progress_bar=True):
        # --- check version
        if nengo.version.version_info in bad_nengo_versions:
            raise ValueError(
                "This simulator does not support Nengo version %s. Upgrade "
                "with 'pip install --upgrade --no-deps nengo'."
                % nengo.__version__)
        elif nengo.version.version_info > latest_nengo_version_info:
            warnings.warn("This version of `nengo_ocl` has not been tested "
                          "with your `nengo` version (%s). The latest fully "
                          "supported version is %s" % (
                              nengo.__version__, latest_nengo_version))
#.........这里部分代码省略.........

class Simulator(sim_npy.Simulator):

    def Array(self, val, dtype=np.float32):
        return to_device(self.queue, np.asarray(val, dtype=dtype))

    def RaggedArray(self, listofarrays, **kwargs):
        return CLRaggedArray.from_arrays(self.queue, listofarrays, **kwargs)

    def __init__(self, network, dt=0.001, seed=None, model=None, context=None,
                 n_prealloc_probes=32, profiling=None, ocl_only=False):
        if context is None:
            print('No context argument was provided to sim_ocl.Simulator')
            print("Calling pyopencl.create_some_context() for you now:")
            context = cl.create_some_context()
        if profiling is None:
            profiling = int(os.getenv("NENGO_OCL_PROFILING", 0))
        self.context = context
        self.profiling = profiling
        if self.profiling:
            self.queue = cl.CommandQueue(context, properties=PROFILING_ENABLE)
        else:
            self.queue = cl.CommandQueue(context)

        self.n_prealloc_probes = n_prealloc_probes
        self.ocl_only = ocl_only
        self.cl_rng_state = None

        # -- allocate data
        sim_npy.Simulator.__init__(
            self, network=network, dt=dt, seed=seed, model=model)

        # -- create object to execute list of plans
        self._plans = Plans(self._plan, self.profiling)

    def _init_cl_rng(self):
        if self.cl_rng_state is None:
            self.cl_rng_state = init_rng(self.queue, self.seed)

    def _prep_all_data(self):
        # -- replace the numpy-allocated RaggedArray with OpenCL one
        self.all_data = CLRaggedArray(self.queue, self.all_data)

    def plan_ragged_gather_gemv(self, *args, **kwargs):
        return plan_ragged_gather_gemv(self.queue, *args, **kwargs)

    def plan_TimeUpdate(self, ops):
        op, = ops
        step = self.all_data[[self.sidx[op.step]]]
        time = self.all_data[[self.sidx[op.time]]]
        return [plan_timeupdate(self.queue, step, time, self.model.dt)]

    def plan_Reset(self, ops):
        targets = self.all_data[[self.sidx[op.dst] for op in ops]]
        values = self.Array([op.value for op in ops])
        return [plan_reset(self.queue, targets, values)]

    def plan_SlicedCopy(self, ops):
        copies, ops = split(
            ops, lambda op: op.a_slice is Ellipsis and op.b_slice is Ellipsis)

        plans = []
        if copies:
            A = self.all_data[[self.sidx[op.a] for op in copies]]
            B = self.all_data[[self.sidx[op.b] for op in copies]]
            incs = np.array([op.inc for op in copies], dtype=np.int32)
            plans.append(plan_copy(self.queue, A, B, incs))

        if ops:
            A = self.all_data[[self.sidx[op.a] for op in ops]]
            B = self.all_data[[self.sidx[op.b] for op in ops]]
            inds = lambda ary, i: np.arange(ary.size, dtype=np.int32)[i]
            Ainds = self.RaggedArray([inds(op.a, op.a_slice) for op in ops])
            Binds = self.RaggedArray([inds(op.b, op.b_slice) for op in ops])
            incs = np.array([op.inc for op in ops], dtype=np.int32)
            plans.append(plan_slicedcopy(self.queue, A, B, Ainds, Binds, incs))

        return plans

    def plan_ElementwiseInc(self, ops):
        A = self.all_data[[self.sidx[op.A] for op in ops]]
        X = self.all_data[[self.sidx[op.X] for op in ops]]
        Y = self.all_data[[self.sidx[op.Y] for op in ops]]
        return [plan_elementwise_inc(self.queue, A, X, Y)]

    def plan_SimPyFunc(self, ops):
        # TODO: test with a hybrid program (Python and OCL)

        # group nonlinearities
        unique_ops = OrderedDict()
        for op in ops:
            # assert op.n_args in (1, 2), op.n_args
            op_key = (op.fn, op.t_in, op.x is not None)
            if op_key not in unique_ops:
                unique_ops[op_key] = {'in': [], 'out': []}
            unique_ops[op_key]['in'].append(op.x)
            unique_ops[op_key]['out'].append(op.output)

        # make plans
        plans = []
        for (fn, t_in, x_in), signals in unique_ops.items():
#.........这里部分代码省略.........

def test_lif_rate(n_elements):
    """Test the `lif_rate` nonlinearity"""
    # n_neurons = [3, 3, 3]
    n_neurons = [123459, 23456, 34567]
    N = len(n_neurons)
    J = RA([np.random.normal(loc=1, scale=10, size=n) for n in n_neurons])
    R = RA([np.zeros(n) for n in n_neurons])

    ref = 2e-3
    taus = list(np.random.uniform(low=15e-3, high=80e-3, size=len(n_neurons)))

    queue = cl.CommandQueue(ctx)
    clJ = CLRA(queue, J)
    clR = CLRA(queue, R)
    clTau = CLRA(queue, RA(taus))

    ### simulate host
    nls = [LIF(n, tau_ref=ref, tau_rc=taus[i])
           for i, n in enumerate(n_neurons)]
    for i, nl in enumerate(nls):
        nl.gain = 1
        nl.bias = 0
        R[i] = nl.rates(J[i].flatten()).reshape((-1,1))

    ### simulate device
    plan = plan_lif_rate(queue, clJ, clR, ref, clTau, dt=1,
                         n_elements=n_elements)
    plan()

    rate_sum = np.sum([np.sum(r) for r in R])
    if rate_sum < 1.0:
        logger.warn("LIF rate was not tested above the firing threshold!")
    assert ra.allclose(J, clJ.to_host())
    assert ra.allclose(R, clR.to_host())

def test_lif_rate(n_elements):
    """Test the `lif_rate` nonlinearity"""
    rng = np.random
    dt = 1e-3

    n_neurons = [123459, 23456, 34567]
    J = RA([rng.normal(loc=1, scale=10, size=n) for n in n_neurons])
    R = RA([np.zeros(n) for n in n_neurons])

    ref = 2e-3
    taus = list(rng.uniform(low=15e-3, high=80e-3, size=len(n_neurons)))

    queue = cl.CommandQueue(ctx)
    clJ = CLRA(queue, J)
    clR = CLRA(queue, R)
    clTau = CLRA(queue, RA(taus))

    # simulate host
    nls = [LIFRate(tau_ref=ref, tau_rc=taus[i])
           for i, n in enumerate(n_neurons)]
    for i, nl in enumerate(nls):
        nl.step_math(dt, J[i], R[i])

    # simulate device
    plan = plan_lif_rate(queue, clJ, clR, ref, clTau, dt=dt,
                         n_elements=n_elements)
    plan()

    rate_sum = np.sum([np.sum(r) for r in R])
    if rate_sum < 1.0:
        logger.warn("LIF rate was not tested above the firing threshold!")
    assert ra.allclose(J, clJ.to_host())
    assert ra.allclose(R, clR.to_host())

def test_small():
    n = 3
    sizes = [3] * 3
    vals = [np.random.normal(size=size) for size in sizes]
    A = RA(vals)

    queue = cl.CommandQueue(ctx)
    clA = CLRA(queue, A)
    assert ra.allclose(A, clA.to_host())

def test_lif_step(upsample, n_elements):
    """Test the lif nonlinearity, comparing one step with the Numpy version."""
    dt = 1e-3
    # n_neurons = [3, 3, 3]
    n_neurons = [12345, 23456, 34567]
    N = len(n_neurons)
    J = RA([np.random.normal(scale=1.2, size=n) for n in n_neurons])
    V = RA([np.random.uniform(low=0, high=1, size=n) for n in n_neurons])
    W = RA([np.random.uniform(low=-5*dt, high=5*dt, size=n) for n in n_neurons])
    OS = RA([np.zeros(n) for n in n_neurons])

    ref = 2e-3
    # tau = 20e-3
    # refs = list(np.random.uniform(low=1.7e-3, high=4.2e-3, size=len(n_neurons)))
    taus = list(np.random.uniform(low=15e-3, high=80e-3, size=len(n_neurons)))

    queue = cl.CommandQueue(ctx)
    clJ = CLRA(queue, J)
    clV = CLRA(queue, V)
    clW = CLRA(queue, W)
    clOS = CLRA(queue, OS)
    # clRef = CLRA(queue, RA(refs))
    clTau = CLRA(queue, RA(taus))

    ### simulate host
    nls = [LIF(n, tau_ref=ref, tau_rc=taus[i])
           for i, n in enumerate(n_neurons)]
    for i, nl in enumerate(nls):
        if upsample <= 1:
            nl.step_math(dt, J[i], V[i], W[i], OS[i])
        else:
            s = np.zeros_like(OS[i])
            for j in xrange(upsample):
                nl.step_math(dt/upsample, J[i], V[i], W[i], s)
                OS[i] = (OS[i] > 0.5) | (s > 0.5)

    ### simulate device
    plan = plan_lif(queue, clJ, clV, clW, clV, clW, clOS, ref, clTau, dt,
                    n_elements=n_elements, upsample=upsample)
    plan()

    if 1:
        a, b = V, clV
        for i in xrange(len(a)):
            nc, _ = not_close(a[i], b[i]).nonzero()
            if len(nc) > 0:
                j = nc[0]
                print "i", i, "j", j
                print "J", J[i][j], clJ[i][j]
                print "V", V[i][j], clV[i][j]
                print "W", W[i][j], clW[i][j]
                print "...", len(nc) - 1, "more"

    n_spikes = np.sum([np.sum(os) for os in OS])
    if n_spikes < 1.0:
        logger.warn("LIF spiking mechanism was not tested!")
    assert ra.allclose(J, clJ.to_host())
    assert ra.allclose(V, clV.to_host())
    assert ra.allclose(W, clW.to_host())
    assert ra.allclose(OS, clOS.to_host())

def test_lif_step(upsample):
    """Test the lif nonlinearity, comparing one step with the Numpy version."""
    rng = np.random

    dt = 1e-3
    n_neurons = [12345, 23456, 34567]
    J = RA([rng.normal(scale=1.2, size=n) for n in n_neurons])
    V = RA([rng.uniform(low=0, high=1, size=n) for n in n_neurons])
    W = RA([rng.uniform(low=-5 * dt, high=5 * dt, size=n) for n in n_neurons])
    OS = RA([np.zeros(n) for n in n_neurons])

    ref = 2e-3
    taus = rng.uniform(low=15e-3, high=80e-3, size=len(n_neurons))

    queue = cl.CommandQueue(ctx)
    clJ = CLRA(queue, J)
    clV = CLRA(queue, V)
    clW = CLRA(queue, W)
    clOS = CLRA(queue, OS)
    clTaus = CLRA(queue, RA([t * np.ones(n) for t, n in zip(taus, n_neurons)]))

    # simulate host
    nls = [nengo.LIF(tau_ref=ref, tau_rc=taus[i])
           for i, n in enumerate(n_neurons)]
    for i, nl in enumerate(nls):
        if upsample <= 1:
            nl.step_math(dt, J[i], OS[i], V[i], W[i])
        else:
            s = np.zeros_like(OS[i])
            for j in range(upsample):
                nl.step_math(dt / upsample, J[i], s, V[i], W[i])
                OS[i] = (1./dt) * ((OS[i] > 0) | (s > 0))

    # simulate device
    plan = plan_lif(
        queue, dt, clJ, clV, clW, clOS, ref, clTaus, upsample=upsample)
    plan()

    if 1:
        a, b = V, clV
        for i in range(len(a)):
            nc, _ = not_close(a[i], b[i]).nonzero()
            if len(nc) > 0:
                j = nc[0]
                print("i", i, "j", j)
                print("J", J[i][j], clJ[i][j])
                print("V", V[i][j], clV[i][j])
                print("W", W[i][j], clW[i][j])
                print("...", len(nc) - 1, "more")

    n_spikes = np.sum([np.sum(os) for os in OS])
    if n_spikes < 1.0:
        logger.warn("LIF spiking mechanism was not tested!")
    assert ra.allclose(J, clJ.to_host())
    assert ra.allclose(V, clV.to_host())
    assert ra.allclose(W, clW.to_host())
    assert ra.allclose(OS, clOS.to_host())

def test_slicedcopy(rng):
    sizes = rng.randint(20, 200, size=10)
    A = RA([rng.normal(size=size) for size in sizes])
    B = RA([rng.normal(size=size) for size in sizes])
    incs = RA([rng.randint(0, 2) for _ in sizes])

    Ainds = []
    Binds = []
    for size in sizes:
        r = np.arange(size, dtype=np.int32)
        u = rng.choice([0, 1, 2])
        if u == 0:
            Ainds.append(r)
            Binds.append(r)
        elif u == 1:
            Ainds.append(r[:10])
            Binds.append(r[-10:])
        elif u == 2:
            n = rng.randint(2, size - 2)
            Ainds.append(rng.permutation(size)[:n])
            Binds.append(rng.permutation(size)[:n])

    Ainds = RA(Ainds)
    Binds = RA(Binds)

    queue = cl.CommandQueue(ctx)
    clA = CLRA(queue, A)
    clB = CLRA(queue, B)
    clAinds = CLRA(queue, Ainds)
    clBinds = CLRA(queue, Binds)
    clincs = CLRA(queue, incs)

    # compute on host
    for i in range(len(sizes)):
        if incs[i]:
            B[i][Binds[i]] += A[i][Ainds[i]]
        else:
            B[i][Binds[i]] = A[i][Ainds[i]]

    # compute on device
    plan = plan_slicedcopy(queue, clA, clB, clAinds, clBinds, clincs)
    plan()

    # check result
    for y, yy in zip(B, clB.to_host()):
        assert np.allclose(y, yy)

def test_elementwise_inc(rng):
    Xsizes = [(3, 3), (32, 64), (457, 342), (1, 100)]
    Asizes = [(3, 3), (1, 1),   (457, 342), (100, 1)]
    A = RA([rng.normal(size=size) for size in Asizes])
    X = RA([rng.normal(size=size) for size in Xsizes])
    Y = RA([a * x for a, x in zip(A, X)])

    queue = cl.CommandQueue(ctx)
    clA = CLRA(queue, A)
    clX = CLRA(queue, X)
    clY = CLRA(queue, RA([np.zeros_like(y) for y in Y]))

    # compute on device
    plan = plan_elementwise_inc(queue, clA, clX, clY)
    plan()

    # check result
    for y, yy in zip(Y, clY.to_host()):
        assert np.allclose(y, yy)

def test_speed(rng):
    try:
        import pyopencl_blas
    except ImportError:
        pyopencl_blas = None

    # enable_out_of_order = (
    #     cl.command_queue_properties.OUT_OF_ORDER_EXEC_MODE_ENABLE)

    k = 300
    # k = 100
    # k = 32
    # k = 16
    ms = [rng.randint(100, 1000) for i in range(k)]
    ns = [rng.randint(100, 1000) for i in range(k)]
    # ms = [4096 for i in range(k)]
    # ns = [4096 for i in range(k)]

    aa = [rng.uniform(-1, 1, size=(m, n)).astype('float32')
          for m, n in zip(ms, ns)]
    xx = [rng.uniform(-1, 1, size=n).astype('float32') for n in ns]
    yy = [rng.uniform(-1, 1, size=m).astype('float32') for m in ms]
    ajs = [np.int32(i) for i in range(k)]
    xjs = [np.int32(i) for i in range(k)]
    # ajs = [rng.randint(k, size=p) for i in range(k)]
    # xjs = [rng.randint(k, size=p) for i in range(k)]

    # alpha = 0.5
    # beta = 0.1
    alpha = 1.0
    beta = 1.0

    # -- prepare initial conditions on device
    queue = cl.CommandQueue(ctx)
    # queue = cl.CommandQueue(ctx, properties=enable_out_of_order)
    clA = CLRA.from_arrays(queue, aa)
    clX = CLRA.from_arrays(queue, xx)
    clY = CLRA.from_arrays(queue, yy)
    A_js = RA(ajs, dtype=np.int32)
    X_js = RA(xjs, dtype=np.int32)

    # -- run cl computation
    prog = plan_ragged_gather_gemv(
        queue, alpha, clA, A_js, clX, X_js, beta, clY)
    plans = prog.choose_plans()

    print('')
    print('-' * 5 + ' Plans ' + '-' * 45)
    for plan in plans:
        print(plan)

    with Timer() as timer:
        for plan in plans:
            plan()
    print("nengo_ocl: %0.3f" % timer.duration)

    # -- speed test in ocl blas
    if pyopencl_blas:
        pyopencl_blas.setup()

        def array(a):
            cla = cl.array.Array(queue, a.shape, a.dtype)
            cla.set(a)
            return cla

        clAs = [array(a) for a in aa]
        clXs = [array(x.ravel()) for x in xx]
        clYs = [array(y.ravel()) for y in yy]

        queues = [cl.CommandQueue(ctx) for _ in range(k)]
        # queues = [cl.CommandQueue(ctx, properties=enable_out_of_order)
        #           for _ in range(k)]

        queue.finish()
        with Timer() as timer:
            if 0:
                # use a single queue
                for A, X, Y in zip(clAs, clXs, clYs):
                    pyopencl_blas.gemv(queue, A, X, Y)
                queue.finish()
            else:
                # use multiple parallel queues
                events = []
                for i, [A, X, Y] in enumerate(zip(clAs, clXs, clYs)):
                    q = queues[i % len(queues)]
                    e = pyopencl_blas.gemv(q, A, X, Y)
                    events.append(e)
                for q in queues:
                    q.flush()
                cl.wait_for_events(events)
        print("clBLAS: %0.3f" % timer.duration)