Adjust arrays and properties references according to module context

Signed-off-by: Rafael Ravedutti <rafaelravedutti@gmail.com>

src/pairs/code_gen/cgen.py

@@ -236,11 +236,11 @@ class CGen:
                 module_params += decl if len(module_params) <= 0 else f", {decl}"
 
             for array in module.arrays():
-                decl = array.name()
+                decl = f"d_{array.name()}" if module.run_on_device else array.name()
                 module_params += decl if len(module_params) <= 0 else f", {decl}"
 
             for prop in module.properties():
-                decl = prop.name()
+                decl = f"d_{prop.name()}" if module.run_on_device else prop.name()
                 module_params += decl if len(module_params) <= 0 else f", {decl}"
 
             self.print(f"{module.name}({module_params});")

src/pairs/ir/block.py

@@ -2,7 +2,7 @@ from pairs.ir.ast_node import ASTNode
 from pairs.ir.module import Module
 
 
-def pairs_block(func):
+def pairs_inline(func):
     def inner(*args, **kwargs):
         sim = args[0].sim # self.sim
         sim.init_block()
@@ -12,6 +12,21 @@ def pairs_block(func):
     return inner
 
 
+def pairs_host_block(func):
+    def inner(*args, **kwargs):
+        sim = args[0].sim # self.sim
+        sim.init_block()
+        func(*args, **kwargs)
+        return Module(sim,
+            name=sim._module_name,
+            block=Block(sim, sim._block),
+            resizes_to_check=sim._resizes_to_check,
+            check_properties_resize=sim._check_properties_resize,
+            run_on_device=False)
+
+    return inner
+
+
 def pairs_device_block(func):
     def inner(*args, **kwargs):
         sim = args[0].sim # self.sim

src/pairs/sim/arrays.py

-from pairs.ir.block import pairs_block
+from pairs.ir.block import pairs_inline
 from pairs.ir.memory import Malloc
 from pairs.ir.arrays import ArrayDecl
 from pairs.sim.lowerable import Lowerable
@@ -8,7 +8,7 @@ class ArraysDecl(Lowerable):
     def __init__(self, sim):
         super().__init__(sim)
 
-    @pairs_block
+    @pairs_inline
     def lower(self):
         for a in self.sim.arrays.all():
             if a.is_static():

src/pairs/sim/cell_lists.py

 from functools import reduce
 import math
 from pairs.ir.bin_op import BinOp
-from pairs.ir.block import pairs_device_block
+from pairs.ir.block import pairs_device_block, pairs_host_block
 from pairs.ir.branches import Branch, Filter
 from pairs.ir.cast import Cast
 from pairs.ir.loops import For, ParticleFor
@@ -36,7 +36,7 @@ class CellListsStencilBuild(Lowerable):
         super().__init__(sim)
         self.cell_lists = cell_lists
 
-    @pairs_device_block
+    @pairs_host_block
     def lower(self):
         sim = self.sim
         cl = self.cell_lists

src/pairs/sim/interaction.py

 from pairs.ir.bin_op import BinOp
-from pairs.ir.block import Block, pairs_block
+from pairs.ir.block import Block, pairs_device_block
 from pairs.ir.branches import Branch, Filter
 from pairs.ir.loops import For, ParticleFor
 from pairs.ir.types import Types
@@ -60,7 +60,7 @@ class ParticleInteraction(Lowerable):
         yield self.i, self.j
         self.sim.leave()
 
-    @pairs_block
+    @pairs_device_block
     def lower(self):
         if self.nbody == 2:
             position = self.sim.position()

src/pairs/sim/lattice.py

-from pairs.ir.block import pairs_block
+from pairs.ir.block import pairs_inline
 from pairs.ir.loops import For
 from pairs.ir.types import Types
 from pairs.sim.lowerable import Lowerable
@@ -12,7 +12,7 @@ class ParticleLattice(Lowerable):
         self.props = props
         self.positions = positions
 
-    @pairs_block
+    @pairs_inline
     def lower(self):
         index = None
         loop_indexes = []

src/pairs/sim/pbc.py

-from pairs.ir.block import pairs_device_block
+from pairs.ir.block import pairs_device_block, pairs_host_block
 from pairs.ir.branches import Branch, Filter
 from pairs.ir.loops import For, ParticleFor
 from pairs.ir.utils import Print
@@ -72,7 +72,7 @@ class SetupPBC(Lowerable):
         super().__init__(sim)
         self.pbc = pbc
 
-    @pairs_device_block
+    @pairs_host_block
     def lower(self):
         sim = self.sim
         ndims = sim.ndims()

src/pairs/sim/properties.py

-from pairs.ir.block import pairs_block, pairs_device_block
+from pairs.ir.block import pairs_device_block, pairs_inline
 from pairs.ir.loops import ParticleFor
 from pairs.ir.memory import Malloc, Realloc
 from pairs.ir.properties import RegisterProperty, UpdateProperty
@@ -14,7 +14,7 @@ class PropertiesAlloc(Lowerable):
         self.sim = sim
         self.realloc = realloc
 
-    @pairs_block
+    @pairs_inline
     def lower(self):
         capacity = sum(self.sim.properties.capacities)
         for p in self.sim.properties.all():

src/pairs/sim/read_from_file.py

-from pairs.ir.block import pairs_block
+from pairs.ir.block import pairs_inline
 from pairs.ir.functions import Call_Int
 from pairs.ir.properties import PropertyList
 from pairs.ir.types import Types
@@ -14,7 +14,7 @@ class ReadFromFile(Lowerable):
         self.grid = MutableGrid(sim, sim.ndims())
         self.grid_buffer = self.sim.add_static_array("grid_buffer", [self.sim.ndims() * 2], Types.Double)
 
-    @pairs_block
+    @pairs_inline
     def lower(self):
         self.sim.nlocal.set(Call_Int(self.sim, "pairs::read_particle_data",
                             [self.filename, self.grid_buffer, self.props, self.props.length()]))

src/pairs/sim/variables.py

-from pairs.ir.block import pairs_block
+from pairs.ir.block import pairs_inline
 from pairs.ir.variables import VarDecl
 from pairs.sim.lowerable import Lowerable
 
@@ -7,7 +7,7 @@ class VariablesDecl(Lowerable):
     def __init__(self, sim):
         super().__init__(sim)
 
-    @pairs_block
+    @pairs_inline
     def lower(self):
         for v in self.sim.vars.all():
             VarDecl(self.sim, v)

src/pairs/sim/vtk.py

 from pairs.ir.ast_node import ASTNode
-from pairs.ir.block import pairs_block
+from pairs.ir.block import pairs_inline
 from pairs.ir.functions import Call_Void
 from pairs.ir.lit import Lit 
 from pairs.sim.lowerable import Lowerable
@@ -11,7 +11,7 @@ class VTKWrite(Lowerable):
         self.filename = filename
         self.timestep = Lit.cvt(sim, timestep)
 
-    @pairs_block
+    @pairs_inline
     def lower(self):
         nlocal = self.sim.nlocal
         npbc = self.sim.pbc.npbc