diff --git a/runtime/devices/cuda.cu b/runtime/devices/cuda.cu
index 2cae5aa89accf7c720ea3f29d922184ab0264830..c411599a74a534dcfdd93e8ebdc88cdc6707bccf 100644
--- a/runtime/devices/cuda.cu
+++ b/runtime/devices/cuda.cu
@@ -88,15 +88,17 @@ __device__ double atomicAdd_double(double* address, double val) {
__device__ int atomic_add(int *addr, int val) { return atomicAdd(addr, val); }
__device__ real_t atomic_add(real_t *addr, real_t val) { return atomicAdd_double(addr, val); }
__device__ int atomic_add_resize_check(int *addr, int val, int *resize, int capacity) {
- const int add_res = *addr + val;
-
- // printf("atomic_add_resize_check::: add_res %d --- val %d --- capacity %d --- resize %d\n", add_res, val, capacity, *resize);
-
- if(add_res >= capacity) {
- *resize = add_res;
- return *addr;
+ if(*resize==0){ // If we haven't reached cap for addr before
+ const int add_res = *addr + val;
+ if(add_res >= capacity){ // Check if addr is going to exceed cap
+ *resize = add_res; // Make resize the new cap
+ return *addr; // Return addr unchanged
+ }
+ }
+ else if (*resize>0){ // If we have reached cap for addr before, resize is the new cap
+ *resize += val; // Increase cap
+ return *addr; // Return addr unchanged
}
-
return atomic_add(addr, val);
}
diff --git a/runtime/devices/device.hpp b/runtime/devices/device.hpp
index c5c406ec7c5c02634e119d4db9b6bbde3d9c0aac..bc418e5cb3ad3e9fc3d5428c9d96ee2ad7feb060 100644
--- a/runtime/devices/device.hpp
+++ b/runtime/devices/device.hpp
@@ -63,12 +63,17 @@ inline __host__ real_t host_atomic_add(real_t *addr, real_t val) {
#endif
inline __host__ int host_atomic_add_resize_check(int *addr, int val, int *resize, int capacity) {
- const int add_res = *addr + val;
- if(add_res >= capacity) {
- *resize = add_res;
- return *addr;
+ if(*resize==0){ // If we haven't reached cap for addr before
+ const int add_res = *addr + val;
+ if(add_res >= capacity){ // Check if addr is going to exceed cap
+ *resize = add_res; // Make resize the new cap
+ return *addr; // Return addr unchanged
+ }
+ }
+ else if (*resize>0){ // If we have reached cap for addr before, resize is the new cap
+ *resize += val; // Increase cap
+ return *addr; // Return addr unchanged
}
-
return host_atomic_add(addr, val);
}
diff --git a/src/pairs/sim/cell_lists.py b/src/pairs/sim/cell_lists.py
index a45ed9566116603f79fc723df7921c705ec53836..d4e571ea2891ab9f388e659ad2886bebb4fb4026 100644
--- a/src/pairs/sim/cell_lists.py
+++ b/src/pairs/sim/cell_lists.py
@@ -42,7 +42,7 @@ class CellLists:
self.nstencil_capacity = self.sim.add_var('nstencil_capacity', Types.Int32, 27)
self.ncells = self.sim.add_var('ncells', Types.Int32, 1)
self.ncells_capacity = self.sim.add_var('ncells_capacity', Types.Int32, 100000)
- self.cell_capacity = self.sim.add_var('cell_capacity', Types.Int32, 64)
+ self.cell_capacity = self.sim.add_var('cell_capacity', Types.Int32, 16)
self.dim_ncells = self.sim.add_array('dim_cells', self.sim.ndims(), Types.Int32)
self.shapes_buffer = self.sim.add_array('shapes_buffer', self.sim.max_shapes(), Types.Int32)
self.cell_particles = self.sim.add_array('cell_particles', [self.ncells_capacity, self.cell_capacity], Types.Int32)
diff --git a/src/pairs/sim/comm.py b/src/pairs/sim/comm.py
index 54e4e0cefe1ade90338c44f69dc9a20e5b30eac9..30ca8e983d7bbc5865fc8bfce1553233d86df77c 100644
--- a/src/pairs/sim/comm.py
+++ b/src/pairs/sim/comm.py
@@ -21,8 +21,8 @@ class Comm:
self.sim = sim
self.dom_part = dom_part
self.nsend_all = sim.add_var('nsend_all', Types.Int32)
- self.send_capacity = sim.add_var('send_capacity', Types.Int32, 200000)
- self.recv_capacity = sim.add_var('recv_capacity', Types.Int32, 200000)
+ self.send_capacity = sim.add_var('send_capacity', Types.Int32, 1000000)
+ self.recv_capacity = sim.add_var('recv_capacity', Types.Int32, 1000000)
self.elem_capacity = sim.add_var('elem_capacity', Types.Int32, 100)
self.nsend = sim.add_array('nsend', [dom_part.nranks_capacity], Types.Int32)
self.send_offsets = sim.add_array('send_offsets', [dom_part.nranks_capacity], Types.Int32)
diff --git a/src/pairs/transformations/modules.py b/src/pairs/transformations/modules.py
index 4b538305087bed716f778b9095d7ff24619f5ea1..5a7704b38b199359501d165a03c3774e3fc7d686 100644
--- a/src/pairs/transformations/modules.py
+++ b/src/pairs/transformations/modules.py
@@ -10,7 +10,7 @@ from pairs.ir.mutator import Mutator
from pairs.ir.properties import ReallocProperty
from pairs.ir.types import Types
from pairs.ir.print import Print
-from pairs.ir.variables import Var, Deref
+from pairs.ir.variables import Var, Deref, DeclareVariable
from functools import reduce
import operator
@@ -83,8 +83,12 @@ class AddResizeLogic(Mutator):
resizes = list(self.module_resizes[module].keys())
capacities = list(self.module_resizes[module].values())
resize_id = resizes[capacities.index(match_capacity)]
- return Branch(ast_node.sim, src + 1 >= match_capacity,
- blk_if=Block(ast_node.sim, Assign(ast_node.sim, ast_node.sim.resizes[resize_id], src)),
+ resize = ast_node.sim.resizes[resize_id]
+ return Branch(ast_node.sim, src >= match_capacity,
+ blk_if=Block(ast_node.sim,
+ Branch(ast_node.sim, ScalarOp.cmp(resize, 0),
+ blk_if=Block(ast_node.sim, Assign(ast_node.sim, resize, match_capacity)),
+ blk_else=Block(ast_node.sim, Assign(ast_node.sim, resize, resize + (src - dest))))),
blk_else=Block(ast_node.sim, ast_node))
return ast_node
@@ -167,8 +171,6 @@ class ReplaceModulesByCalls(Mutator):
def mutate_Module(self, ast_node):
ast_node._block = self.mutate(ast_node._block)
- if ast_node.name == 'main':
- return ast_node
sim = ast_node.sim
call = ModuleCall(sim, ast_node)
@@ -180,7 +182,12 @@ class ReplaceModulesByCalls(Mutator):
branch_cond = None
for resize_id, capacity in self.module_resizes[ast_node].items():
+ checked_size = ast_node._resizes_to_check[capacity]
+ original_size = sim.add_temp_var(0)
+ init_stmts.append(DeclareVariable(sim, original_size))
+ init_stmts.append(Assign(sim, original_size, checked_size)) # Get a temp backup of the original value
init_stmts.append(Assign(sim, sim.resizes[resize_id], 1))
+ reset_stmts.append(Assign(sim, checked_size, original_size)) # Reset size to the original value
reset_stmts.append(Assign(sim, sim.resizes[resize_id], 0))
cond = ScalarOp.inline(sim.resizes[resize_id] > 0)
branch_cond = cond if branch_cond is None else ScalarOp.or_op(cond, branch_cond)
@@ -199,7 +206,6 @@ class ReplaceModulesByCalls(Mutator):
[Assign(sim, capacity, self.grow_fn(sim.resizes[resize_id]))] +
[a.realloc() for a in capacity.bonded_arrays()] +
props_realloc)))
-
- return Block(sim, init_stmts + [While(sim, branch_cond, Block(sim, reset_stmts + [call] + resize_stmts))])
+ return Block(sim, init_stmts + [While(sim, branch_cond, Block(sim, reset_stmts + [call] + resize_stmts))])
return call