"""
Module implementing the C++ "standalone" `CodeObject`
"""
from brian2.codegen.codeobject import CodeObject, constant_or_scalar
from brian2.codegen.generators.cpp_generator import CPPCodeGenerator, c_data_type
from brian2.codegen.targets import codegen_targets
from brian2.codegen.templates import Templater
from brian2.core.functions import DEFAULT_FUNCTIONS
from brian2.core.preferences import prefs
from brian2.devices.device import get_device
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().__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",
)