"""
Module implementing the C++ "standalone" `CodeObject`
"""
from brian2.codegen.codeobject import CodeObject, constant_or_scalar
from brian2.codegen.targets import codegen_targets
from brian2.codegen.templates import Templater
from brian2.codegen.generators.cpp_generator import (CPPCodeGenerator,
c_data_type)
from brian2.devices.device import get_device
from brian2.core.preferences import prefs
from brian2.core.functions import DEFAULT_FUNCTIONS
from brian2.utils.stringtools import replace
__all__ = ['CPPStandaloneCodeObject']
[docs]def openmp_pragma(pragma_type):
nb_threads = prefs.devices.cpp_standalone.openmp_threads
openmp_on = not (nb_threads == 0)
## First we need to deal with some special cases that have to be handle in case
## openmp is not activated
if pragma_type == 'set_num_threads':
if not openmp_on:
return ''
elif nb_threads > 0:
# We have to fix the exact number of threads in all parallel sections
return f'omp_set_dynamic(0);\nomp_set_num_threads({int(nb_threads)});'
elif pragma_type == 'get_thread_num':
if not openmp_on:
return '0'
else:
return 'omp_get_thread_num()'
elif pragma_type == 'get_num_threads':
if not openmp_on:
return '1'
else:
return f'{int(nb_threads)}'
elif pragma_type == 'with_openmp':
# The returned value is a proper Python boolean, i.e. not something
# that should be included in the generated code but rather for use
# in {% if ... %} statements in the template
return openmp_on
## Then by default, if openmp is off, we do not return any pragma statement in the templates
elif not openmp_on:
return ''
## Otherwise, we return the correct pragma statement
elif pragma_type == 'static':
return '#pragma omp for schedule(static)'
elif pragma_type == 'single':
return '#pragma omp single'
elif pragma_type == 'single-nowait':
return '#pragma omp single nowait'
elif pragma_type == 'critical':
return '#pragma omp critical'
elif pragma_type == 'atomic':
return '#pragma omp atomic'
elif pragma_type == 'once':
return '#pragma once'
elif pragma_type == 'parallel-static':
return '#pragma omp parallel for schedule(static)'
elif pragma_type == 'static-ordered':
return '#pragma omp for schedule(static) ordered'
elif pragma_type == 'ordered':
return '#pragma omp ordered'
elif pragma_type == 'include':
return '#include <omp.h>'
elif pragma_type == 'parallel':
return '#pragma omp parallel'
elif pragma_type == 'master':
return '#pragma omp master'
elif pragma_type == 'barrier':
return '#pragma omp barrier'
elif pragma_type == 'compilation':
return '-fopenmp'
elif pragma_type == 'sections':
return '#pragma omp sections'
elif pragma_type == 'section':
return '#pragma omp section'
else:
raise ValueError(f'Unknown OpenMP pragma "{pragma_type}"')
[docs]class CPPStandaloneCodeObject(CodeObject):
"""
C++ standalone code object
The ``code`` should be a `~brian2.codegen.templates.MultiTemplate`
object with two macros defined, ``main`` (for the main loop code) and
``support_code`` for any support code (e.g. function definitions).
"""
templater = Templater('brian2.devices.cpp_standalone', '.cpp',
env_globals={'c_data_type': c_data_type,
'openmp_pragma': openmp_pragma,
'constant_or_scalar': constant_or_scalar,
'prefs': prefs,
'zip': zip})
generator_class = CPPCodeGenerator
def __init__(self, *args, **kwds):
super(CPPStandaloneCodeObject, self).__init__(*args, **kwds)
#: Store whether this code object defines before/after blocks
self.before_after_blocks = []
[docs] def __call__(self, **kwds):
return self.run()
[docs] def compile_block(self, block):
pass # Compilation will be handled in device
[docs] def run_block(self, block):
if block == 'run':
get_device().main_queue.append((f"{block}_code_object", (self,)))
else:
# Check the C++ code whether there is anything to run
cpp_code = getattr(self.code, f"{block}_cpp_file")
if len(cpp_code) and 'EMPTY_CODE_BLOCK' not in cpp_code:
get_device().main_queue.append((f"{block}_code_object", (self,)))
self.before_after_blocks.append(block)
codegen_targets.add(CPPStandaloneCodeObject)
# At module initialization time, we do not yet know whether the code will be
# run with OpenMP or not. We therefore use a "dynamic implementation" which
# generates the rand/randn implementation during code generation.
[docs]def generate_rand_code(rand_func, owner):
nb_threads = prefs.devices.cpp_standalone.openmp_threads
if nb_threads == 0: # no OpenMP
thread_number = '0'
else:
thread_number = 'omp_get_thread_num()'
if rand_func == 'rand':
rk_call = 'rk_double'
elif rand_func == 'randn':
rk_call = 'rk_gauss'
else:
raise AssertionError(rand_func)
code = """
double _%RAND_FUNC%(const int _vectorisation_idx) {
return %RK_CALL%(brian::_mersenne_twister_states[%THREAD_NUMBER%]);
}
"""
code = replace(code, {'%THREAD_NUMBER%': thread_number,
'%RAND_FUNC%': rand_func,
'%RK_CALL%': rk_call})
return {'support_code': code}
rand_impls = DEFAULT_FUNCTIONS['rand'].implementations
rand_impls.add_dynamic_implementation(CPPStandaloneCodeObject,
code=lambda owner:
generate_rand_code('rand', owner),
namespace=lambda owner: {},
name='_rand')
randn_impls = DEFAULT_FUNCTIONS['randn'].implementations
randn_impls.add_dynamic_implementation(CPPStandaloneCodeObject,
code=lambda owner:
generate_rand_code('randn', owner),
namespace=lambda owner: {},
name='_randn')