diff --git a/examples/lj_func.py b/examples/lj_func.py
index 61c8119f8ef5f5f1e92729f7ab4b51ca08e2efa1..c7e5618878aa24d1fe318338b5cb97c18330488a 100644
--- a/examples/lj_func.py
+++ b/examples/lj_func.py
@@ -31,6 +31,6 @@ psim.periodic(2.8)
psim.vtk_output("output/test")
psim.compute(lj, cutoff_radius, {'sigma6': sigma6, 'epsilon': epsilon})
psim.compute(euler, symbols={'dt': dt})
-psim.target(pairs.target_cpu())
-#psim.target(pairs.target_gpu())
+#psim.target(pairs.target_cpu())
+psim.target(pairs.target_gpu())
psim.generate()
diff --git a/src/pairs/analysis/modules.py b/src/pairs/analysis/modules.py
index 4da4846e1100f9f0eb93923d77ea88c86f247f11..6cf31a5d594bf35d2a8db006bb0007f4a6f70e47 100644
--- a/src/pairs/analysis/modules.py
+++ b/src/pairs/analysis/modules.py
@@ -24,11 +24,17 @@ class FetchModulesReferences(Visitor):
def visit_Array(self, ast_node):
for m in self.module_stack:
m.add_array(ast_node, self.writing)
+ if m.run_on_device:
+ ast_node.device_flag = True
def visit_Property(self, ast_node):
for m in self.module_stack:
m.add_property(ast_node, self.writing)
+ if m.run_on_device:
+ ast_node.device_flag = True
def visit_Var(self, ast_node):
for m in self.module_stack:
m.add_variable(ast_node, self.writing)
+ if m.run_on_device:
+ ast_node.device_flag = True
diff --git a/src/pairs/code_gen/cgen.py b/src/pairs/code_gen/cgen.py
index 7191b0af6fb1df98be004f7ba4ca9a6dde149b14..e5bb5aee2bb5f032f66b0377ae65b1bbf158994c 100644
--- a/src/pairs/code_gen/cgen.py
+++ b/src/pairs/code_gen/cgen.py
@@ -25,14 +25,18 @@ from pairs.code_gen.printer import Printer
class CGen:
temp_id = 0
- def __init__(self, output, debug=False):
+ def __init__(self, output, target, debug=False):
self.sim = None
+ self.target = None
self.debug = debug
self.print = Printer(output)
def assign_simulation(self, sim):
self.sim = sim
+ def assign_target(self, target):
+ self.target = target
+
def generate_program(self, ast_node):
self.print.start()
self.print("#include <math.h>")
@@ -101,7 +105,7 @@ class CGen:
t = ast_node.array.type()
tkw = Types.ctype2keyword(t)
size = self.generate_expression(BinOp.inline(ast_node.array.alloc_size()))
- if ast_node.array.is_static() and ast_node.array.init_value is not None:
+ if ast_node.array.init_value is not None:
v_str = str(ast_node.array.init_value)
if t == Types.Int64:
v_str += "LL"
@@ -180,10 +184,12 @@ class CGen:
self.print(f"{call};")
if isinstance(ast_node, CopyToDevice):
- self.print(f"pairs::copy_to_device({ast_node.prop.name()})")
+ array_name = ast_node.prop.name()
+ self.print(f"pairs::copy_to_device({array_name}, d_{array_name})")
if isinstance(ast_node, CopyToHost):
- self.print(f"pairs::copy_to_host({ast_node.prop.name()})")
+ array_name = ast_node.prop.name()
+ self.print(f"pairs::copy_to_host(d_{array_name}, {array_name})")
if isinstance(ast_node, For):
iterator = self.generate_expression(ast_node.iterator)
@@ -211,8 +217,12 @@ class CGen:
if ast_node.decl:
self.print(f"{tkw} *{array_name} = ({tkw} *) malloc({size});")
+ if self.target.is_gpu() and ast_node.array.device_flag:
+ self.print(f"{tkw} *d_{array_name} = ({tkw} *) pairs::device_alloc({size});")
else:
self.print(f"{array_name} = ({tkw} *) malloc({size});")
+ if self.target.is_gpu() and ast_node.array.device_flag:
+ self.print(f"d_{array_name} = ({tkw} *) pairs::device_alloc({size});")
if isinstance(ast_node, ModuleCall):
module = ast_node.module
@@ -244,6 +254,8 @@ class CGen:
size = self.generate_expression(ast_node.size)
array_name = ast_node.array.name()
self.print(f"{array_name} = ({tkw} *) realloc({array_name}, {size});")
+ if self.target.is_gpu() and ast_node.array.device_flag:
+ self.print(f"d_{array_name} = ({tkw} *) pairs::device_realloc(d_{array_name}, {size});")
if isinstance(ast_node, RegisterProperty):
p = ast_node.property()
diff --git a/src/pairs/ir/arrays.py b/src/pairs/ir/arrays.py
index 6573c12f513749ed0eeeb0d35fc4eb4ca9f47142..8316fc074cecffc578ca897008fea8d705164f8c 100644
--- a/src/pairs/ir/arrays.py
+++ b/src/pairs/ir/arrays.py
@@ -46,6 +46,7 @@ class Array(ASTNode):
self.arr_layout = a_layout
self.arr_ndims = len(self.arr_sizes)
self.static = False
+ self.device_flag = False
for var in [s for s in self.arr_sizes if isinstance(s, Var)]:
var.add_bonded_array(self)
diff --git a/src/pairs/ir/properties.py b/src/pairs/ir/properties.py
index 5ee78f4c656fed4dc0278f43a98f8ab6c2cab20b..7d85b19c230ddd065b576a309e8726e941a59681 100644
--- a/src/pairs/ir/properties.py
+++ b/src/pairs/ir/properties.py
@@ -57,6 +57,7 @@ class Property(ASTNode):
self.prop_layout = layout
self.default_value = default
self.volatile = volatile
+ self.device_flag = False
Property.last_prop_id += 1
def __str__(self):
diff --git a/src/pairs/ir/variables.py b/src/pairs/ir/variables.py
index d35453540dc18a936f3ebfb8ba682936564d8882..da13df11693590f953dd55a230306ee0d095ebf5 100644
--- a/src/pairs/ir/variables.py
+++ b/src/pairs/ir/variables.py
@@ -33,6 +33,7 @@ class Var(ASTTerm):
self.var_init_value = init_value
self.mutable = True
self.var_bonded_arrays = []
+ self.device_flag = False
def __str__(self):
return f"Var<{self.var_name}>"
diff --git a/src/pairs/sim/simulation.py b/src/pairs/sim/simulation.py
index d5ad6fe8f49b013d99ed58df13e92df34a61b37a..7cca2fcd377426bd33e87c1d2e8c2e682b14f17f 100644
--- a/src/pairs/sim/simulation.py
+++ b/src/pairs/sim/simulation.py
@@ -209,6 +209,7 @@ class Simulation:
def target(self, target):
self._target = target
+ self.code_gen.assign_target(target)
def generate(self):
assert self._target is not None, "Target not specified!"