AgentPython capture external constants.

swig/python/codegen/__init__.py

@@ -37,9 +37,26 @@ def translate(function: Union[str, Callable]) -> str:
         # get source for each function
         for d_f in d_functions:
             prepend_source += inspect.getsource(d_f);
-        # get source for function and preprend device functions
+        # get source for function and prepend device functions
         function_source = prepend_source + inspect.getsource(function)
         tree = ast.parse(function_source)
-        return codegen(tree)
+        # Filter constants
+        module_members = inspect.getmembers(module);
+        # Emulate inspect.get_annotations() (requires python 3.10+)
+        module_annontations = {}
+        for mem in module_members:
+            if mem[0] == "__annotations__":
+                module_annontations = mem[1]
+                break
+        prepend_c_source = ""
+        # Find all annotated variables
+        for key, val in module_annontations.items():
+            if val.__name__ == "Final" or val.__name__ == "constant":
+                # Locate the literal for that variable (Python will precompute anything e.g. math.sqrt(12.5))
+                for mem in module_members:
+                    if key == mem[0]:
+                        prepend_c_source += f"constexpr auto {mem[0]} = {mem[1]};\n"
+                        break
+        return prepend_c_source + codegen(tree)
     else:
         raise CodeGenException(f"Error: translate function requires either a source string or Callable")

swig/python/codegen/codegen.py

@@ -625,6 +625,9 @@ def _Assign(self, t):
             if hasattr(t.value, "func") and isinstance(t.value.func, ast.Attribute) and t.value.func.attr == 'at' :
                 if t.value.func.value.id == self._input_message_var :
                     self._standalone_message_var.append(t.targets[0].id)
+            # Special case, definitions outside of agent fn are made const
+            if self._indent == 0:
+                self.write("constexpr ")
             self.write("auto ")
             self._locals.append(t.targets[0].id)
         self.dispatch(t.targets[0])

swig/python/flamegpu.i

@@ -1182,4 +1182,10 @@ TEMPLATE_VARIABLE_INSTANTIATE_INTS(poisson, flamegpu::HostRandom::poisson)
     #define GLM true
 #else
     #define GLM false
-#endif
+#endif
+
+// Declare an empty type we can use as an attribute for constants to be pulled in by codegen
+%pythoncode {
+class constant:
+  pass;
+}

tests/python/codegen/test_codegen.py

@@ -133,7 +133,7 @@ class Foo: pass
 a += 3
 """
 cpp_var_existing = """\
-auto a = 1;
+constexpr auto a = 1;
 a = 2;
 a += 3;
 """
@@ -170,76 +170,98 @@ async def f():
 """
 
 # FGPU functionality
