diff --git a/swig/python/codegen/__init__.py b/swig/python/codegen/__init__.py index a51633e1f..cf27dd11d 100644 --- a/swig/python/codegen/__init__.py +++ b/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") \ No newline at end of file diff --git a/swig/python/codegen/codegen.py b/swig/python/codegen/codegen.py index e7781b391..4c04359e9 100644 --- a/swig/python/codegen/codegen.py +++ b/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]) diff --git a/swig/python/flamegpu.i b/swig/python/flamegpu.i index 364995d91..8ed346388 100644 --- a/swig/python/flamegpu.i +++ b/swig/python/flamegpu.i @@ -1182,4 +1182,10 @@ TEMPLATE_VARIABLE_INSTANTIATE_INTS(poisson, flamegpu::HostRandom::poisson) #define GLM true #else #define GLM false -#endif \ No newline at end of file +#endif + +// Declare an empty type we can use as an attribute for constants to be pulled in by codegen +%pythoncode { +class constant: + pass; +} diff --git a/tests/python/codegen/test_codegen.py b/tests/python/codegen/test_codegen.py index 0fc5da0d4..5a0cfc8cd 100644 --- a/tests/python/codegen/test_codegen.py +++ b/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("f"); -auto d = FLAMEGPU->getVariable("d"); -auto i = FLAMEGPU->getVariable("i"); -auto ui = FLAMEGPU->getVariable("ui"); -auto i8 = FLAMEGPU->getVariable("i8"); -auto ui8 = FLAMEGPU->getVariable("ui8"); -auto c = FLAMEGPU->getVariable("c"); -auto uc = FLAMEGPU->getVariable("uc"); -auto i16 = FLAMEGPU->getVariable("i16"); -auto ui16 = FLAMEGPU->getVariable("ui16"); -auto i32 = FLAMEGPU->getVariable("i32"); -auto ui32 = FLAMEGPU->getVariable("ui32"); -auto i64 = FLAMEGPU->getVariable("i64"); -auto ui64 = FLAMEGPU->getVariable("ui64"); +FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){ + auto f = FLAMEGPU->getVariable("f"); + auto d = FLAMEGPU->getVariable("d"); + auto i = FLAMEGPU->getVariable("i"); + auto ui = FLAMEGPU->getVariable("ui"); + auto i8 = FLAMEGPU->getVariable("i8"); + auto ui8 = FLAMEGPU->getVariable("ui8"); + auto c = FLAMEGPU->getVariable("c"); + auto uc = FLAMEGPU->getVariable("uc"); + auto i16 = FLAMEGPU->getVariable("i16"); + auto ui16 = FLAMEGPU->getVariable("ui16"); + auto i32 = FLAMEGPU->getVariable("i32"); + auto ui32 = FLAMEGPU->getVariable("ui32"); + auto i64 = FLAMEGPU->getVariable("i64"); + auto ui64 = FLAMEGPU->getVariable("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("f"); -auto d = FLAMEGPU->getVariable("d"); -auto i = FLAMEGPU->getVariable("i"); -auto ui = FLAMEGPU->getVariable("ui"); -auto i8 = FLAMEGPU->getVariable("i8"); -auto ui8 = FLAMEGPU->getVariable("ui8"); -auto c = FLAMEGPU->getVariable("c"); -auto uc = FLAMEGPU->getVariable("uc"); -auto i16 = FLAMEGPU->getVariable("i16"); -auto ui16 = FLAMEGPU->getVariable("ui16"); -auto i32 = FLAMEGPU->getVariable("i32"); -auto ui32 = FLAMEGPU->getVariable("ui32"); -auto i64 = FLAMEGPU->getVariable("i64"); -auto ui64 = FLAMEGPU->getVariable("ui64"); +FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){ + auto f = FLAMEGPU->getVariable("f"); + auto d = FLAMEGPU->getVariable("d"); + auto i = FLAMEGPU->getVariable("i"); + auto ui = FLAMEGPU->getVariable("ui"); + auto i8 = FLAMEGPU->getVariable("i8"); + auto ui8 = FLAMEGPU->getVariable("ui8"); + auto c = FLAMEGPU->getVariable("c"); + auto uc = FLAMEGPU->getVariable("uc"); + auto i16 = FLAMEGPU->getVariable("i16"); + auto ui16 = FLAMEGPU->getVariable("ui16"); + auto i32 = FLAMEGPU->getVariable("i32"); + auto ui32 = FLAMEGPU->getVariable("ui32"); + auto i64 = FLAMEGPU->getVariable("i64"); + auto ui64 = FLAMEGPU->getVariable("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("foo"); +FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){ + auto m = FLAMEGPU->message_in.at(1); + ; + auto n = m.getVariable("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("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("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): diff --git a/tests/python/codegen/test_codegen_integration.py b/tests/python/codegen/test_codegen_integration.py index 58380a902..17a5369d4 100644 --- a/tests/python/codegen/test_codegen_integration.py +++ b/tests/python/codegen/test_codegen_integration.py @@ -3,11 +3,16 @@ 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: @@ -15,7 +20,7 @@ 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): """