diff --git a/brian2/codegen/generators/cppyy_generator.py b/brian2/codegen/generators/cppyy_generator.py new file mode 100644 index 000000000..ccb79de8b --- /dev/null +++ b/brian2/codegen/generators/cppyy_generator.py @@ -0,0 +1,191 @@ +""" +cppyy Code Generator +==================== +This generator converts Brian2's abstract code into C++ code that can be +JIT-compiled by cppyy. It inherits from CPPCodeGenerator to reuse most +C++ generation logic. + +WHY: We need this to translate Brian's abstract syntax tree (AST) into C++ code. +""" + +import typing + +from brian2.codegen.generators.cpp_generator import CPPCodeGenerator +from brian2.core.preferences import BrianPreference, prefs + +# Register cppyy-specific preferences +prefs.register_preferences( + "codegen.generators.cppyy", + "cppyy codegen preferences", + restrict_keyword=BrianPreference( + default="", # No restrict keyword for cppyy + docs=""" + The restrict keyword for cppyy. Empty by default as cppyy + doesn't always handle __restrict well. + """, + ), + flush_denormals=BrianPreference( + default=False, + docs=""" + Whether to add denormal flushing code. Disabled for cppyy + as it's handled at runtime. + """, + ), +) + + +class CppyyCodeGenerator(CPPCodeGenerator): + """ + cppyy code generator - generates C++ code for JIT compilation + + This class handles the conversion of Brian2's abstract code representation + into C++ code that cppyy can compile and execute. + """ + + class_name = "cppyy" + + def __init__(self, *args, **kwds): + """ + Initialize the cppyy code generator. + + WHY: We call the parent class (CPPCodeGenerator) because we want to + reuse majority of its functionality. We're not reinventing C++ generation, + just adapting it for cppyy. + """ + super().__init__(*args, **kwds) + + # # Track what headers we need for cppyy + # # WHY: cppyy needs to know which C++ headers to include + # self.headers_to_include = [ + # "", # For mathematical functions + # "", # For min, max, etc. + # "", # For int32_t, int64_t types + # ] + + @property + def restrict(self): + """Override to use cppyy-specific preference""" + # Return empty string or use cppyy preference + return f"{prefs['codegen.generators.cppyy.restrict_keyword']}" + + @property + def flush_denormals(self): + """Override to use cppyy-specific preference.""" + return prefs["codegen.generators.cppyy.flush_denormals"] + + def translate(self, abstract_code, dtype): + """Override to flatten the generated code structure for cppyy templates""" + + # Get the standard CPP generator result + scalar_code, vector_code, kwds = super().translate(abstract_code, dtype) + + print("\n=== DEBUGGING CODE TRANSLATION (Before Flattening) ===") + print(f"Raw scalar_code type: {type(scalar_code)}") + print(f"Raw scalar_code: {scalar_code}") + print(f"Raw vector_code type: {type(vector_code)}") + print(f"Raw vector_code: {vector_code}") + + # Flatten the code structures into simple strings + flattened_scalar = self._flatten_code_block(scalar_code) + flattened_vector = self._flatten_code_block(vector_code) + + print(f"Flattened scalar_code: '{flattened_scalar}'") + print(f"Flattened vector_code: '{flattened_vector}'") + print("=" * 60) + + return flattened_scalar, flattened_vector, kwds + + def _flatten_code_block(self, code_block): + """ + Convert Brian2's multi-block code structure into a simple string. + + This handles the conversion from: + {None: ['line1', 'line2', 'line3']} + To: + "line1\nline2\nline3" + """ + + if isinstance(code_block, str): + # Already a simple string, return as-is + return code_block + + if isinstance(code_block, dict): + # This is the multi-block structure we need to flatten + all_lines = [] + + # Process each block (usually just None for simple cases) + for _, line_list in code_block.items(): + if isinstance(line_list, list): + # Join all lines in this block + for line in line_list: + if line.strip(): # Skip empty lines + all_lines.append(line) + elif isinstance(line_list, str): + # Sometimes it's already a string + if line_list.strip(): + all_lines.append(line_list) + + # Join all lines with newlines to create proper C++ code + return "\n".join(all_lines) + + if isinstance(code_block, list): + # Sometimes it's just a list directly + return "\n".join(line for line in code_block if line.strip()) + + # Fallback: convert to string + return str(code_block) + + def translate_expression(self, expr): + """ + Translate a Brian2 expression to C++ code. + + Example: + Input: "v + 1" + Output: "_ptr_array_neurongroup_v[_idx] + 1" + + WHY: Brian2 uses simple variable names (v), but in C++ we need to + access them as array elements with proper indexing. + """ + # Use parent class method - it already handles this well + return super().translate_expression(expr) + + def determine_keywords(self) -> dict[str, typing.Any]: + """ + Determine which C++ keywords are used in the generated code. + + WHY: This helps optimize the generated code by only including + necessary type definitions and functions. + """ + # Get all standard CPP keywords + keywords: dict[str, typing.Any] = ( + super().determine_keywords() + ) # satisfy Pyright + + keywords.update( + { + "is_cppyy_target": True, + "is_standalone": False, + "cppyy_function_name": f"brian_kernel_{self.name}", + # These help templates know they're in runtime mode + "runtime_mode": True, + "needs_main_function": False, + } + ) + + # Modify support code for cppyy + # We don't need file I/O or main function setup + keywords["support_code_lines"] = self._adapt_support_code( + keywords.get("support_code_lines", "") + ) + print(f"Final keywords: {list(keywords.keys())}") + print("=" * 50) + return keywords + + def _adapt_support_code(self, support_code): + """ + Adapt support code for cppyy runtime. + Remove file I/O, adapt for JIT compilation. + """ + # TODO: For cppyy, we compile support code separately, so we might + # want to split it into header-like and implementation parts + return support_code diff --git a/brian2/codegen/runtime/cppyy_rt/__init__.py b/brian2/codegen/runtime/cppyy_rt/__init__.py new file mode 100644 index 000000000..ae779ccda --- /dev/null +++ b/brian2/codegen/runtime/cppyy_rt/__init__.py @@ -0,0 +1,9 @@ +""" +cppyy Runtime Module +==================== +This module provides JIT C++ compilation for Brian2 using cppyy. +""" + +from .cppyy_rt import CppyyCodeObject + +__all__ = ["CppyyCodeObject"] diff --git a/brian2/codegen/runtime/cppyy_rt/cppyy_rt.py b/brian2/codegen/runtime/cppyy_rt/cppyy_rt.py new file mode 100644 index 000000000..7d928db9d --- /dev/null +++ b/brian2/codegen/runtime/cppyy_rt/cppyy_rt.py @@ -0,0 +1,692 @@ +""" +cppyy Runtime Code Object +========================= +This is the core of our implementation. It handles: +1. Taking generated C++ code +2. Compiling it with cppyy +3. Executing it with numpy arrays +""" + +import cppyy +import cppyy.ll +import numpy as np + +from brian2.codegen.codeobject import ( + CodeObject, + check_compiler_kwds, + constant_or_scalar, +) +from brian2.codegen.generators.cpp_generator import c_data_type +from brian2.codegen.generators.cppyy_generator import CppyyCodeGenerator +from brian2.codegen.targets import codegen_targets +from brian2.codegen.templates import Templater +from brian2.core.base import BrianObjectException, weakproxy_with_fallback +from brian2.core.functions import DEFAULT_FUNCTIONS, Function +from brian2.core.preferences import BrianPreference, prefs +from brian2.core.variables import ( + ArrayVariable, + AuxiliaryVariable, + DynamicArrayVariable, + Subexpression, +) +from brian2.utils.logger import get_logger +from brian2.utils.stringtools import get_identifiers + +logger = get_logger(__name__) + +# Configure cppyy for better performance +# WHY: These settings optimize cppyy for numerical computing +cppyy.add_include_path(np.get_include()) # Include numpy headers +cppyy.ll.set_signals_as_exception( + True +) # to log failures , https://cppyy.readthedocs.io/en/latest/debugging.html + +# Register preferences +prefs.register_preferences( + "codegen.runtime.cppyy", + "cppyy runtime codegen preferences", + cache_compiled_functions=BrianPreference( + default=True, + docs=""" + Whether to cache JIT compiled functions. This avoids recompilation + but uses more memory. + """, + ), + debug_mode=BrianPreference( + default=False, + docs=""" + Whether to print generated C++ code before compilation for debugging. + """, + ), +) + +# Global flag to track if support code has been compiled +_support_code_compiled = False + + +def compile_support_code(): + """Compile common support code once for all cppyy code objects.""" + global _support_code_compiled + if _support_code_compiled: + return + + support_code = """ + #include + #include + #include + #include + + // Brian2 type definitions + typedef int32_t int32; + typedef int64_t int64; + typedef float float32; + typedef double float64; + + // Random number generation + namespace brian { + thread_local std::mt19937 _random_generator; + + inline double _rand(int idx) { + std::uniform_real_distribution dist(0.0, 1.0); + return dist(_random_generator); + } + + inline double _randn(int idx) { + std::normal_distribution dist(0.0, 1.0); + return dist(_random_generator); + } + + template + inline T _clip(T x, T low, T high) { + return std::min(std::max(x, low), high); + } + + inline int _int(double x) { return (int)x; } + inline int _int(bool x) { return x ? 1 : 0; } + + template + inline T _brian_mod(T x, T y) { + return x - y * floor(x/y); + } + } + + using namespace brian; + """ + + try: + cppyy.cppdef(support_code) + _support_code_compiled = True + logger.debug("Compiled cppyy support code") + except Exception as e: + logger.error(f"Failed to compile support code: {e}") + raise + + +class CppyyCodeObject(CodeObject): + """ + Execute Brian2 code using cppyy JIT compilation. + + This class is responsible for: + 1. Taking the C++ code generated by CppyyCodeGenerator + 2. Compiling it using cppyy's JIT compiler + 3. Executing the compiled function with Brian2's data + """ + + # Tell Brian2 which templates and generator to use + templater = Templater( + "brian2.codegen.runtime.cppyy_rt", + ".cpp", # We use .cpp extension for our templates + env_globals={ + "dtype": np.dtype, + "c_data_type": c_data_type, + "constant_or_scalar": constant_or_scalar, + "prefs": prefs, + }, + ) + + generator_class = CppyyCodeGenerator + class_name = "cppyy" + + def __init__( + self, + owner, + code, + variables, + variable_indices, + template_name, + template_source, + compiler_kwds, + name="cppyy_code_object*", + ): + """ + Initialize a cppyy code object. + + Parameters: + ----------- + owner : Group + The NeuronGroup or Synapses object that owns this code + code : str + The generated C++ code + variables : dict + Dictionary of Variable objects used in the code + variable_indices : dict + Mapping of variables to their index arrays + template_name : str + Name of the template used (e.g., 'stateupdate') + """ + logger.debug(f"Creating cppyy code object: {name}") + # Check compiler keywords (we don't use many for cppyy) + check_compiler_kwds( + compiler_kwds, + ["include_dirs", "libraries"], # Minimal set for cppyy + "cppyy", + ) + + print("\n=== DEBUGGING TEMPLATE VARIABLE PASSING ===") + print(f"Template name: {template_name}") + print(f"Template source type: {type(template_source)}") + print(f"Code parameter type: {type(code)}") + print(f"Code parameter content: '{str(code)[:500]}...'") + + # This is crucial - let's see what the template actually received + if hasattr(template_source, "globals"): + print(f"Template globals: {template_source .globals.keys()}") + + # Let's also check if Brian2 is calling our generator correctly + print(f"Variables passed to code object: {list(variables.keys())}") + print("=" * 60) + + super().__init__( + owner, + code, + variables, + variable_indices, + template_name, + template_source, + compiler_kwds={}, # Don't pass compiler args to numpy + name=name, + ) + + # Ensure support code is compiled + compile_support_code() + + # Store compiled functions for each code block + self.compiled_funcs = {} + + # Lists for tracking non-constant values (like CythonCodeObject) + self.nonconstant_values = [] + self.namespace = {"_owner": weakproxy_with_fallback(owner)} + + # Add default functions and constants + self.namespace.update(DEFAULT_FUNCTIONS) + + from brian2.devices.device import get_device + + self.device = get_device() + # Process variables into namespace + self.variables_to_namespace() + + @classmethod + def is_available(cls): + """Check if cppyy is available and working.""" + try: + import cppyy + + cppyy.cppdef("void test_func() { }") + cppyy.gbl.test_func() + return True + except ImportError: + return False + + # def compile_block(self, block): + # """Compile a specific code block - CORRECTED VERSION""" + + # print(f"\n=== DEBUGGING BLOCK ACCESS ===") + # print(f"Looking for block: {block}") + # print(f"Code object type: {type(self.code)}") + # print(f"Code object length: {len(str(self.code))}") + + # # For runtime code objects, the code is usually a single string + # # Let's first check if it's a simple string containing all the code + # if isinstance(self.code, str): + # print("Code is a single string - this is typical for runtime targets") + # code = self.code.strip() + # print(f"Full code content (first 500 chars):\n{code[:500]}...") + + # # If it's an object with attributes, try the standard approach + # elif hasattr(self.code, block): + # code = getattr(self.code, block, "").strip() + # print(f"Found {block} attribute with content: {code[:200]}...") + + # # Try the cpp_file variant (like standalone mode as our generator inherits from cppone) + # elif hasattr(self.code, f"{block}_cpp_file"): + # code = getattr(self.code, f"{block}_cpp_file", "").strip() + # print(f"Found {block}_cpp_file attribute with content: {code[:200]}...") + + # else: + # print("Could not find code content using any known method") + # print(f"Available attributes: {dir(self.code)}") + # return None + + # if not code or "EMPTY_CODE_BLOCK" in code: + # print(f"Block {block} is empty or marked as empty") + # return None + # # Clean up the template code before building the function + # if "['template" in code: + # # This is a stringified list - extract and join it + # import ast + # try: + # # Find the list representation in the code + # start = code.index("['") + # end = code.index("]", start) + 1 + # list_str = code[start:end] + # actual_list = ast.literal_eval(list_str) + # joined_code = '\n'.join(actual_list) + # code = code[:start] + joined_code + code[end:] + # except: + # pass + + # # Extract the relevant part of the code for this specific block + # # Since we have the full rendered template, we need to extract the part we want + # if block == "run": + # # For the main run block, we want the vector code section + # # Look for the main computation loop + # if "for(int _idx" in code and "vector_code" not in code.lower(): + # # This looks like rendered template code, use it as-is + # extracted_code = code + # else: + # # If we can't identify the right section, use the whole thing + # extracted_code = code + + # elif block in ["before_run", "after_run"]: + # # These blocks might be empty for simple neuron groups + # if len(code.strip()) < 50: # Very short code, likely empty + # print(f"Block {block} appears to be empty") + # return None + # extracted_code = code + + # else: + # extracted_code = code + + # print(f"Extracted code for {block} (length: {len(extracted_code)}):") + # print(f"Preview: {extracted_code[:300]}...") + # print("=" * 60) + + # # Generate unique function name + # func_name = f"brian_{self.name.replace('*', '').replace('-', '_')}_{block}" + + # # Build complete C++ function + # cpp_code = self._build_block_function(func_name, extracted_code) + + # if prefs["codegen.runtime.cppyy.debug_mode"]: + # print(f"Complete C++ function to compile:") + # print("-" * 60) + # print(cpp_code) + # print("-" * 60) + + # try: + # # Compile with cppyy + # cppyy.cppdef(cpp_code) + # compiled_func = getattr(cppyy.gbl, func_name) + + # print(f"Successfully compiled {func_name}") + # return compiled_func + + # except Exception as e: + # print(f"Compilation failed for {block}: {e}") + # if prefs["codegen.runtime.cppyy.debug_mode"]: + # print(f"Failed code was:\n{cpp_code}") + # raise + + # def _build_block_function(self, func_name, template_code): + # """ + # Build a complete C++ function from template code. + + # WHY: Templates generate code fragments, we need to wrap them + # in a proper C++ function that cppyy can compile. + # """ + # # Extract array parameters from variables + # params = [] + # param_setup = [] + + # for varname, var in self.variables.items(): + # if isinstance(var, ArrayVariable): + # # Get the exact array name that Brian2's code generator uses + # brian2_array_name = self.device.get_array_name(var) + + # # Create function parameter using simplified name + # simplified_param_name = f"{brian2_array_name}_void" + # params.append(f"void* {simplified_param_name}") + + # # Create the pointer variable with Brian2's expected name + # dtype = c_data_type(var.dtype) + # param_setup.append(f"{dtype}* {brian2_array_name} = ({dtype}*){simplified_param_name};") + + # # Add size parameter for dynamic arrays + # if isinstance(var, DynamicArrayVariable): + # params.append(f"int _num_{varname}") + + # # Add standard parameters + # params.extend(["double t", "double dt", "int N"]) + + # # Build the function + # function_code = f""" + # extern "C" void {func_name}({', '.join(params)}) {{ + # // Cast void pointers to proper types + # {''.join(param_setup)} + + # // Original template code + # {template_code} + # }} + # """ + + # return function_code + # + + def compile_block(self, block): + """focus on just getting it working""" + # Only handle the main 'run' block for the now + if block != "run": + return None + + code = str(self.code).strip() + if not code: + return None + + print(f"Template generated code:\n{code}") + + # Generate unique function name with timestamp to avoid conflicts + import time + + timestamp = str(int(time.time() * 1000000)) # microsecond precision + base_name = self.name.replace("*", "").replace("-", "_") + func_name = f"brian_poc_{base_name}_{timestamp}" + + # Check if function already exists in cppyy + try: + existing_func = getattr(cppyy.gbl, func_name, None) + if existing_func is not None: + print(f"Function {func_name} already exists, reusing it") + return existing_func + except AttributeError: + pass # Function doesn't exist, which is what we want + + # Create the function with EXACT variable names Brian2 expects + cpp_code = self._create_poc_function(func_name, code) + + if prefs["codegen.runtime.cppyy.debug_mode"]: + print("POC C++ function to compile:") + print("-" * 60) + print(cpp_code) + print("-" * 60) + + try: + cppyy.cppdef(cpp_code) + compiled_func = getattr(cppyy.gbl, func_name) + print(f"Successfully compiled {func_name}") + return compiled_func + except Exception as e: + print(f"Compilation failed: {e}") + raise + + def _create_poc_function(self, func_name, template_code): + """Create a simple working function for POC""" + + # for now as we know the exact variables needed + # for testing we create them with the exact names Brian2 uses + + # Remove the problematic block markers that create scoping issues + processed_code = template_code.replace("// Scalar code\n{", "// Scalar code") + processed_code = processed_code.replace("// Vector code\n{", "// Vector code") + + # Remove closing braces that aren't part of for loops + lines = processed_code.split("\n") + cleaned_lines = [] + + for line in lines: + stripped = line.strip() + # Keep braces that are part of for loops or have content after them + if stripped == "}" or "{" and len(cleaned_lines) > 0: + # Check if this is likely a block-ending brace we want to remove + prev_line = cleaned_lines[-1].strip() if cleaned_lines else "" + if not (prev_line.endswith(";") or prev_line.endswith("}")): + continue # Skip this closing brace + cleaned_lines.append(line) + + processed_code = "\n".join(cleaned_lines) + + function_code = f""" + extern "C" void {func_name}(void* dt_ptr, void* tau_ptr, void* v_ptr, + double t_val, double dt_val, int N_val) {{ + // Create variables with EXACT names that Brian2 expects + double* _ptr_array_defaultclock_dt = (double*)dt_ptr; + double* _array_defaultclock_dt = (double*)dt_ptr; + double* _ptr_array_neurongroup_tau = (double*)tau_ptr; + double* _ptr_array_neurongroup_v = (double*)v_ptr; + + // Set up scalar variables + double t = t_val; + int N = N_val; + + // Execute flattened template code (no problematic scoping) + {processed_code} + }} + """ + + return function_code + + # def run_block(self, block): + # """ + # Run a compiled code block. + + # This is called by Brian2's execution system. + # """ + # compiled_func = self.compiled_code.get(block) + # print('test',getattr(self.code, block, "").strip()) + # if compiled_func is None: + # return # Nothing to run + + # try: + # # Prepare arguments for the C++ function + # args = [] + + # # Add array pointers + # for _, var in self.variables.items(): + # if isinstance(var, ArrayVariable): + # # Get the numpy array + # value = var.get_value() + # if isinstance(value, np.ndarray): + # # Pass the data pointer + # args.append(value.ctypes.data) + + # # Add size for dynamic arrays + # if isinstance(var, DynamicArrayVariable): + # args.append(len(value)) + + # # Add scalar values + # args.append(self.namespace.get("t", 0.0)) + # args.append(self.namespace.get("dt", 0.0001)) + # args.append(self.namespace.get("N", 0)) + + # # Call the compiled function + # compiled_func(*args) + + # except Exception as exc: + # message = ( + # f"An exception occurred during execution of the " + # f"'{block}' block of code object '{self.name}'.\n" + # ) + # raise BrianObjectException(message, self.owner) from exc + # + def run_block(self, block): + + if block not in self.compiled_funcs: + self.compiled_funcs[block] = self.compile_block(block) + + compiled_func = self.compiled_funcs[block] + if compiled_func is None: + return + + try: + # Get the numpy arrays for your specific variables + dt_array = None + tau_array = None + v_array = None + N_val = 0 + + # Find the arrays we need + for _, var in self.variables.items(): + if isinstance(var, ArrayVariable): + value = var.get_value() + if isinstance(value, np.ndarray): + array_name = self.device.get_array_name(var) + # print(f"Found array {name} -> {array_name}: shape={value.shape}") + + if "defaultclock" in array_name and "dt" in array_name: + dt_array = value + elif "tau" in array_name: + tau_array = value + N_val = len(value) + elif "v" in array_name: + v_array = value + + if dt_array is None or tau_array is None or v_array is None: + return + + # CRITICAL FIX: Convert numpy arrays to cppyy-compatible pointers + import ctypes + + # Ensure arrays are contiguous and have the right data type + dt_array = np.ascontiguousarray(dt_array, dtype=np.float64) + tau_array = np.ascontiguousarray(tau_array, dtype=np.float64) + v_array = np.ascontiguousarray(v_array, dtype=np.float64) + + # Create ctypes pointers that cppyy can understand + dt_ptr = dt_array.ctypes.data_as(ctypes.POINTER(ctypes.c_double)) + tau_ptr = tau_array.ctypes.data_as(ctypes.POINTER(ctypes.c_double)) + v_ptr = v_array.ctypes.data_as(ctypes.POINTER(ctypes.c_double)) + + # Convert ctypes pointers to cppyy void* pointers + dt_void_ptr = ctypes.cast(dt_ptr, ctypes.c_void_p) + tau_void_ptr = ctypes.cast(tau_ptr, ctypes.c_void_p) + v_void_ptr = ctypes.cast(v_ptr, ctypes.c_void_p) + + print("Calling compiled function with properly converted pointers...") + print(f" dt_array: {dt_array} -> pointer: {dt_void_ptr}") + print(f" tau_array shape: {tau_array.shape}") + print(f" v_array shape: {v_array.shape}") + print(f" N_val: {N_val}") + + # Store original values for comparison + v_original = v_array.copy() + + # Call the function with properly converted pointers + compiled_func( + dt_void_ptr, # dt pointer + tau_void_ptr, # tau pointer + v_void_ptr, # v pointer + 0.0, # t value + float(dt_array[0]) if len(dt_array) > 0 else 0.0001, # dt value + N_val, # N value + ) + + print("Execution completed successfully!") + print(f"Original v values (first 5): {v_original[:5]}") + print(f"Updated v values (first 5): {v_array[:5]}") + print(f"Values changed: {not np.array_equal(v_original, v_array)}") + + except Exception as e: + print(f"✗ Execution failed: {e}") + print(f"Error type: {type(e)}") + import traceback + + traceback.print_exc() + raise BrianObjectException( + f"Execution failed in {block} block", self.owner + ) from e + + def _insert_func_namespace(self, func): + """Insert function namespace (copied from CythonCodeObject).""" + impl = func.implementations[self] + func_namespace = impl.get_namespace(self.owner) + if func_namespace is not None: + self.namespace.update(func_namespace) + if impl.dependencies is not None: + for dep in impl.dependencies.values(): + self._insert_func_namespace(dep) + + def variables_to_namespace(self): + """ + Convert variables to namespace (adapted from CythonCodeObject). + + WHY: We need to make variables accessible to the compiled code + and track which ones change during execution. + """ + # Lists for tracking non-constant values + self.nonconstant_values = [] + + for name, var in self.variables.items(): + if isinstance(var, Function): + self._insert_func_namespace(var) + if isinstance(var, (AuxiliaryVariable, Subexpression)): + continue + + try: + value = var.get_value() + except (TypeError, AttributeError): + # A dummy Variable without value or a function + self.namespace[name] = var + continue + + if isinstance(var, ArrayVariable): + # Store array reference + array_name = self.device.get_array_name(var, self.variables) + self.namespace[array_name] = value + self.namespace[f"_num_{name}"] = var.get_len() + + if var.scalar and var.constant: + self.namespace[name] = value.item() + else: + self.namespace[name] = value + + # Handle dynamic arrays + if isinstance(var, DynamicArrayVariable): + dyn_array_name = self.generator_class.get_array_name( + var, access_data=False + ) + self.namespace[dyn_array_name] = self.device.get_value( + var, access_data=False + ) + + # Store Variable object itself + self.namespace[f"_var_{name}"] = var + + # Filter namespace to only include used identifiers + all_identifiers = set() + for block in ["before_run", "run", "after_run"]: + if hasattr(self.code, block): + code = getattr(self.code, block) + all_identifiers.update(get_identifiers(code)) + + self.namespace = { + k: v + for k, v in self.namespace.items() + if k in all_identifiers or k.startswith("_") + } + + # Track dynamic arrays that need updates + for name, var in self.variables.items(): + if isinstance(var, DynamicArrayVariable) and var.needs_reference_update: + array_name = self.device.get_array_name(var, self.variables) + if array_name in self.namespace: + self.nonconstant_values.append((array_name, var.get_value)) + if f"_num_{name}" in self.namespace: + self.nonconstant_values.append((f"_num_{name}", var.get_len)) + + def update_namespace(self): + """Update the values of non-constant values in the namespace.""" + for name, func in self.nonconstant_values: + self.namespace[name] = func() + + +codegen_targets.add(CppyyCodeObject) diff --git a/brian2/codegen/runtime/cppyy_rt/templates/common.cpp b/brian2/codegen/runtime/cppyy_rt/templates/common.cpp new file mode 100644 index 000000000..116763379 --- /dev/null +++ b/brian2/codegen/runtime/cppyy_rt/templates/common.cpp @@ -0,0 +1,29 @@ +{# Base template for groups - handles most common cases #} + + +// Scalar code +{% if scalar_code %} +{ + const int _vectorisation_idx = -1; + {{scalar_code|autoindent}} +} +{% endif %} + +// Vector code +{% if vector_code %} +{ + // Get N if available + {% if 'N' in variables %} + const int _N = {{constant_or_scalar('N', variables['N'])}}; + {% else %} + const int _N = 1; // Default + {% endif %} + + // Main loop + for(int _idx=0; _idx<_N; _idx++) + { + const int _vectorisation_idx = _idx; + {{vector_code|autoindent}} + } +} +{% endif %} diff --git a/brian2/codegen/runtime/cppyy_rt/templates/reset.cpp b/brian2/codegen/runtime/cppyy_rt/templates/reset.cpp new file mode 100644 index 000000000..f5c4b3f3a --- /dev/null +++ b/brian2/codegen/runtime/cppyy_rt/templates/reset.cpp @@ -0,0 +1,26 @@ +{# USES_VARIABLES { N, _spikespace, t } #} +{# ALLOWS_SCALAR_WRITE #} + + +// Scalar code +{% if scalar_code %} +{ + const int _vectorisation_idx = -1; + {{scalar_code|autoindent}} +} +{% endif %} + +// Vector code - reset spiking neurons +{% if vector_code %} +{ + // Reset code runs only for neurons that spiked + // _spikespace contains indices of neurons that spiked + const int _N = {{constant_or_scalar('N', variables['N'])}}; + + for(int _idx=0; _idx<_N; _idx++) + { + const int _vectorisation_idx = _idx; + {{vector_code|autoindent}} + } +} +{% endif %} diff --git a/brian2/codegen/runtime/cppyy_rt/templates/stateupdate.cpp b/brian2/codegen/runtime/cppyy_rt/templates/stateupdate.cpp new file mode 100644 index 000000000..c142ea8e3 --- /dev/null +++ b/brian2/codegen/runtime/cppyy_rt/templates/stateupdate.cpp @@ -0,0 +1,24 @@ +{# USES_VARIABLES { N, dt } #} +{# ALLOWS_SCALAR_WRITE #} + +// Scalar code +{% if scalar_code %} +{ + const int _vectorisation_idx = -1; + {{scalar_code|autoindent}} +} +{% endif %} + +// Vector code +{% if vector_code %} +{ + const int _N = {{constant_or_scalar('N', variables['N'])}}; + const double _dt = {{constant_or_scalar('dt', variables['dt'])}}; + + for(int _idx=0; _idx<_N; _idx++) + { + const int _vectorisation_idx = _idx; + {{vector_code|autoindent}} + } +} +{% endif %} diff --git a/brian2/codegen/runtime/cppyy_rt/templates/threshold.cpp b/brian2/codegen/runtime/cppyy_rt/templates/threshold.cpp new file mode 100644 index 000000000..012ac0e86 --- /dev/null +++ b/brian2/codegen/runtime/cppyy_rt/templates/threshold.cpp @@ -0,0 +1,24 @@ +{# USES_VARIABLES { N, _spikespace, t } #} +{# ALLOWS_SCALAR_WRITE #} + +// Scalar code +{% if scalar_code %} +{ + const int _vectorisation_idx = -1; + {{scalar_code|autoindent}} +} +{% endif %} + +// Vector code - check threshold condition +{% if vector_code %} +{ + const int _N = {{constant_or_scalar('N', variables['N'])}}; + + // Check threshold for each neuron + for(int _idx=0; _idx<_N; _idx++) + { + const int _vectorisation_idx = _idx; + {{vector_code|autoindent}} + } +} +{% endif %} diff --git a/brian2/devices/device.py b/brian2/devices/device.py index 5227928cc..c338fab45 100644 --- a/brian2/devices/device.py +++ b/brian2/devices/device.py @@ -8,6 +8,7 @@ import numpy as np +from brian2.codegen.runtime.cppyy_rt import CppyyCodeObject from brian2.codegen.targets import codegen_targets from brian2.core.names import find_name from brian2.core.preferences import prefs @@ -271,6 +272,7 @@ def code_object_class(self, codeobj_class=None, fallback_pref="codegen.target"): for target in codegen_targets: if target.class_name == codeobj_class: return target + print("codegen_targets", codegen_targets) # No target found targets = ["auto"] + [ target.class_name for target in codegen_targets if target.class_name @@ -785,3 +787,15 @@ def seed(seed=None): runtime_device = RuntimeDevice() all_devices["runtime"] = runtime_device + + +class CppyyDevice(RuntimeDevice): + """ + A Brian 2 device for running simulations using the cppyy JIT compiler. + """ + + code_object_class = CppyyCodeObject + + def __init__(self, *args, **kwds): + super().__init__(*args, **kwds) + self.name = "cppyy"