RTC test and fix for PBM

FLAMEGPU · Jan 4, 2023 · 86dfeaf · 86dfeaf
1 parent 10b9be6
commit 86dfeaf
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 2 deletions.
diff --git a/src/flamegpu/runtime/detail/curve/curve_rtc.cpp b/src/flamegpu/runtime/detail/curve/curve_rtc.cpp
@@ -854,12 +854,11 @@ void CurveRTCHost::initHeaderGetters() {
     {
         size_t ct = 0;
         std::stringstream getDirectedGraphPBMImpl;
-        // static_cast<unsigned int*>(static_cast<void*>(getVariablePtr<unsigned int, 1>("_pbm", graphHash ^ Curve::variableHash("_environment_directed_graph_vertex"), 0)));
         for (const auto& element : directedGraph_vertexProperties) {
             RTCVariableProperties props = element.second;
             if (props.read  && element.first.second == "_pbm") {
                 getDirectedGraphPBMImpl << "            if (graphHash == " << Curve::variableRuntimeHash(element.first.first) << ") {\n";
-                getDirectedGraphPBMImpl << "                return (static_cast<unsigned int*>(static_cast<void*>(flamegpu::detail::curve::" << getVariableSymbolName() << " + " << directedGraphVertex_data_offset + (ct++ * sizeof(void*)) << ")));\n";
+                getDirectedGraphPBMImpl << "                return (*static_cast<unsigned int**>(static_cast<void*>(flamegpu::detail::curve::" << getVariableSymbolName() << " + " << directedGraphVertex_data_offset + (ct++ * sizeof(void*)) << ")));\n";
                 getDirectedGraphPBMImpl << "            }\n";
             } else { ++ct; }
         }

diff --git a/tests/python/runtime/environment/test_environment_directed_graph.py b/tests/python/runtime/environment/test_environment_directed_graph.py
@@ -411,6 +411,18 @@ class always_exit(pyflamegpu.HostCondition):
     def run(self, FLAMEGPU):
         return pyflamegpu.EXIT;
 
+class HostTestPBM(pyflamegpu.HostFunction):
+    def run(self, FLAMEGPU):
+        graph = FLAMEGPU.environment.getDirectedGraph("graph");
+        graph.setVertexCount(5);
+        graph.setEdgeCount(15);
+        k = int(0);
+        for i in range(5):
+            for j in range(i + 1):
+                graph.setEdgePropertyID("src_copy", k, j);
+                graph.setEdgeSourceDestination(k, j, i);
+                k += 1;
+
 class EnvironmentDirectedGraphTest(TestCase):
     def test_HostGetResetGet(self):
         model = pyflamegpu.ModelDescription("GraphTest");
@@ -793,3 +805,56 @@ def test_MasterModelSet_SubModelDeviceGet(self):
         sim.setPopulationData(pop);
 
         sim.step();
+
+    IterateEdges_func = """
+    FLAMEGPU_AGENT_FUNCTION(IterateEdges, flamegpu::MessageNone, flamegpu::MessageNone) {
+        flamegpu::id_t src = FLAMEGPU->getIndex();
+        unsigned int ct = 0;
+        bool src_all_correct = true;
+        auto filter = FLAMEGPU->environment.getDirectedGraph("graph").outEdges(src);
+        FLAMEGPU->setVariable<int>("count2", filter.size());
+        for (auto &edge : filter) {
+            src_all_correct &= edge.getProperty<flamegpu::id_t>("src_copy") == src;
+            FLAMEGPU->setVariable<flamegpu::id_t, 5>("dests", ct, edge.getEdgeDestination());
+            ++ct;
+        }
+        FLAMEGPU->setVariable<int>("count", ct);
+        FLAMEGPU->setVariable<int>("src_all_correct", src_all_correct ? 1 : 0);
+        return flamegpu::ALIVE;
+    }
+    """
+
+    def test_PBM(self):
+        model = pyflamegpu.ModelDescription("GraphTest");
+        graph = model.Environment().newDirectedGraph("graph");
+
+        graph.newEdgePropertyID("src_copy");
+
+        agent = model.newAgent("agent");
+        agent.newVariableArrayID("dests", 5);
+        agent.newVariableInt("count");
+        agent.newVariableInt("count2");
+        agent.newVariableInt("src_all_correct");
+        IterateEdges = agent.newRTCFunction("iterate_edges", self.IterateEdges_func);
+
+        # Init graph with known data
+        model.newLayer().addHostFunction(HostTestPBM());
+        model.newLayer().addAgentFunction(IterateEdges);
+
+        # Create enough agents, to copy all data from the 2nd graph init
+        pop = pyflamegpu.AgentVector(agent, 5);
+
+        sim = pyflamegpu.CUDASimulation(model);
+        sim.setPopulationData(pop);
+
+        sim.step();
+
+        sim.getPopulationData(pop);
+        k = int(0);
+        for agt in pop:
+            assert agt.getVariableInt("src_all_correct") == 1;
+            assert agt.getVariableInt("count") == 5 - k;
+            assert agt.getVariableInt("count2") == 5 - k;
+            for i in range(5 - k):
+                assert agt.getVariableID("dests", i) == k + i;
+            k += 1;