+py_fgpu_constexpr = """\
+a = 12
 
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    b = 13
+"""
+
+cpp_fgpu_constexpr = """\
+constexpr auto a = 12;
+
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto b = 13;
+}
+"""
 
 
 
 py_fgpu_types = """\
-f = pyflamegpu.getVariableFloat("f")
-d = pyflamegpu.getVariableDouble("d")
-i = pyflamegpu.getVariableInt("i")
-ui = pyflamegpu.getVariableUInt("ui")
-i8 = pyflamegpu.getVariableInt8("i8")
-ui8 = pyflamegpu.getVariableUInt8("ui8")
-c = pyflamegpu.getVariableChar("c")
-uc = pyflamegpu.getVariableUChar("uc")
-i16 = pyflamegpu.getVariableInt16("i16")
-ui16 = pyflamegpu.getVariableUInt16("ui16")
-i32 = pyflamegpu.getVariableInt32("i32")
-ui32 = pyflamegpu.getVariableUInt32("ui32")
-i64 = pyflamegpu.getVariableInt64("i64")
-ui64 = pyflamegpu.getVariableUInt64("ui64")
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    f = pyflamegpu.getVariableFloat("f")
+    d = pyflamegpu.getVariableDouble("d")
+    i = pyflamegpu.getVariableInt("i")
+    ui = pyflamegpu.getVariableUInt("ui")
+    i8 = pyflamegpu.getVariableInt8("i8")
+    ui8 = pyflamegpu.getVariableUInt8("ui8")
+    c = pyflamegpu.getVariableChar("c")
+    uc = pyflamegpu.getVariableUChar("uc")
+    i16 = pyflamegpu.getVariableInt16("i16")
+    ui16 = pyflamegpu.getVariableUInt16("ui16")
+    i32 = pyflamegpu.getVariableInt32("i32")
+    ui32 = pyflamegpu.getVariableUInt32("ui32")
+    i64 = pyflamegpu.getVariableInt64("i64")
+    ui64 = pyflamegpu.getVariableUInt64("ui64")
 """
 
 cpp_fgpu_types = """\
-auto f = FLAMEGPU->getVariable<float>("f");
-auto d = FLAMEGPU->getVariable<double>("d");
-auto i = FLAMEGPU->getVariable<int>("i");
-auto ui = FLAMEGPU->getVariable<unsigned int>("ui");
-auto i8 = FLAMEGPU->getVariable<int_8>("i8");
-auto ui8 = FLAMEGPU->getVariable<uint_8>("ui8");
-auto c = FLAMEGPU->getVariable<char>("c");
-auto uc = FLAMEGPU->getVariable<unsigned char>("uc");
-auto i16 = FLAMEGPU->getVariable<int_16>("i16");
-auto ui16 = FLAMEGPU->getVariable<uint_16>("ui16");
-auto i32 = FLAMEGPU->getVariable<int_32>("i32");
-auto ui32 = FLAMEGPU->getVariable<uint_32>("ui32");
-auto i64 = FLAMEGPU->getVariable<int_64>("i64");
-auto ui64 = FLAMEGPU->getVariable<uint_64>("ui64");
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto f = FLAMEGPU->getVariable<float>("f");
+    auto d = FLAMEGPU->getVariable<double>("d");
+    auto i = FLAMEGPU->getVariable<int>("i");
+    auto ui = FLAMEGPU->getVariable<unsigned int>("ui");
+    auto i8 = FLAMEGPU->getVariable<int_8>("i8");
+    auto ui8 = FLAMEGPU->getVariable<uint_8>("ui8");
+    auto c = FLAMEGPU->getVariable<char>("c");
+    auto uc = FLAMEGPU->getVariable<unsigned char>("uc");
+    auto i16 = FLAMEGPU->getVariable<int_16>("i16");
+    auto ui16 = FLAMEGPU->getVariable<uint_16>("ui16");
+    auto i32 = FLAMEGPU->getVariable<int_32>("i32");
+    auto ui32 = FLAMEGPU->getVariable<uint_32>("ui32");
+    auto i64 = FLAMEGPU->getVariable<int_64>("i64");
+    auto ui64 = FLAMEGPU->getVariable<uint_64>("ui64");
+}
 """
 
 py_fgpu_array_types = """\
-f = pyflamegpu.getVariableFloatArray1("f")
-d = pyflamegpu.getVariableDoubleArray2("d")
-i = pyflamegpu.getVariableIntArray3("i")
-ui = pyflamegpu.getVariableUIntArray4("ui")
-i8 = pyflamegpu.getVariableInt8Array5("i8")
-ui8 = pyflamegpu.getVariableUInt8Array6("ui8")
-c = pyflamegpu.getVariableCharArray7("c")
-uc = pyflamegpu.getVariableUCharArray8("uc")
-i16 = pyflamegpu.getVariableInt16Array9("i16")
-ui16 = pyflamegpu.getVariableUInt16Array10("ui16")
-i32 = pyflamegpu.getVariableInt32Array11("i32")
-ui32 = pyflamegpu.getVariableUInt32Array12("ui32")
-i64 = pyflamegpu.getVariableInt64Array13("i64")
-ui64 = pyflamegpu.getVariableUInt64Array14("ui64")
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    f = pyflamegpu.getVariableFloatArray1("f")
+    d = pyflamegpu.getVariableDoubleArray2("d")
+    i = pyflamegpu.getVariableIntArray3("i")
+    ui = pyflamegpu.getVariableUIntArray4("ui")
+    i8 = pyflamegpu.getVariableInt8Array5("i8")
+    ui8 = pyflamegpu.getVariableUInt8Array6("ui8")
+    c = pyflamegpu.getVariableCharArray7("c")
+    uc = pyflamegpu.getVariableUCharArray8("uc")
+    i16 = pyflamegpu.getVariableInt16Array9("i16")
+    ui16 = pyflamegpu.getVariableUInt16Array10("ui16")
+    i32 = pyflamegpu.getVariableInt32Array11("i32")
+    ui32 = pyflamegpu.getVariableUInt32Array12("ui32")
+    i64 = pyflamegpu.getVariableInt64Array13("i64")
+    ui64 = pyflamegpu.getVariableUInt64Array14("ui64")
 """
 
 cpp_fgpu_array_types = """\
-auto f = FLAMEGPU->getVariable<float, 1>("f");
-auto d = FLAMEGPU->getVariable<double, 2>("d");
-auto i = FLAMEGPU->getVariable<int, 3>("i");
-auto ui = FLAMEGPU->getVariable<unsigned int, 4>("ui");
-auto i8 = FLAMEGPU->getVariable<int_8, 5>("i8");
-auto ui8 = FLAMEGPU->getVariable<uint_8, 6>("ui8");
-auto c = FLAMEGPU->getVariable<char, 7>("c");
-auto uc = FLAMEGPU->getVariable<unsigned char, 8>("uc");
-auto i16 = FLAMEGPU->getVariable<int_16, 9>("i16");
-auto ui16 = FLAMEGPU->getVariable<uint_16, 10>("ui16");
-auto i32 = FLAMEGPU->getVariable<int_32, 11>("i32");
-auto ui32 = FLAMEGPU->getVariable<uint_32, 12>("ui32");
-auto i64 = FLAMEGPU->getVariable<int_64, 13>("i64");
-auto ui64 = FLAMEGPU->getVariable<uint_64, 14>("ui64");
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto f = FLAMEGPU->getVariable<float, 1>("f");
+    auto d = FLAMEGPU->getVariable<double, 2>("d");
+    auto i = FLAMEGPU->getVariable<int, 3>("i");
+    auto ui = FLAMEGPU->getVariable<unsigned int, 4>("ui");
+    auto i8 = FLAMEGPU->getVariable<int_8, 5>("i8");
+    auto ui8 = FLAMEGPU->getVariable<uint_8, 6>("ui8");
+    auto c = FLAMEGPU->getVariable<char, 7>("c");
+    auto uc = FLAMEGPU->getVariable<unsigned char, 8>("uc");
+    auto i16 = FLAMEGPU->getVariable<int_16, 9>("i16");
+    auto ui16 = FLAMEGPU->getVariable<uint_16, 10>("ui16");
+    auto i32 = FLAMEGPU->getVariable<int_32, 11>("i32");
+    auto ui32 = FLAMEGPU->getVariable<uint_32, 12>("ui32");
+    auto i64 = FLAMEGPU->getVariable<int_64, 13>("i64");
+    auto ui64 = FLAMEGPU->getVariable<uint_64, 14>("ui64");
+}
 """
 
 py_fgpu_unknown_type = """\
@@ -324,14 +346,18 @@ def movement_request(message_in: pyflamegpu.MessageArray2D, message_out: pyflame
 """
 
 py_fgpu_standalone_msg_input = """\
-m = message_in.at(1)
-pass
-n = m.getVariableInt("foo")
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    m = message_in.at(1)
+    pass
+    n = m.getVariableInt("foo")
 """
 cpp_fgpu_standalone_msg_input = """\
-auto m = FLAMEGPU->message_in.at(1);
-;
-auto n = m.getVariable<int>("foo");
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto m = FLAMEGPU->message_in.at(1);
+    ;
+    auto n = m.getVariable<int>("foo");
+}
 """
 
 py_fgpu_for_msg_input_args = """\
@@ -356,20 +382,24 @@ def movement_request(message_in: pyflamegpu.MessageArray2D, message_out: pyflame
 """
 
 py_fgpu_macro_env_permitted = """\
-a = pyflamegpu.environment.getMacroPropertyInt('a') 
-a += 1
-a.exchange(b) 
-a.CAS(b, c)
-a.min(b)
-a.max(b)
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    a = pyflamegpu.environment.getMacroPropertyInt('a') 
+    a += 1
+    a.exchange(b) 
+    a.CAS(b, c)
+    a.min(b)
+    a.max(b)
 """
 cpp_fgpu_macro_env_permitted = """\
-auto a = FLAMEGPU->environment.getMacroProperty<int>("a");
-a += 1;
-a.exchange(b);
-a.CAS(b, c);
-a.min(b);
-a.max(b);
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto a = FLAMEGPU->environment.getMacroProperty<int>("a");
+    a += 1;
+    a.exchange(b);
+    a.CAS(b, c);
+    a.min(b);
+    a.max(b);
+}
 """
 
 py_fgpu_macro_env_function = """\
@@ -733,6 +763,7 @@ def test_check_name_not_attr(self):
 # FLAME GPU specific functionality
 
     # numpy types
+
 
     def test_fgpu_supported_types(self):
         self._checkExpected("a: numpy.byte", "char a;")
@@ -771,6 +802,9 @@ def test_fgpu_supported_types(self):
         self._checkException("a: numpy.unsupported", "Not a supported numpy type")
 
 
+    def test_fgpu_constexpr(self):
+        self._checkExpected(py_fgpu_constexpr, cpp_fgpu_constexpr)
+
     # getVariable and types
 
     def test_fgpu_types(self):

tests/python/codegen/test_codegen_integration.py

@@ -3,19 +3,24 @@
 from pyflamegpu import *
 import pyflamegpu.codegen
 from random import randint
+import typing
 
+def return_two():
+    return 2
 
 AGENT_COUNT = 100
 STEPS = 5
 
+TEN: pyflamegpu.constant = 10
+TWO: typing.Final = return_two()
 
 @pyflamegpu.device_function
 def add_2(a : int) -> int:
     """
     Pure python agent device function that can be called from a @pyflamegpu.agent_function
     Function adds two to an int variable
     """
-    return a + 2
+    return a + TWO
 
 @pyflamegpu.agent_function
 def add_func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone):
@@ -26,7 +31,7 @@ def add_func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.Message
 @pyflamegpu.agent_function_condition
 def cond_func() -> bool:
     i = pyflamegpu.getVariableInt("i")
-    return i < 10
+    return i < TEN
 
 class GPUTest(TestCase):
     """